Эффективные пакетные производные операции в PyTorch - PullRequest
Новая
       4

Эффективные пакетные производные операции в PyTorch

0 голосов
/ 19 июня 2020

Я использую 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 по одному элементу входного вектора (скажем, 5-му элементу) для каждого примера в пакете. Я знаю, что могу вычислить производные одного элемента вывода относительно всех входов, используя torch.autograd.grad, и я мог бы использовать это следующим образом: 

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] по отношению к каждому элементу в пакете, а затем отбрасывает все, кроме одного. Есть ли лучший способ сделать это?

1 Ответ

1 голос
/ 19 июня 2020
• 1000 .

Так что, возможно, это не оптимальный способ, но на самом деле он не создает больших накладных расходов, особенно если в вашей сети много уровней.

Кстати, есть ли причина, по которой вы нужно л oop больше nbatch? Если вам нужен градиент каждого элемента пакета по параметру, я мог бы это понять, потому что pytorch объединяет их вместе, но, похоже, вас интересует исключительно ввод ...

...