有没有办法加快这个训练过程？

Question

我有 5 个数据集，我从一个批次中的每个数据集中获取 3 张图像，因此批次大小为 5 * 3 = 15。

for batch in batches:
   for i, d in enumerate(batch):
       optimizer.zero_grad()
       r = net.forward(d)

我需要分别计算和保存每个数据集的梯度：

loss = loss_fn(r)
loss.backward()
for param in net.parameters():
    if param.grad is not None:
       grads[i].append(Variable(param.grad.data.clone(), requires_grad=False))

然后我需要对这些梯度进行一些计算，以获取加权损失的系数：

scale = get_scales(grads)
optimizer.zero_grad()
for i, d in enumerate(batch):
    r = net.forward(d)
    loss_t = loss_fn(r)
    if i > 0:
       loss = loss + scale[i]*loss_t
     else:
       loss = scale[i]*loss_t
 loss.backward()
 optimizer.step()

所以这里发生了什么，我做了 5 次前向和后向传递来计算梯度，然后我计算系数，然后再进行 5 次前向和后向传递，损失求和，最后是优化器步骤。 所以向前和向后传递的数量是x2。而且需要的时间太长了。 GPU 的内存有限，仅适合前向和后向传递中的这 3 个图像。 这是目前我想出的唯一简单的解决方案。计算梯度并丢弃图形。我可以以某种方式保存第一次通过的计算吗？但我想我很快就会失去记忆。

是否有一些解决方案可以在内存受限的环境中加快此过程？我对 pytorch 比较陌生，所以我问这个问题是希望我不知道可以在这里提供帮助的东西。

完整的伪代码示例：

for batch in batches:
   for i, d in enumerate(batch):
       optimizer.zero_grad()
       r = net.forward(d)

       loss = loss_fn(r)
       loss.backward()
       for param in net.parameters():
           if param.grad is not None:
              grads[i].append(Variable(param.grad.data.clone(), requires_grad=False))

   scale = get_scales(grads)
   optimizer.zero_grad()
   for i, d in enumerate(batch):
       r = net.forward(d)
       loss_t = loss_fn(r)
       if i > 0:
          loss = loss + scale[i]*loss_t
       else:
          loss = scale[i]*loss_t
    loss.backward()
    optimizer.step()

Answer 1

据我了解，您使用相同的 r 来计算 loss_ft 和 loss_fn。您至少应该能够重用 last r 来计算 loss_ft 和 loss_fn，因为这不需要额外的内存。 您可以尝试在backward() 中使用retain_graph=True 来第二次将图形保留为反向传播（尽管文档建议使用此标志的解决方法我在这里看不到任何其他内容https://pytorch.org/docs/1.9.1/generated/torch.Tensor.backward.html）