Hi Nikhil!



 Nikhil_Chhabra:

Input= [batch_size, channels, height,width]
predicted_mask = [batch_size, 1, height, width]
gt_mask = [batch_size, 1, height, width]


This is fine.


Assuming your data has 25% positive pixels and 75% negative pixels.
positive_weight = torch.ones([1])*3.0
loss = nn.BCEWithLogitsLoss(pos_weight=positive_weight)


This is fine.


My only confusion is do i need to have extra one dimension for class or it doesn’t matter ?
predicted_mask = [batch_size, 1, height, width] OR  [batch_size, height, width]
gt_mask = [batch_size, 1, height, width] OR  [batch_size, height, width]


I would call the “extra one dimension” (singleton dimension) a “channel”
dimension (rather than “class”).
Regardless of terminology, as long as the shapes of predicted_mask
and gt_mask match, both ways will work and are essentially equivalent.
If the output of the model and gt_mask (from, say, your dataloader) already
match, I wouldn’t bother doing anything.  If one has the singleton dimension
and the other doesn’t, I would probably squeeze() the extra dimension away.
Best.
K. Frank