模型并行性，Pytorch 中 CUDA 内存不足

Question

我正在尝试构建自动编码器模型，其中 input/output 是大小为 256 x 256 的 RGB 图像。我尝试在具有 12 GB 内存的 1 个 GPU 上训练模型，但我总是遇到 CUDA OOM（我尝试了不同的批大小，甚至 1 的批大小都失败了）。所以我阅读了 Pytorch 中的模型并行性并尝试了这个：

class Autoencoder(nn.Module):
    def __init__(self, input_output_size):
        super(Autoencoder, self).__init__()
        self.encoder = nn.Sequential(
            nn.Linear(input_output_size, 1024),
            nn.ReLU(True),
            nn.Linear(1024, 200),
            nn.ReLU(True)
            ).cuda(0)
      
        self.decoder = nn.Sequential(
           nn.Linear(200, 1024),
           nn.ReLU(True),
           nn.Linear(1024, input_output_size),
           nn.Sigmoid()).cuda(1)
        
        print(self.encoder.get_device())
        print(self.decoder.get_device())

    def forward(self, x):
        x = x.cuda(0)
        x = self.encoder(x)
        x = x.cuda(1)
        x = self.decoder(x)
        return x

所以我将我的编码器和解码器移到了不同的 GPU 上。但现在我得到这个例外：

Expected tensor for 'out' to have the same device as tensor for argument #2 'mat1'; but device 0 does not equal 1 (while checking arguments for addmm)

当我在 forward 方法中执行 x = x.cuda(1) 时出现。

此外，这是我的“train”代码，您可以给我一些优化建议吗？ 3 x 256 x 256 的图像对于训练来说太大了吗？（我不能减少它们）。提前谢谢你。

培训：

input_output_size = 3 * 256 * 256
model = Autoencoder(input_output_size).to(device)
optimizer = optim.Adam(model.parameters(), lr=1e-4)
criterion = nn.MSELoss()

for epoch in range(100):
    epoch_loss = 0
    for batch_idx, (images, _) in enumerate(dataloader):
        images = torch.flatten(images, start_dim=1).to(device)
        output_images = model(images).to(device)
        train_loss = criterion(output_images, images)
            
        
        train_loss.backward()
        optimizer.step()

        if batch_idx % 5 == 0:
            with torch.no_grad():
                model.eval()
                pred = model(test_set).to(device)
                model.train()

                test_loss = criterion(pred, test_set)

                wandb.log({"MSE train": train_loss})
                wandb.log({"MSE test": test_loss})
                del pred, test_loss

        if batch_idx % 200 == 0:
            # here I send testing images from output to W&B
            with torch.no_grad():
                model.eval()
                pred = model(test_set).to(device)
                model.train()
                wandb.log({"PRED": [wandb.Image((pred[i].cpu().reshape((3, 256, 256)).permute(1, 2, 0) * 255).numpy().astype(np.uint8), caption=str(i)) for i in range(20)]})
                del pred
        gc.collect()
        torch.cuda.empty_cache()
        epoch_loss += train_loss.item()
        del output_images, train_loss
    epoch_loss = epoch_loss / len(dataloader)
    wandb.log({"Epoch MSE train": epoch_loss})
    del epoch_loss

Answer 1

我看到了三个问题：

model(test_set)

此时您将 整个测试集（可能很大）作为单个批次发送到您的模型中。

我不知道 wandb 是什么，但内存增长的另一个可能来源是这些行：

wandb.log({"MSE train": train_loss})
wandb.log({"MSE test": test_loss})

您似乎在节省 train_loss 和 test_loss，但它们不仅包含数字本身，还包含反向传播所需的计算图（存在于 GPU 上）。在保存它们之前，您想将它们转换为 float 或 numpy.

您的模型包含两个 3*256*256 x 1024 配重块。当在 Adam 中使用时，这些将需要 3*256*256 x 1024 * 3 * 4 bytes = 2.25GB 的 VRAM 每个（如果实施效率低下，可能更多）由于其他原因，这看起来也是一个糟糕的架构。

模型并行性，Pytorch 中 CUDA 内存不足

Model parallelism, CUDA out of memory in Pytorch

python

gpu

out-of-memory

pytorch