如何在 loss.backward() 之后立即获得梯度总和？

Question

我是 Pytorch 的新手，我正在尝试做一些重要性抽样实验： 在评估时期，我计算每个训练样本的损失，并获得该训练样本的梯度总和。最后，我将根据他们引入的梯度对训练样本进行排序。例如，如果样本 A 显示出非常高的梯度和，则它必须是训练的重要样本。否则，它不是一个非常重要的样本。

请注意，此处计算的梯度不会用于更新参数。也就是说，它们只用于选择重要性样本。

我知道梯度会在 loss.backward() 之后的某个地方准备好。但是，在整个模型中获取总梯度的最简单方法是什么？在我目前的实现中，我只允许修改一个只有损失可用的小模块，所以我没有“输入”或“模型”。是否可以仅从“损失”中获得梯度？

Answer 1

backward 后的梯度存储为需要梯度的张量的grad 属性。你可以找到所有涉及的张量并总结他们的grads。一种更简洁的方法可能是编写一个反向挂钩，以便在反向传播时将梯度累积到某个全局变量。

一个例子是

import torch
import torch.nn as nn

model = nn.Linear(5, 3)
print(model.weight.grad)  # None, since the grads have not been computed yet
print(model.bias.grad)

x = torch.randn(5, 5)
y = model(x)
loss = y.sum()
loss.backward()

print(model.weight.grad)
print(model.bias.grad)

输出：

None
None
tensor([[-0.6164,  1.1585, -3.4117, -4.3192, -3.7273],
        [-0.6164,  1.1585, -3.4117, -4.3192, -3.7273],
        [-0.6164,  1.1585, -3.4117, -4.3192, -3.7273]])
tensor([5., 5., 5.])

如您所见，您可以使用param.grad 访问渐变。如果model 是torch.nn.Module 对象，您可以使用for param in model.parameters() 对其参数进行迭代。

也许您也可以使用后向挂钩，但我对它们不太熟悉，无法给出代码示例。