模型并行，数据自动化系统内存不足

Question

我正在尝试构建自动编码器模型，其中输入/输出是RGB图像，大小为256 x 256。我试着用12 GB的内存在1GPU上训练模型，但我总是抓到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)

在正向方法中，当我执行x= x.cuda(1)时，它就出现了。

此外，这是我的“火车”代码，也许你能给我一些关于优化的建议吗？3x256x256的图像对训练来说是不是太大了？(我不能减少)。提前谢谢你。

培训：

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 --每个(可能更多，如果实现效率不高)--这看起来像一个糟糕的体系结构，原因还有其他原因。