Fork me on GitHub

pytorch 深度学习训练过程gpu内存溢出问题

Pytorch GPU运算过程中会出现:“cuda runtime error(2): out of memory”这样的错误。通常,这种错误是由于在循环中使用全局变量当做累加器,且累加梯度信息的缘故,用官方的说法就是:"accumulate history across your training loop"。在默认情况下,开启梯度计算的Tensor变量是会在GPU保持他的历史数据的,所以在编程或者调试过程中应该尽力避免在循环中累加梯度信息。

代码:
total_loss=0
for i in range(10000):
optimizer.zero_grad()
output=model(input)
loss=criterion(output)
loss.backward()
optimizer.step()
total_loss+=loss
#这里total_loss是跨越循环的变量,起着累加的作用,
#loss变量是带有梯度的tensor,会保持历史梯度信息,在循环过程中会不断积累梯度信息到tota_loss,占用内存

以上例子的修正方法是在循环中的最后一句修改为:total_loss+=float(loss),利用类型变换解除梯度信息,这样,多次累加不会累加梯度信息。

posted @ 2021-07-27 11:13  sy-  阅读(1455)  评论(0编辑  收藏  举报