When you do segmentation you will usually do pixel-wise MSE or something as a loss function.
If you generate masks with this “Unlabeled” class you can adjust your loss functions as follows:
Assuming you have your predictions p
(prediction) and segmentation mask t
(target) you can find
t != UNLABELED
(whatever value this unlabeled class has) which is a tensor containing booleans.
You can then either index your p
and t
tensors using this resulting tensor as a mask for both tensors or just do
hasLabel = (t != UNLABELED).float()
loss = mse(p * hasLabel, t * hasLabel)
which will basically replace every pixel in your prediction and label with 0
, leading to no loss contribution. Using indexes is probably faster, but you will have to look for the appropriate function in the pytorch documentation.
Note that you do NOT want to remove stuff from your input because of that, because your convolutions will need to take the unlabeled parts into consideration.