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

Question

我正在尝试构建自动编码器模型，其中输入/输出是大小为 256 x 256 的 RGB 图像。我尝试在 1 个具有 12 GB 内存的 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) 时出现。

此外，这是我的“火车”代码，你可以给我一些关于优化的建议吗？ 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 每个（可能更多，如果实现效率低下）由于其他原因，这看起来像是一个糟糕的架构。