Merge pull request #177 from Zehaos/fix

Support class agnostic in mask_cross_entropy loss

Merge pull request #177 from Zehaos/fix
Support class agnostic in mask_cross_entropy loss
cd35200f · Kai Chen · GitHub · c1f9fa0f · b4af1c94 · cd35200f
Unverified Commit cd35200f authored Dec 12, 2018 by Kai Chen Committed by GitHub Dec 12, 2018
Hide whitespace changes
Inline Side-by-side

Showing with 5 additions and 1 deletion

mmdet/models/mask_heads/fcn_mask_head.py mmdet/models/mask_heads/fcn_mask_head.py +5 -1

No files found.
--- a/mmdet/models/mask_heads/fcn_mask_head.py
+++ b/mmdet/models/mask_heads/fcn_mask_head.py
@@ -97,7 +97,11 @@ class FCNMaskHead(nn.Module):

    def loss(self, mask_pred, mask_targets, labels):
        loss = dict()
-        loss_mask = mask_cross_entropy(mask_pred, mask_targets, labels)
+        if self.class_agnostic:
+            loss_mask = mask_cross_entropy(mask_pred, mask_targets,
+                                           torch.zeros_like(labels))
+        else:
+            loss_mask = mask_cross_entropy(mask_pred, mask_targets, labels)
        loss['loss_mask'] = loss_mask
        return loss