I have to implement a loss in backward of convolution layer as illustrated in below code. But I don’t know how can I get grad_weight = grad_weight + cont_loss_weight such that the shape of cont_loss_weight should be same as that of grad_weight and grad_bias = grad_bias + cont_loss_bias such that the shape of cont_loss_bias should be same as that of grad_bias.
class Custom_Convolution(torch.autograd.Function):
@staticmethod
def forward(ctx, input, weight, bias, stride, padding): #input's shape = ([batch_size=100, 96, 8, 8])
output = torch.nn.functional.conv2d(input, weight, bias, stride, padding)
ctx.save_for_backward(input, weight, bias, output)
return output #output's shape = ([[batch_size= 100,128, 4, 4])
@staticmethod
def backward(ctx, grad_output): #grad_output size = ([batch_size= 100, 128, 4, 4])
input, weight, bias, output = ctx.saved_tensors #input's size = ([batch_size=100 , 96,8,8])
features = output
features = features.view(features.shape[0], features.shape[1], -1)
features = torch.nn.functional.normalize(features,dim=1)
#Total_features= features.shape[0]* features.shape[1]
cont_loss = torch.tensor([0.]).requires_grad_(requires_grad=True).to(dev) # shape: ([1])
for ..... :
# My code for loss... includes some operations like torch.div,exp,sum...
# Calculation of loss for each feature 'i' (Li) out of Total_features
# Total loss = sum of all Li (Number of Li values = features.shape[0]* features.shape[1])
cont_loss = cont_loss.add(cont_loss_img) # Total contrastive loss of the batch
grad_input = grad_weight = grad_bias = None
if ctx.needs_input_grad[0]:
grad_input = torch.nn.grad.conv2d_input( input.shape, weight, grad_output) #shape = ([batch_size,96,8,8])
if ctx.needs_input_grad[1]:
grad_weight = torch.nn.grad.conv2d_weight(input, weight.shape, grad_output) #shape = ([128,96,5,5])
# cont_loss_weight = torch.nn.grad.conv2d_weight(cont_loss, weight.shape, grad_output)
# grad_weight = grad_weight + cont_loss_weight
if bias is not None and ctx.needs_input_grad[2]:
grad_bias = grad_output.sum((0,2,3)) #shape = ([128])
#grad_bias = grad_bias + cont_loss_bias
if bias is not None:
return grad_input, grad_weight, grad_bias, None, None
else:
return grad_input, grad_weight, None, None