Is the gradient of the sum equal to the sum of the gradients for a neural network in pytorch?

Viewed 514

Let's suppose I have the code below and I want to calculate the jacobian of L, which is the prediction made by a neural network in Pytorch, L is of size nx1 where n is the number of samples in a mini batch. In order to avoid a for loop for each entry of L (n entries) to calculate the jacobian for each sample in the mini batch some codes I found just sum the n predictions of the neural network (L) with respect with the inputs and then calculate the gradient of the sum. First I can't understand why is the gradient of the sum the same of the sum of the gradients for each sample in pytorch architecture. Second I tried both with the sum and with a for loop and the results diverge. Could it be due to numerical approximations or because the sum just doesn't make sense?

The code is below, where both functions belong to a nn.module:

def forward(self, x):
        with torch.set_grad_enabled(True):
            def function(x,t):
                 self.n = n = x.shape[1]//2

                 qqd = x.requires_grad_(True)
                 L = self._lagrangian(qqd).sum()
                 J = grad(L, qqd, create_graph=True)[0]

        
def _lagrangian(self, qqd):
    x = F.softplus(self.fc1(qqd))
    x = F.softplus(self.fc2(x))
    x = F.softplus(self.fc3(x))
    L = self.fc_last(x)
    return L
1 Answers

I think it should, this is just a toy example

w = torch.tensor([2.], requires_grad=True)
x1 = torch.tensor([3.], requires_grad=True)
x2 = torch.tensor([4.], requires_grad=True)
y = w * a + w * b
y.backward() # calculate gradient

return

>>> w.grad
tensor([7.])
Related