python - 在 PyTorch 中，如何通过损失列表中的平均梯度更新神经网络？

Question

我有一个基于 REINFORCE 算法(此处为 PyTorch's implementation)的玩具强化学习项目，我想向其添加批量更新。在 RL 中，“目标”只能在做出“预测”后创建，因此标准的批处理技术不适用。因此，我为每一集计算损失并将它们附加到列表 l_losses 中，其中每个项目都是零维张量。我推迟调用 .backward() 或 optimizer.step() 直到一定数量的剧集已经过去以创建一种伪批处理。

鉴于此损失列表，我如何让 PyTorch 根据它们的平均梯度更新网络？或者基于平均梯度的更新是否与基于平均损失的更新相同(我似乎在其他地方读过)？

我目前的方法是从 torch.stack(l_losses) 创建一个新的张量 t_loss，然后运行 ​​t_loss = t_loss.mean()、t_loss.backward()、optimizer.step() 和梯度归零，但我不确定这是否等同于我的意图？我也不清楚我是否应该在每个单独的损失上运行 .backward() 而不是将它们连接在列表中(但保留 .step() 部分到最后？

Answer 1

梯度是线性运算，所以梯度的平均值与梯度的平均值相同。

拿一些例子数据

import torch
a = torch.randn(1, 4, requires_grad=True);
b = torch.randn(5, 4);

您可以存储所有损失并计算平均值，

a.grad = None
x = (a * b).mean(axis=1)
x.mean().backward() # gradient of the mean
print(a.grad)

或者每次迭代计算反向传播以获得该损失对梯度的贡献。

a.grad = None
for bi in b:
    (a * bi / len(b)).mean().backward()
print(a.grad)

性能

我不知道pytorch后向实现的内部细节，但我可以说

(1) 在反向传递 ratain_graph=True 或 create_graph=True 到 backward() 之后，图默认被销毁.

(2) 除叶张量外，不保留梯度，除非您指定 retain_grad ;

(3) 如果您使用不同的输入对模型进行两次评估，您可以对单个变量执行反向传递，这意味着它们具有单独的图形。这可以通过以下代码验证。

a.grad = None
# compute all the variables in advance
r = [ (a * b / len(b)).mean() for bi in b ]
for ri in r:
    # This depends on the graph of r[i] but the graph or r[i-1]
    # was already destroyed, it means that r[i] graph is independent
    # of r[i-1] graph, hence they require separate memory.
    ri.backward()  # this will remove the graph of ri
print(a.grad)

因此，如果您在每一集之后更新梯度，它将累积叶节点的梯度，这就是下一个优化步骤所需的所有信息，因此您可以丢弃该损失，从而释放资源用于进一步的计算。我希望内存使用量减少，如果内存分配可以有效地使用刚刚释放的页面进行下一次分配，则执行速度可能会更快。