Gradient of tensor wrt input tensor during batch processing

Hi,

I have the following data

input: torch.Size([10000, 1, 28, 28])
output: torch.Size([10000, 10])

The input is from the MNIST data set and the output is the tensor consisting of the output from the classification of all this MNIST data.

I cannot figure out how to calculate the Jacobian for each of these outputs wrt to the input that created the output.

So the Jacobian of the first output in the output tensor would need to be calculated with respect to the first input image. Second Jacobian of the second output with respect to the second input image, and so on.

I can easily do this by looping over the whole data structure, but I am unsure if this is possible in one go with the backward method. But I’d rather not have to do this in raw Python code if the backward method is capable of accommodating this.

I think the following should work.

outputs.backward(torch.ones_like(outputs))

But my grads on the input tensor seem to stay None.

This is the code I am using.

net = Net()
net.load_state_dict(torch.load(PATH))

testset = torchvision.datasets.MNIST(root='./data',
                                     train=False,
                                     download=True,
                                     transform=transforms.ToTensor())
testloader = torch.utils.data.DataLoader(testset,
                                         batch_size=10000,
                                         shuffle=True)

optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

correct = 0
total = 0
for data in testloader:
    images, labels = data
    outputs = net(images)
    _, predicted = torch.max(outputs.data, 1)
    total += labels.size(0)
    correct += (predicted == labels).sum().item()
    optimizer.zero_grad()
    outputs.backward(torch.ones_like(outputs), retain_graph=True)
    print(images.grad)

Okay so I figured out that the input images were not set to require gradient, I enabled that and get out a gradient now.

But I am wondering why the following code results in two different gradients.

optimizer.zero_grad()
outputs.backward(torch.ones_like(outputs), retain_graph=True)
print(images.grad[0][0][0])
optimizer.zero_grad()
outputs[0][0].backward(retain_graph=True)
print(images.grad[0][0][0])

This results in the output

tensor([ 0.0035,  0.0133,  0.0018,  0.0005, -0.0046, -0.0143,  0.0113,  0.0151,
         0.0017,  0.0022, -0.0222,  0.0082, -0.0164, -0.0168, -0.0114,  0.0023,
        -0.0067,  0.0034, -0.0148,  0.0263,  0.0039,  0.0343,  0.0054,  0.0106,
        -0.0175, -0.0160,  0.0196, -0.0109])
tensor([ 0.0044,  0.0073,  0.0097, -0.0619, -0.0163, -0.0065,  0.0255, -0.0144,
         0.0138,  0.0228, -0.0363, -0.0176,  0.0077, -0.0022,  0.0227, -0.0188,
         0.0081,  0.0045, -0.0278,  0.0578,  0.0228,  0.0451,  0.0041,  0.0102,
         0.0040, -0.0396,  0.0285, -0.0171])

Does it calculate a gradient for my outputs based on some inputs that is unrelated to it?
I would assume the last output is the correct one. I have a feeling that it’s the torch.ones_like(outputs) messing it up here, but I am unsure of what it should be otherwise.