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