Remov Argmax Computation for torchmetrics in Classification and Segme…

…ntation (#1777) * remove y_hat_hard * argmax vs softmax :)
microsoft · Mar 2, 2024 · f713ba8 · f713ba8
1 parent b2369d6
commit f713ba8
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 14 deletions.
diff --git a/torchgeo/trainers/classification.py b/torchgeo/trainers/classification.py
@@ -161,10 +161,9 @@ def training_step(
         x = batch["image"]
         y = batch["label"]
         y_hat = self(x)
-        y_hat_hard = y_hat.argmax(dim=1)
         loss: Tensor = self.criterion(y_hat, y)
         self.log("train_loss", loss)
-        self.train_metrics(y_hat_hard, y)
+        self.train_metrics(y_hat, y)
         self.log_dict(self.train_metrics)
 
         return loss
@@ -182,10 +181,9 @@ def validation_step(
         x = batch["image"]
         y = batch["label"]
         y_hat = self(x)
-        y_hat_hard = y_hat.argmax(dim=1)
         loss = self.criterion(y_hat, y)
         self.log("val_loss", loss)
-        self.val_metrics(y_hat_hard, y)
+        self.val_metrics(y_hat, y)
         self.log_dict(self.val_metrics)
 
         if (
@@ -198,7 +196,7 @@ def validation_step(
         ):
             try:
                 datamodule = self.trainer.datamodule
-                batch["prediction"] = y_hat_hard
+                batch["prediction"] = y_hat.argmax(dim=-1)
                 for key in ["image", "label", "prediction"]:
                     batch[key] = batch[key].cpu()
                 sample = unbind_samples(batch)[0]
@@ -223,10 +221,9 @@ def test_step(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> None
         x = batch["image"]
         y = batch["label"]
         y_hat = self(x)
-        y_hat_hard = y_hat.argmax(dim=1)
         loss = self.criterion(y_hat, y)
         self.log("test_loss", loss)
-        self.test_metrics(y_hat_hard, y)
+        self.test_metrics(y_hat, y)
         self.log_dict(self.test_metrics)
 
     def predict_step(

diff --git a/torchgeo/trainers/segmentation.py b/torchgeo/trainers/segmentation.py
@@ -217,10 +217,9 @@ def training_step(
         x = batch["image"]
         y = batch["mask"]
         y_hat = self(x)
-        y_hat_hard = y_hat.argmax(dim=1)
         loss: Tensor = self.criterion(y_hat, y)
         self.log("train_loss", loss)
-        self.train_metrics(y_hat_hard, y)
+        self.train_metrics(y_hat, y)
         self.log_dict(self.train_metrics)
         return loss
 
@@ -237,10 +236,9 @@ def validation_step(
         x = batch["image"]
         y = batch["mask"]
         y_hat = self(x)
-        y_hat_hard = y_hat.argmax(dim=1)
         loss = self.criterion(y_hat, y)
         self.log("val_loss", loss)
-        self.val_metrics(y_hat_hard, y)
+        self.val_metrics(y_hat, y)
         self.log_dict(self.val_metrics)
 
         if (
@@ -253,7 +251,7 @@ def validation_step(
         ):
             try:
                 datamodule = self.trainer.datamodule
-                batch["prediction"] = y_hat_hard
+                batch["prediction"] = y_hat.argmax(dim=1)
                 for key in ["image", "mask", "prediction"]:
                     batch[key] = batch[key].cpu()
                 sample = unbind_samples(batch)[0]
@@ -278,10 +276,9 @@ def test_step(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> None
         x = batch["image"]
         y = batch["mask"]
         y_hat = self(x)
-        y_hat_hard = y_hat.argmax(dim=1)
         loss = self.criterion(y_hat, y)
         self.log("test_loss", loss)
-        self.test_metrics(y_hat_hard, y)
+        self.test_metrics(y_hat, y)
         self.log_dict(self.test_metrics)
 
     def predict_step(