Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Remove Argmax Computation for torchmetrics in Classification and Segmentation #1777

Merged
merged 2 commits into from
Dec 15, 2023
Merged

Remove Argmax Computation for torchmetrics in Classification and Segmentation #1777

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
d2b8258
remove y_hat_hard
nilsleh Dec 15, 2023
e60cd10
argmax vs softmax :)
nilsleh Dec 15, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 4 additions & 7 deletions torchgeo/trainers/classification.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -162,10 +162,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
Expand All @@ -183,10 +182,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 (
Expand All @@ -198,7 +196,7 @@ def validation_step(
and hasattr(self.logger.experiment, "add_figure")
):
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]
Expand Down Expand Up @@ -227,10 +225,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(
Expand Down
11 changes: 4 additions & 7 deletions torchgeo/trainers/segmentation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -218,10 +218,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

Expand All @@ -238,10 +237,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 (
Expand All @@ -253,7 +251,7 @@ def validation_step(
and hasattr(self.logger.experiment, "add_figure")
):
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]
Expand Down Expand Up @@ -282,10 +280,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(
Expand Down