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[:,[1]]
    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