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 参数,这样每次更新完成后会释放占用的内存,也就不会出现越来越慢的情况了

Logo

GitCode 天启AI是一款由 GitCode 团队打造的智能助手,基于先进的LLM(大语言模型)与多智能体 Agent 技术构建,致力于为用户提供高效、智能、多模态的创作与开发支持。它不仅支持自然语言对话,还具备处理文件、生成 PPT、撰写分析报告、开发 Web 应用等多项能力,真正做到“一句话,让 Al帮你完成复杂任务”。

更多推荐