[Fix] fix loss computation in MSPNHead (#2993)

open-mmlab · Mar 26, 2024 · e0eb5d4 · e0eb5d4
1 parent de67839
commit e0eb5d4
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 1 deletion.
diff --git a/mmpose/models/heads/heatmap_heads/mspn_head.py b/mmpose/models/heads/heatmap_heads/mspn_head.py
@@ -394,7 +394,9 @@ def loss(self,
 
         keypoint_weights = torch.cat([
             d.gt_instance_labels.keypoint_weights for d in batch_data_samples
-        ])  # shape: [B*N, L, K]
+        ],
+                                     dim=1)
+        keypoint_weights = keypoint_weights.transpose(0, 1)  # [B*N, L, K]
 
         # calculate losses over multiple stages and multiple units
         losses = dict()

diff --git a/tests/test_models/test_heads/test_heatmap_heads/test_mspn_head.py b/tests/test_models/test_heads/test_heatmap_heads/test_mspn_head.py
@@ -44,6 +44,7 @@ def _get_data_samples(self,
             with_heatmap=True,
             with_reg_label=False,
             num_levels=num_levels)['data_samples']
+
         return batch_data_samples
 
     def test_init(self):
@@ -153,6 +154,10 @@ def test_loss(self):
                          (unit_channels, 32, 24), (unit_channels, 64, 48)])
         batch_data_samples = self._get_data_samples(
             batch_size=2, heatmap_size=(48, 64), num_levels=4)
+        for ds in batch_data_samples:
+            ds.gt_instance_labels = InstanceData(
+                keypoint_weights=ds.gt_instance_labels.keypoint_weights.
+                transpose(0, 1))
         losses = head.loss(feats, batch_data_samples)
 
         self.assertIsInstance(losses['loss_kpt'], torch.Tensor)
@@ -189,6 +194,10 @@ def test_loss(self):
                          (unit_channels, 32, 24), (unit_channels, 64, 48)])
         batch_data_samples = self._get_data_samples(
             batch_size=2, heatmap_size=(48, 64), num_levels=16)
+        for ds in batch_data_samples:
+            ds.gt_instance_labels = InstanceData(
+                keypoint_weights=ds.gt_instance_labels.keypoint_weights.
+                transpose(0, 1))
         losses = head.loss(feats, batch_data_samples)
 
         self.assertIsInstance(losses['loss_kpt'], torch.Tensor)