如何高效计算多个样本的逐样本梯度？

huangapple
go评论110阅读模式
英文:

How do I compute multiple per-sample gradients efficiently?

问题

  1. 我正在尝试在PyTorch中高效地计算多个损失梯度（不使用for循环）。给定：
  3. ```python
  4. import torch
  5. from torch import nn
  7. class NeuralNetwork(nn.Module):
  8.     def __init__(self):
  9.         super().__init__()
  10.         self.linear = nn.Sequential(
  11.             nn.Linear(input_size, 16, bias=False),
  12.             nn.Linear(16, output_size, bias=False),
  13.         )
  15.     def forward(self, x):
  16.         return self.linear(x)
  17.     
  18. device = "cpu"
  19. input_size = 2
  20. output_size = 2
  22. x = torch.randn(10, 1, input_size).to(device)
  23. y = torch.randn(10, 1, output_size).to(device)
  25. model = NeuralNetwork().to(device)
  26. loss_fn = nn.MSELoss()
  28. def loss_grad(x, label):
  29.     y = model(x)
  30.     loss = loss_fn(y, label)
  31.     grads = torch.autograd.grad(loss, model.parameters(), retain_graph=True)
  32.     return grads

以下部分有效，但使用了for循环：

  1. # 有效但效率低下
  2. def compute_for():
  3.     grads = [loss_grad(x[i], y[i]) for i in range(x.shape[0])]
  4.     print(grads)
  6. compute_for()

为了提高效率，我尝试使用torch.vmap：

  1. # 可能更高效但不起作用
  2. def compute_vmap():
  3.     grads = torch.vmap(loss_grad)(x, y)
  4.     print(grads)
  6. compute_vmap()

我原本希望它计算每个元素在x, y中关于参数的损失梯度。但是，我收到了错误消息：

  1. RuntimeError: element 0 of tensors does not require grad

据我了解，这意味着来自张量x的元素将被计算，它们不需要独立地梯度。

我该如何修改这段代码以计算所有梯度？或者是否有其他方法可以做到这一点？

  2. <details>
  3. <summary>英文:</summary>
  5. I am trying to compute multiple loss gradients efficiently (without a for loop) in PyTorch. Given:
  7. ```python
  8. import torch
  9. from torch import nn
  11. class NeuralNetwork(nn.Module):
  12.     def __init__(self):
  13.         super().__init__()
  14.         self.linear = nn.Sequential(
  15.             nn.Linear(input_size, 16, bias=False),
  16.             nn.Linear(16, output_size, bias=False),
  17.         )
  19.     def forward(self, x):
  20.         return self.linear(x)
  21.     
  22. device = &quot;cpu&quot;
  23. input_size = 2
  24. output_size = 2
  26. x = torch.randn(10, 1, input_size).to(device)
  27. y = torch.randn(10, 1, output_size).to(device)
  29. model = NeuralNetwork().to(device)
  30. loss_fn = nn.MSELoss()
  32. def loss_grad(x, label):
  33.     y = model(x)
  34.     loss = loss_fn(y, label)
  35.     grads = torch.autograd.grad(loss, model.parameters(), retain_graph=True)
  36.     return grads

The following works, but uses a for loop:

  1. # inefficient but works
  2. def compute_for():
  3.     grads = [loss_grad(x[i], y[i]) for i in range(x.shape[0])]
  4.     print(grads)
  6. compute_for()

For efficiency, I tried using torch.vmap instead:

  1. # potentially more efficient but doesn&#39;t work
  2. def compute_vmap():
  3.     grads = torch.vmap(loss_grad)(x, y)
  4.     print(grads)
  6. compute_vmap()

I was expecting it to compute the gradients of the losses w.r.t. the parameters for each element in x, y. Instead, I get an error:

  1. RuntimeError: element 0 of tensors does not require grad

As I understand, this means that elements from the tensor x will be computed and they don't individually require grad.

How can I modify this code so that it computes all gradients? Or is there another method to do that?

答案1

得分: 0

per-sample gradients 可以使用 vmap 计算，如相关的 教程 中所示：

  1. from torch.func import functional_call, vmap, grad
  3. def compute_loss(params, buffers, sample, target):
  4.     batch = sample.unsqueeze(0)
  5.     targets = target.unsqueeze(0)
  6.     predictions = functional_call(model, (params, buffers), (batch,))
  7.     loss = loss_fn(predictions, targets)
  8.     return loss
  10. params = {k: v.detach() for k, v in model.named_parameters()}
  11. buffers = {k: v.detach() for k, v in model.named_buffers()}
  13. ft_compute_grad = grad(compute_loss)
  14. ft_compute_sample_grad = vmap(ft_compute_grad, in_dims=(None, None, 0, 0))
  15. ft_per_sample_grads = ft_compute_sample_grad(params, buffers, x, y)
  17. print(ft_per_sample_grads)

这些与为每个对 (x[i], y[i]) 分别计算的梯度相匹配。

英文:

The per-sample gradients may be computed using vmap as shown in the relevant tutorial:

  1. from torch.func import functional_call, vmap, grad
  3. def compute_loss(params, buffers, sample, target):
  4.     batch = sample.unsqueeze(0)
  5.     targets = target.unsqueeze(0)
  6.     predictions = functional_call(model, (params, buffers), (batch,))
  7.     loss = loss_fn(predictions, targets)
  8.     return loss
  10. params = {k: v.detach() for k, v in model.named_parameters()}
  11. buffers = {k: v.detach() for k, v in model.named_buffers()}
  13. ft_compute_grad = grad(compute_loss)
  14. ft_compute_sample_grad = vmap(ft_compute_grad, in_dims=(None, None, 0, 0))
  15. ft_per_sample_grads = ft_compute_sample_grad(params, buffers, x, y)
  17. print(ft_per_sample_grads)

These match the gradients computed individually for each pair (x[i], y[i]).

huangapple
  • 本文由 发表于 2023年5月29日 04:49:11
  • 转载请务必保留本文链接：https://go.coder-hub.com/76353546.html
  • python
  • pytorch
  • vectorization
匿名

发表评论

匿名网友

:?: :razz: :sad: :evil: :!: :smile: :oops: :grin: :eek: :shock: :???: :cool: :lol: :mad: :twisted: :roll: :wink: :idea: :arrow: :neutral: :cry: :mrgreen:

确定