fix inconsistent in loss function and training

GingerHuang · Aug 14, 2019 · 7a1b216 · 7a1b216
1 parent 3b33b0a
commit 7a1b216
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 5 deletions.
diff --git a/model/loss.py b/model/loss.py
@@ -120,10 +120,10 @@ def forward(self, logits, target,
                 use_hierar=False,
                 is_multi=False,
                 *argvs):
+        device = logits.device
         if use_hierar:
             assert self.loss_type in [LossType.BCE_WITH_LOGITS,
                                       LossType.SIGMOID_FOCAL_CROSS_ENTROPY]
-            device = logits.device
             if not is_multi:
                 target = torch.eye(self.label_size)[target].to(device)
             hierar_penalty, hierar_paras, hierar_relations = argvs[0:3]
@@ -132,9 +132,13 @@ def forward(self, logits, target,
                                                                   hierar_relations,
                                                                   device)
         else:
-            if not is_multi:
-                device = logits.device
-                target = torch.eye(self.label_size)[target].to(device)
+            if is_multi:
+                assert self.loss_type in [LossType.BCE_WITH_LOGITS,
+                                        LossType.SIGMOID_FOCAL_CROSS_ENTROPY]
+            else:
+                if self.loss_type not in [LossType.SOFTMAX_CROSS_ENTROPY,
+                                          LossType.SOFTMAX_FOCAL_CROSS_ENTROPY]
+                    target = torch.eye(self.label_size)[target].to(device)
             return self.criterion(logits, target)
 
     def cal_recursive_regularize(self, paras, hierar_relations, device="cpu"):

diff --git a/train.py b/train.py
@@ -130,7 +130,9 @@ def run(self, data_loader, model, optimizer, stage,
             else:  # flat classification
                 loss = self.loss_fn(
                     logits,
-                    batch[ClassificationDataset.DOC_LABEL].to(self.conf.device))
+                    batch[ClassificationDataset.DOC_LABEL].to(self.conf.device),
+                    False,
+                    is_multi)
             if mode == ModeType.TRAIN:
                 optimizer.zero_grad()
                 loss.backward()