Conditional loss for multi-class multi-output model

1

Hello all!

I am working on a room-image classifier, which for a given image predicts: (1) what type of room it is (e.g. bedroom); (2) what style does the room have (e.g. modern); (3) does it look high or low budget.

I want to add a category for an “empty room” in the room types, but empty rooms should not have a style or budget associated with them. Therefore, in those cases where the room label is “empty”, I want to ignore the Cross Entropy losses associated with the style and budget. However, I cannot figure out how to do this for the entire batch at once.

For clarity’s sake, the losses (ignoring the condition) are computed as follows (based on one-hot encoding):

loss_r = nn.CrossEntropyLoss()(r_p, r_t)
loss_s = nn.CrossEntropyLoss()(s_p, s_t)
loss_b = nn.CrossEntropyLoss()(b_p, b_t)

where r_p, s_p, and b_p are the predictions (in logits), and r_t, s_t, and b_t are the targets.

Cheers!

2

Hi Tal,

when you load your data, on top of creating the one hot encoded classes, create an additional tensor that will mask your loss. If your room is empty, this tensor will be [1, 0, 0], whilst if it is not empty, it will be like [1, 1, 1].

When you compute your final loss (which is probably the sum of loss_r, loss_s, loss_b), just perform the internal (dot) product with the mask I told you about.

In this case you will have the following situations:

# empty room, mask = [1, 0, 0]
loss = torch.dot(torch.stack([loss_r, loss_s, loss_b], 0), mask) # loss_r * 1 + loss_s * 0 + loss_b * 0 = loss_r

# non-empty room, mask = [1, 1, 1]
loss = torch.dot(torch.stack([loss_r, loss_s, loss_b], 0), mask) # loss_r * 1 + loss_s * 1 + loss_b * 1

Therefore, when you have an empty room the loss will ignore the contribution from the style and budget.

1 Like