My task is to create a single object detection model, It is doing 2 task: 1. Classification between 2 classes [cat, no object]
2: Regression x and y co_ordinates.
So I want to design a loss function such that it gives me 2 losses for classification and regression each.
for classification I am using CrossEntropy Loss, but for regression I want to calculate the MSE loss only for the images containing cat. So here is a dummy example of such loss function I have created please check if it is right or wrong.
output = torch.tensor([[0.7,0.3,0.5,0.6], #(prob_cat, prob_no_object, x,y) [0.2,0.8,100,100], [0.8,0.2,0.4,0.5]]) target = torch.tensor([[1,0,0.4,0.8], [0,1,0,0], [1,0,0.3,0.4]]) def myloss(output,target): #classification loss (cat, no object) target_class = target[:,:2].argmax(1) #True class output_class = output[:,:2] #Predicted Class l1 = nn.CrossEntropyLoss(reduction='sum') Loss1 = l1(output_class,target_class) #regression loss (x,y) prob_no_object = target[:,] prob_object=(prob_no_object==0).float() target_bbx = target[:,[2,3]] output_bbx = output[:,[2,3]] l2 = nn.MSELoss(reduction='none') L2 = l2(output_bbx,target_bbx) # print(L2.shape) Loss2 = torch.sum(torch.mul(prob_object,L2)) return(Loss1,Loss2/torch.sum(prob_object)) #classification loss, regression_loss_per_cat_img_in_batch