使用torch.autograd.grad()后各epoch运行时间持续递增的原因排查

核心原因 #

你遇到的epoch运行时间逐次递增问题，本质是**retain_graph=True导致计算图未被及时释放，内存/显存持续堆积**。每次调用torch.autograd.grad时设置该参数，会强制保留整个计算图结构；后续loss.backward()又会生成新的计算图，旧图无法被垃圾回收，随着epoch推进，内存占用越来越高，最终拖慢运行速度。

修复方案 #

根据你的代码逻辑，只需要在计算loss2_grads时保留计算图（因为后续还要执行loss.backward()），loss1_grads计算完成后不需要保留，因此修改代码如下：

# loss1的梯度计算后不需要保留图，设为False
loss1_grads = torch.autograd.grad(loss1, model.parameters(), retain_graph=False) 
# loss2的梯度计算需要保留图，供后续loss.backward()使用
loss2_grads = torch.autograd.grad(loss2, model.parameters(), retain_graph=True)

# 计算加权损失并反向传播
loss = loss1 + alpha * loss2  # alpha由loss1_grads和loss2_grads推导
loss.backward() 

额外注意事项 #

• 非必要情况下不要随意设置retain_graph=True，backward()默认会在执行后销毁计算图，避免内存泄漏。
• 可以在每个epoch结束时通过torch.cuda.memory_allocated()（GPU环境）或torch.utils.memory_summary()查看内存占用，验证问题是否解决。

内容的提问来源于stack exchange，提问作者Knotnet