Compute pixel wise gradient of loss function - most efficient way?

autograd
1

Hey everyone,

I have a fully convolutional segmentation model and I am computing the binary cross entropy loss at every pixel. Now I want to get the gradient of each of those losses wrt. to the weights in the network to process them further. At the moment I am iterating over every pixel, basically doing the following:

import torch.nn.functional as F
output = net(input)
for i in range(imgheight):
    for j in range(imwidth):
        net.zero_grad()
        loss = F.binary_cross_entropy(output[:,i,j], label[:,i,j])
        loss.backward(retain_graph=True)
        grad=[]
        for param in net.parameters():
            if param.grad is not None:
                grad.append(param.grad)
        # further processing of the gradient

Is there a more efficient way of doing this?

Thank you for any suggestions