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

Question

我正在使用 PyTorch 闪电，所以闪电控制 GPU/CPU 分配和 return 我得到了简单的多 GPU 训练支持。

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

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

然后当我 select 一个批次的子集时，将它发送到 GPU

GPU_tensor = embedding(idx)

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

Answer 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

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

Pytorch embedding too big for GPU but fits in CPU

pytorch-lightning