Pytorch中有多次backward时需要retain_graph参数

Pytorch中有多次backward时需要retain_graph参数

背景介绍

Pytorch中的机制是每次调用loss.backward()时都会free掉计算图中所有缓存的buffers，当模型中可能有多次backward()时，因为前一次调用backward()时已经释放掉了buffer，所以下一次调用时会因为buffers不存在而报错

解决办法

loss.backward(retain_graph=True)

错误使用

optimizer.zero_grad() 清空过往梯度；
loss1.backward(retain_graph=True) 反向传播，计算当前梯度；
loss2.backward(retain_graph=True) 反向传播，计算当前梯度；
optimizer.step() 根据梯度更新网络参数

因为每次调用bckward时都没有将buffers释放掉，所以会导致内存溢出，迭代越来越慢（因为梯度都保存了，没有free)

正确使用

optimizer.zero_grad() 清空过往梯度；
loss1.backward(retain_graph=True) 反向传播，计算当前梯度；
loss2.backward() 反向传播，计算当前梯度；
optimizer.step() 根据梯度更新网络参数

最后一个 backward() 不要加 retain_graph 参数，这样每次更新完成后会释放占用的内存，也就不会出现越来越慢的情况了