PyTorch中的高效批量导数操作

Question

我正在使用 Pytorch 来实现具有（比如说）5 个输入和 2 个输出的神经网络

class myNetwork(nn.Module):
   def __init__(self):
      super(myNetwork,self).__init__()
      self.layer1 = nn.Linear(5,32)
      self.layer2 = nn.Linear(32,2)
   def forward(self,x):
      x = torch.relu(self.layer1(x))
      x = self.layer2(x)
      return x

显然，我可以输入一个 (N x 5) 张量并得到 (N x 2) 结果，

net = myNetwork()
nbatch = 100
inp = torch.rand([nbatch,5])
inp.requires_grad = True
out = net(inp)

我现在想为批处理中的每个示例计算 NN output 关于输入向量的一个元素（假设是第 5 个元素）的导数。 我知道我可以使用torch.autograd.grad计算 output 的一个元素相对于所有输入的导数，我可以按如下方式使用它：

deriv = torch.zeros([nbatch,2])
for i in range(nbatch):
   for j in range(2):
      deriv[i,j] = torch.autograd.grad(out[i,j],inp,retain_graph=True)[0][i,4]

然而，这似乎非常低效：它计算out[i,j]关于批次中每个元素的梯度，然后丢弃除一个之外的所有元素。 有一个更好的方法吗？

Answer 1

凭借反向传播，如果您只计算梯度 w.r.ta 单个输入，则计算节省不一定很多，您只会在第一层节省一些，之后的所有层都需要反向传播.

因此，这可能不是最佳方式，但实际上并不会产生太多开销，尤其是在您的网络有很多层的情况下。

顺便说一句，是否有理由需要循环nbatch ？ 如果您想要批次 w.r.ta 参数的每个元素的梯度，我可以理解，因为 pytorch 会将它们放在一起，但您似乎只对输入感兴趣......