From 92fd1c77164ddd2e43df9cb2fa5ec07d839a8f86 Mon Sep 17 00:00:00 2001
From: Ezra-Yu <1105212286@qq.com>
Date: Tue, 3 Aug 2021 17:56:11 +0800
Subject: [PATCH] Support default target_weight as None in losses (#829)

---
 mmpose/models/losses/classfication_loss.py |  3 ++-
 mmpose/models/losses/regression_loss.py    | 18 ++++++++++++------
 2 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/mmpose/models/losses/classfication_loss.py b/mmpose/models/losses/classfication_loss.py
index ba87306cb7..35fc0930dc 100644
--- a/mmpose/models/losses/classfication_loss.py
+++ b/mmpose/models/losses/classfication_loss.py
@@ -14,7 +14,7 @@ def __init__(self, use_target_weight=False, loss_weight=1.):
         self.use_target_weight = use_target_weight
         self.loss_weight = loss_weight
 
-    def forward(self, output, target, target_weight):
+    def forward(self, output, target, target_weight=None):
         """Forward function.
 
         Note:
@@ -29,6 +29,7 @@ def forward(self, output, target, target_weight):
         """
 
         if self.use_target_weight:
+            assert target_weight is not None
             loss = self.criterion(output, target, reduction='none')
             if target_weight.dim() == 1:
                 target_weight = target_weight[:, None]
diff --git a/mmpose/models/losses/regression_loss.py b/mmpose/models/losses/regression_loss.py
index 599824ab21..6160ea24ce 100644
--- a/mmpose/models/losses/regression_loss.py
+++ b/mmpose/models/losses/regression_loss.py
@@ -23,7 +23,7 @@ def __init__(self, use_target_weight=False, loss_weight=1.):
         self.use_target_weight = use_target_weight
         self.loss_weight = loss_weight
 
-    def forward(self, output, target, target_weight):
+    def forward(self, output, target, target_weight=None):
         """Forward function.
 
         Note:
@@ -38,6 +38,7 @@ def forward(self, output, target, target_weight):
                 Weights across different joint types.
         """
         if self.use_target_weight:
+            assert target_weight is not None
             loss = self.criterion(output * target_weight,
                                   target * target_weight)
         else:
@@ -91,7 +92,7 @@ def criterion(self, pred, target):
             self.omega * torch.log(1.0 + delta / self.epsilon), delta - self.C)
         return torch.mean(torch.sum(losses, dim=[1, 2]), dim=0)
 
-    def forward(self, output, target, target_weight):
+    def forward(self, output, target, target_weight=None):
         """Forward function.
 
         Note:
@@ -106,6 +107,7 @@ def forward(self, output, target, target_weight):
                 Weights across different joint types.
         """
         if self.use_target_weight:
+            assert target_weight is not None
             loss = self.criterion(output * target_weight,
                                   target * target_weight)
         else:
@@ -129,7 +131,7 @@ def __init__(self, use_target_weight=False, loss_weight=1.):
         self.use_target_weight = use_target_weight
         self.loss_weight = loss_weight
 
-    def forward(self, output, target, target_weight):
+    def forward(self, output, target, target_weight=None):
         """Forward function.
 
         Note:
@@ -145,6 +147,7 @@ def forward(self, output, target, target_weight):
         """
 
         if self.use_target_weight:
+            assert target_weight is not None
             loss = torch.mean(
                 torch.norm((output - target) * target_weight, dim=-1))
         else:
@@ -163,7 +166,7 @@ def __init__(self, use_target_weight=False, loss_weight=1.):
         self.use_target_weight = use_target_weight
         self.loss_weight = loss_weight
 
-    def forward(self, output, target, target_weight):
+    def forward(self, output, target, target_weight=None):
         """Forward function.
 
         Note:
@@ -177,6 +180,7 @@ def forward(self, output, target, target_weight):
                 Weights across different joint types.
         """
         if self.use_target_weight:
+            assert target_weight is not None
             loss = self.criterion(output * target_weight,
                                   target * target_weight)
         else:
@@ -195,7 +199,7 @@ def __init__(self, use_target_weight=False, loss_weight=1.):
         self.use_target_weight = use_target_weight
         self.loss_weight = loss_weight
 
-    def forward(self, output, target, target_weight):
+    def forward(self, output, target, target_weight=None):
         """Forward function.
 
         Note:
@@ -209,6 +213,7 @@ def forward(self, output, target, target_weight):
                 Weights across different joint types.
         """
         if self.use_target_weight:
+            assert target_weight is not None
             loss = self.criterion(output * target_weight,
                                   target * target_weight)
         else:
@@ -239,7 +244,7 @@ def __init__(self, joint_parents, use_target_weight=False, loss_weight=1.):
             if i != self.joint_parents[i]:
                 self.non_root_indices.append(i)
 
-    def forward(self, output, target, target_weight):
+    def forward(self, output, target, target_weight=None):
         """Forward function.
 
         Note:
@@ -260,6 +265,7 @@ def forward(self, output, target, target_weight):
             target - target[:, self.joint_parents, :],
             dim=-1)[:, self.non_root_indices]
         if self.use_target_weight:
+            assert target_weight is not None
             loss = torch.mean(
                 torch.abs((output_bone * target_weight).mean(dim=0) -
                           (target_bone * target_weight).mean(dim=0)))