Hello  For binary segmentation with high class imbalance, I believe the Dice loss or Intersection over Union loss are better than the Cross Entropy.
Here is an implementation of both

  
      github.com
  
  
    francois-rozet/adopptrs/blob/4b680f72748033ffc37ab67926f5444b243e4922/python/criterions.py#L57-L85

class DiceLoss(nn.Module):	'''Dice Loss (F-score, ...)'''	def __init__(self, smooth=1.):		super().__init__()		self.smooth = smooth	def forward(self, outputs, targets):		inter = (outputs * targets).sum()		dice = (2. * inter + self.smooth) / (outputs.sum() + targets.sum() + self.smooth)		return 1. - diceclass IOULoss(nn.Module):	'''Intersection Over Union Loss'''	def __init__(self, smooth=1.):		super().__init__()


  This file has been truncated. show original

  
  
    
    
  
  


Note that the mask should be 0 or 1 and the prediction a single number between 0 and 1 (not two like cross entropy).