Merge pull request #101 from hzdr-MedImaging/feature-eval-activation

added final activation before eval score calculation
wolny · Dec 4, 2023 · b70a54e · b70a54e
2 parents 85ae308 + d3f16ea
commit b70a54e
Showing 1 changed file with 11 additions and 1 deletion.
diff --git a/pytorch3dunet/unet3d/trainer.py b/pytorch3dunet/unet3d/trainer.py
@@ -209,7 +209,17 @@ def train(self):
             if self.num_iterations % self.log_after_iters == 0:
                 # compute eval criterion
                 if not self.skip_train_validation:
-                    eval_score = self.eval_criterion(output, target)
+                    # apply final activation before calculating eval score
+                    if isinstance(self.model, nn.DataParallel):
+                        final_activation = self.model.module.final_activation
+                    else:
+                        final_activation = self.model.final_activation
+
+                    if final_activation is not None:
+                        act_output = final_activation(output)
+                    else:
+                        act_output = output
+                    eval_score = self.eval_criterion(act_output, target)
                     train_eval_scores.update(eval_score.item(), self._batch_size(input))
 
                 # log stats, params and images