add transpose for loss computation

dygraph/models/hrnet.py

@@ -179,6 +179,8 @@ class HRNet(fluid.dygraph.Layer):
             return pred, score_map
 
     def _get_loss(self, logit, label):
+        logit = fluid.layers.transpose(logit, [0, 2, 3, 1])
+        label = fluid.layers.transpose(label, [0, 2, 3, 1])
         mask = label != self.ignore_index
         mask = fluid.layers.cast(mask, 'float32')
         loss, probs = fluid.layers.softmax_with_cross_entropy(
@@ -186,7 +188,7 @@ class HRNet(fluid.dygraph.Layer):
             label,
             ignore_index=self.ignore_index,
             return_softmax=True,
-            axis=1)
+            axis=-1)
 
         loss = loss * mask
         avg_loss = fluid.layers.mean(loss) / (

dygraph/models/unet.py

@@ -43,6 +43,8 @@ class UNet(fluid.dygraph.Layer):
             return pred, score_map
 
     def _get_loss(self, logit, label):
+        logit = fluid.layers.transpose(logit, [0, 2, 3, 1])
+        label = fluid.layers.transpose(label, [0, 2, 3, 1])
         mask = label != self.ignore_index
         mask = fluid.layers.cast(mask, 'float32')
         loss, probs = fluid.layers.softmax_with_cross_entropy(