fix: Fixed binary classification evaluation (#228)

* fix: Fixed binary classification evaluation * style: Fixed style * fix: Fixed evaluate of binary classification
frgfm · Jul 16, 2022 · bc0d972 · bc0d972
1 parent 9871899
commit bc0d972
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/holocron/trainer/classification.py b/holocron/trainer/classification.py
@@ -118,7 +118,7 @@ def evaluate(self) -> Dict[str, float]:
 
         self.model.eval()
 
-        val_loss, top1, num_samples, num_valid_batches = 0.0, 0, 0, 0
+        val_loss, top1, num_samples, num_valid_batches = 0.0, 0.0, 0, 0
         for x, target in self.val_loader:
             x, target = self.to_cuda(x, target)
 
@@ -129,7 +129,7 @@ def evaluate(self) -> Dict[str, float]:
                 val_loss += _loss.item()
                 num_valid_batches += 1
 
-            top1 += int(torch.sum((target >= 0.5) == (torch.sigmoid(out) >= 0.5)).item())
+            top1 += torch.sum((target.view_as(out) >= 0.5) == (torch.sigmoid(out) >= 0.5)).item() / out[0].numel()
 
             num_samples += x.shape[0]