[Fix] Fix udp regress (#1682)

open-mmlab · Oct 14, 2022 · a3b71ee · a3b71ee
1 parent af2af15
commit a3b71ee
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 15 deletions.
diff --git a/...y_2d_keypoint/topdown_heatmap/coco/td-hm_hrnet-w32_udp-regress-8xb64-210e_coco-256x192.py b/...y_2d_keypoint/topdown_heatmap/coco/td-hm_hrnet-w32_udp-regress-8xb64-210e_coco-256x192.py
@@ -83,8 +83,9 @@
         in_channels=32,
         out_channels=3 * 17,
         deconv_out_channels=None,
-        loss=dict(type='KeypointMSELoss', use_target_weight=True),
+        loss=dict(type='CombinedTargetMSELoss', use_target_weight=True),
         decoder=codec),
+    train_cfg=dict(compute_acc=False),
     test_cfg=dict(
         flip_test=True,
         flip_mode='udp_combined',

diff --git a/mmpose/codecs/utils/offset_heatmap.py b/mmpose/codecs/utils/offset_heatmap.py
@@ -23,7 +23,7 @@ def generate_offset_heatmap(
             (N, K)
         radius_factor (float): The radius factor of the binary label
             map. The positive region is defined as the neighbor of the
-            keypoit with the radius :math:`r=radius_factor*max(W, H)`
+            keypoint with the radius :math:`r=radius_factor*max(W, H)`
 
     Returns:
         tuple:

diff --git a/mmpose/models/heads/heatmap_heads/heatmap_head.py b/mmpose/models/heads/heatmap_heads/heatmap_head.py
@@ -332,13 +332,14 @@ def loss(self,
         losses.update(loss_kpt=loss)
 
         # calculate accuracy
-        _, avg_acc, _ = pose_pck_accuracy(
-            output=to_numpy(pred_fields),
-            target=to_numpy(gt_heatmaps),
-            mask=to_numpy(keypoint_weights) > 0)
-
-        acc_pose = torch.tensor(avg_acc, device=gt_heatmaps.device)
-        losses.update(acc_pose=acc_pose)
+        if train_cfg.get('compute_acc', True):
+            _, avg_acc, _ = pose_pck_accuracy(
+                output=to_numpy(pred_fields),
+                target=to_numpy(gt_heatmaps),
+                mask=to_numpy(keypoint_weights) > 0)
+
+            acc_pose = torch.tensor(avg_acc, device=gt_heatmaps.device)
+            losses.update(acc_pose=acc_pose)
 
         return losses
 

diff --git a/mmpose/models/losses/__init__.py b/mmpose/models/losses/__init__.py
@@ -2,16 +2,17 @@
 from .classification_loss import BCELoss, JSDiscretLoss, KLDiscretLoss
 from .heatmap_loss import AdaptiveWingLoss
 from .loss_wrappers import MultipleLossWrapper
-from .mse_loss import KeypointMSELoss, KeypointOHKMMSELoss
+from .mse_loss import (CombinedTargetMSELoss, KeypointMSELoss,
+                       KeypointOHKMMSELoss)
 from .multi_loss_factory import AELoss, HeatmapLoss, MultiLossFactory
 from .regression_loss import (BoneLoss, L1Loss, MPJPELoss, MSELoss, RLELoss,
                               SemiSupervisionLoss, SmoothL1Loss, SoftWingLoss,
                               WingLoss)
 
 __all__ = [
-    'KeypointMSELoss', 'KeypointOHKMMSELoss', 'HeatmapLoss', 'AELoss',
-    'MultiLossFactory', 'SmoothL1Loss', 'WingLoss', 'MPJPELoss', 'MSELoss',
-    'L1Loss', 'BCELoss', 'BoneLoss', 'SemiSupervisionLoss', 'SoftWingLoss',
-    'AdaptiveWingLoss', 'RLELoss', 'KLDiscretLoss', 'MultipleLossWrapper',
-    'JSDiscretLoss'
+    'KeypointMSELoss', 'KeypointOHKMMSELoss', 'CombinedTargetMSELoss',
+    'HeatmapLoss', 'AELoss', 'MultiLossFactory', 'SmoothL1Loss', 'WingLoss',
+    'MPJPELoss', 'MSELoss', 'L1Loss', 'BCELoss', 'BoneLoss',
+    'SemiSupervisionLoss', 'SoftWingLoss', 'AdaptiveWingLoss', 'RLELoss',
+    'KLDiscretLoss', 'MultipleLossWrapper', 'JSDiscretLoss'
 ]