在 PyTorch 中为批处理中的每个单独样本计算梯度

Question

我正在尝试实现一个版本的差异私有随机梯度下降（例如this），如下所示：

计算大小为 L 的批次中每个点的梯度，然后分别裁剪 L 个梯度中的每一个，然后将它们平均在一起，最后执行（噪声）梯度下降步骤。

在 pytorch 中执行此操作的最佳方法是什么？

最好有一种方法可以同时计算批处理中每个点的梯度：

x # inputs with batch size L
y #true labels
y_output = model(x)
loss = loss_func(y_output,y) #vector of length L
loss.backward() #stores L distinct gradients in each param.grad, magically

但是失败了，分别计算每个梯度，然后在累积之前裁剪范数，但是

x # inputs with batch size L
y #true labels
y_output = model(x)
loss = loss_func(y_output,y) #vector of length L   
for i in range(loss.size()[0]):
    loss[i].backward(retain_graph=True)
    torch.nn.utils.clip_grad_norm(model.parameters(), clip_size)

累积第 i 个梯度，然后剪辑，而不是先剪辑再累积到梯度中。解决此问题的最佳方法是什么？

Answer 1

我认为在计算效率方面你不能比第二种方法做得更好，你正在失去批处理的好处backward，这是事实。关于裁剪的顺序，autograd 将梯度存储在.grad参数张量中。一个粗略的解决方案是添加一个字典，例如

clipped_grads = {name: torch.zeros_like(param) for name, param in net.named_parameters()}

像这样运行你的 for 循环

for i in range(loss.size(0)):
    loss[i].backward(retain_graph=True)
    torch.nn.utils.clip_grad_norm_(net.parameters())
    for name, param in net.named_parameters():
        clipped_grads[name] += param.grad / loss.size(0)
    net.zero_grad()

for name, param in net.named_parameters():
    param.grad = clipped_grads[name]

optimizer.step()

其中我省略了大部分和 类似的业务detach，requires_grad=False这些业务可能是使其按预期运行所必需的。

上述方法的缺点是您最终会为参数梯度存储 2 倍的内存。原则上，您可以采用“原始”梯度，对其进行剪辑，添加到clipped_gradient，然后在下游操作不需要它时立即丢弃，而在这里您保留原始值，grad直到向后传递结束。如果您违反指南并实际修改了，则register_backward_hook可能允许您这样做，但您必须与更熟悉 autograd 的人进行验证。grad_input