pytorch训练解决显存爆炸问题

1.torch.cuda.empty_cache()
每轮epoch后加这个代码实测没有用
2.存储loss时将loss改为loss.item(),问题解决。

loss_list=[]

在我的代码中,在训练前建立了一个loss的list来存储每轮的loss,目的为了在N个EPOCH后绘制图像。

loss = criterion(outputs, dec_outputs.view(-1).reshape(-1,1))
optimizer.zero_grad()
loss.backward()
optimizer.step()
loss_list.append(loss)

原本代码截取一段如上,在loss_list中每次append loss

loss_list.append(loss)

改为:

loss_list.append(loss.item())

后显存爆炸问题解决,再每个epoch后不再增加GPU显存占用。原理目前不清楚,但是这个是在网上的很多方法中亲测有用的唯一一个方法,大家可以尝试一下。

你可能感兴趣的:(pytorch学习)