Hi,
for semantic segmentation I wouldn’t use cross entropy. See Sudre et al. 29 (see also Crum et al. 12 )
that has a generalized dice coefficient, weighted for class imbalance. It’s a good start. There exist TF implementation of the generalized dice coefficient, that you can easily port in pytorch, here 62.