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

Question

我正在尝试实现一个差分私有(private)随机梯度下降的版本(例如 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如果您违反指南并实际修改 grad_input，则允许您这样做，但您必须与更熟悉 autograd 的人核实。