fixbug: miss cls loss weight in center head (open-mmlab#772)

ajunlonglive · Jan 24, 2022 · 274c90c · 274c90c
1 parent fa78c4f
commit 274c90c
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/pcdet/models/dense_heads/center_head.py b/pcdet/models/dense_heads/center_head.py
@@ -232,6 +232,7 @@ def get_loss(self):
         for idx, pred_dict in enumerate(pred_dicts):
             pred_dict['hm'] = self.sigmoid(pred_dict['hm'])
             hm_loss = self.hm_loss_func(pred_dict['hm'], target_dicts['heatmaps'][idx])
+            hm_loss *= self.model_cfg.LOSS_CONFIG.LOSS_WEIGHTS['cls_weight']
 
             target_boxes = target_dicts['target_boxes'][idx]
             pred_boxes = torch.cat([pred_dict[head_name] for head_name in self.separate_head_cfg.HEAD_ORDER], dim=1)