Custom Loss for single object detection task

1

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
2

The code looks generally alright. The output should contain raw logits, not probabilities for the nn.CrossEntropyLoss, but I assume you just used these values for easier debugging.

If you are concerned about the indexing, you could just use separate target tensors (and potentially different output layers in your model) for both tasks.

1 Like