Pytorch 嵌入对于 GPU 来说太大但适合 CPU答案

【问题标题】：Pytorch embedding too big for GPU but fits in CPUPytorch 嵌入对于 GPU 来说太大但适合 CPU
【发布时间】：2022-03-23 21:46:52
【问题描述】：

我正在使用 PyTorch 闪电，所以闪电控制 GPU/CPU 分配并在 return 我可以轻松获得多 GPU 支持来进行训练。

我想创建一个不适合 GPU 内存的嵌入。

fit_in_cpu = torch.nn.Embedding(too_big_for_GPU, embedding_dim)

然后当我为一个批次选择子集时，将其发送到 GPU

GPU_tensor = embedding(idx)

如何在 Pytorch Lightning 中执行此操作？

【问题讨论】：

标签： pytorch-lightning

【解决方案1】：

Lightning 会将注册为模型参数的任何内容发送到 GPU，即：层的权重（torch.nn.* 中的任何内容）和使用torch.nn.parameter.Parameter 注册的变量。

但是，如果您想在 CPU 中声明某些内容，然后在运行时将其移至 GPU，您可以采用两种方式：

在__init__ 内创建too_big_for_GPU，而不将其注册为模型参数（使用torch.zeros 或torch.randn 或任何其他初始化函数）。然后在前向传播中将其移动到 GPU

class MyModule(pl.LightningModule):
    def __init__():
        self.too_big_for_GPU = torch.zeros(4, 1000, 1000, 1000)
    def forward(self, x):
        # Move tensor to same GPU as x and operate with it
        y = self.too_big_for_GPU.to(x.device) * x**2
        return y

创建too_big_for_GPU，它将默认在 CPU 中创建，然后您需要将其移动到 GPU

class MyModule(pl.LightningModule):
    def forward(self, x):
        # Create the tensor on the fly and move it to x GPU
        too_big_for_GPU = torch.zeros(4, 1000, 1000, 1000).to(x.device)
        # Operate with it
        y = too_big_for_GPU * x**2
        return y

【讨论】：