使用 GPU 时，pytorch 中的 loss.backward() 停止响应答案

【问题标题】：loss.backward() in pytorch stops responding when using GPU使用 GPU 时，pytorch 中的 loss.backward() 停止响应
【发布时间】：2021-10-31 19:20:57
【问题描述】：

我在基于 WSL 的环境中使用 pytorch。当我在 GPU(RTX3090) 上训练深度学习模型时，如下所示，它在 loss.backward() 处停止响应。

pytorch停止后，GPU占用率为0%，nvidia-smi停止响应。

def train_loop(model, optimizer, scheduler, loader, device):
   losses, lrs = [], []
   model.train()
   optimizer.zero_grad()
   for i, d in enumerate(loader):
       print(f"{i}-start")
       out, loss = model(d['X'].to(device), d['y'].to(device))
       print(f"{i}-goal")
       losses.append(loss.item())
       step_lr = np.array([param_group["lr"] for param_group in optimizer.param_groups]).mean()
       lrs.append(step_lr)
     
       loss.backward()
       optimizer.step()
       optimizer.zero_grad()
       scheduler.step() 

   return True

例如。第 45 批停止。

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
device
-> device(type='cuda', index=0)

train_loop(model, optimizer, scheduler, train_loader, device)
-> 0-start
   0-goal
   ...
   43-goal
   44-start

我想让你告诉我为什么它不起作用。

【问题讨论】：

“停止响应”是否意味着没有错误消息并且您陷入了某个无限循环？
+ 可能与您的问题无关，但 scheduler.step() 预计会在每个时期被调用，而不是每个步骤（除非自定义定义）。
是的。不显示错误消息，并且该过程不会继续。
+ 谢谢。注释掉 scheduler.step() 没有帮助。
你检查过GPU内存使用吗？用了多少？它是在训练期间保持增长还是保持不变？

标签： python pytorch

【解决方案1】：

在 PyTorch 中，我们在开始对每个小批量进行反向传播之前明确地将梯度设置为零。

我无法确认在后续通道中累积过多渐变是导致问题的原因，但很有可能。此外，通过这种方式，您会混淆 mini-batch 之间的梯度，这会影响模型性能

所以，optimizer.zero_grad() 必须在您的训练循环中。

【讨论】：

在我的train_loop 函数中，看起来我正在训练循环中运行optimizer.zero_grad()。有什么问题吗？
是的；但是，我的意思是在小批量循环中。
如何在小批量循环中做optimizer.zero_grad()？
在for i, d in enumerate(loader):循环内移动行？