Support default target_weight as None in losses (#829)

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]

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)))