BaseTask: explicitly pass batch size to logger

torchgeo/trainers/byol.py

@@ -364,6 +364,7 @@ def training_step(
             AssertionError: If channel dimensions are incorrect.
         """
         x = batch["image"]
+        batch_size = x.shape[0]
 
         in_channels = self.hparams["in_channels"]
         assert x.size(1) == in_channels or x.size(1) == 2 * in_channels
@@ -387,7 +388,7 @@ def training_step(
 
         loss = torch.mean(normalized_mse(pred1, targ2) + normalized_mse(pred2, targ1))
 
-        self.log("train_loss", loss)
+        self.log("train_loss", loss, batch_size=batch_size)
         self.model.update_target()
 
         return loss

torchgeo/trainers/classification.py

@@ -178,11 +178,12 @@ def training_step(
         """
         x = batch["image"]
         y = batch["label"]
+        batch_size = x.shape[0]
         y_hat = self(x)
         loss: Tensor = self.criterion(y_hat, y)
-        self.log("train_loss", loss)
+        self.log("train_loss", loss, batch_size=batch_size)
         self.train_metrics(y_hat, y)
-        self.log_dict(self.train_metrics)
+        self.log_dict(self.train_metrics, batch_size=batch_size)
 
         return loss
 
@@ -198,11 +199,12 @@ def validation_step(
         """
         x = batch["image"]
         y = batch["label"]
+        batch_size = x.shape[0]
         y_hat = self(x)
         loss = self.criterion(y_hat, y)
-        self.log("val_loss", loss)
+        self.log("val_loss", loss, batch_size=batch_size)
         self.val_metrics(y_hat, y)
-        self.log_dict(self.val_metrics)
+        self.log_dict(self.val_metrics, batch_size=batch_size)
 
         if (
             batch_idx < 10
@@ -241,11 +243,12 @@ def test_step(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> None
         """
         x = batch["image"]
         y = batch["label"]
+        batch_size = x.shape[0]
         y_hat = self(x)
         loss = self.criterion(y_hat, y)
-        self.log("test_loss", loss)
+        self.log("test_loss", loss, batch_size=batch_size)
         self.test_metrics(y_hat, y)
-        self.log_dict(self.test_metrics)
+        self.log_dict(self.test_metrics, batch_size=batch_size)
 
     def predict_step(
         self, batch: Any, batch_idx: int, dataloader_idx: int = 0
@@ -317,10 +320,11 @@ def training_step(
         """
         x = batch["image"]
         y = batch["label"]
+        batch_size = x.shape[0]
         y_hat = self(x)
         y_hat_hard = torch.sigmoid(y_hat)
         loss: Tensor = self.criterion(y_hat, y.to(torch.float))
-        self.log("train_loss", loss)
+        self.log("train_loss", loss, batch_size=batch_size)
         self.train_metrics(y_hat_hard, y)
         self.log_dict(self.train_metrics)
 
@@ -338,12 +342,13 @@ def validation_step(
         """
         x = batch["image"]
         y = batch["label"]
+        batch_size = x.shape[0]
         y_hat = self(x)
         y_hat_hard = torch.sigmoid(y_hat)
         loss = self.criterion(y_hat, y.to(torch.float))
-        self.log("val_loss", loss)
+        self.log("val_loss", loss, batch_size=batch_size)
         self.val_metrics(y_hat_hard, y)
-        self.log_dict(self.val_metrics)
+        self.log_dict(self.val_metrics, batch_size=batch_size)
 
         if (
             batch_idx < 10
@@ -381,12 +386,13 @@ def test_step(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> None
         """
         x = batch["image"]
         y = batch["label"]
+        batch_size = x.shape[0]
         y_hat = self(x)
         y_hat_hard = torch.sigmoid(y_hat)
         loss = self.criterion(y_hat, y.to(torch.float))
-        self.log("test_loss", loss)
+        self.log("test_loss", loss, batch_size=batch_size)
         self.test_metrics(y_hat_hard, y)
-        self.log_dict(self.test_metrics)
+        self.log_dict(self.test_metrics, batch_size=batch_size)
 
     def predict_step(
         self, batch: Any, batch_idx: int, dataloader_idx: int = 0

torchgeo/trainers/moco.py

@@ -372,6 +372,7 @@ def training_step(
             The loss tensor.
         """
         x = batch["image"]
+        batch_size = x.shape[0]
 
         in_channels = self.hparams["in_channels"]
         assert x.size(1) == in_channels or x.size(1) == 2 * in_channels
@@ -420,8 +421,8 @@ def training_step(
         output_std = torch.mean(output_std, dim=0)
         self.avg_output_std = 0.9 * self.avg_output_std + (1 - 0.9) * output_std.item()
 
-        self.log("train_ssl_std", self.avg_output_std)
-        self.log("train_loss", loss)
+        self.log("train_ssl_std", self.avg_output_std, batch_size=batch_size)
+        self.log("train_loss", loss, batch_size=batch_size)
 
         return loss
 

torchgeo/trainers/regression.py

@@ -159,15 +159,16 @@ def training_step(
             The loss tensor.
         """
         x = batch["image"]
+        batch_size = x.shape[0]
         # TODO: remove .to(...) once we have a real pixelwise regression dataset
         y = batch[self.target_key].to(torch.float)
         y_hat = self(x)
         if y_hat.ndim != y.ndim:
             y = y.unsqueeze(dim=1)
         loss: Tensor = self.criterion(y_hat, y)
-        self.log("train_loss", loss)
+        self.log("train_loss", loss, batch_size=batch_size)
         self.train_metrics(y_hat, y)
-        self.log_dict(self.train_metrics)
+        self.log_dict(self.train_metrics, batch_size=batch_size)
 
         return loss
 
@@ -182,15 +183,16 @@ def validation_step(
             dataloader_idx: Index of the current dataloader.
         """
         x = batch["image"]
+        batch_size = x.shape[0]
         # TODO: remove .to(...) once we have a real pixelwise regression dataset
         y = batch[self.target_key].to(torch.float)
         y_hat = self(x)
         if y_hat.ndim != y.ndim:
             y = y.unsqueeze(dim=1)
         loss = self.criterion(y_hat, y)
-        self.log("val_loss", loss)
+        self.log("val_loss", loss, batch_size=batch_size)
         self.val_metrics(y_hat, y)
-        self.log_dict(self.val_metrics)
+        self.log_dict(self.val_metrics, batch_size=batch_size)
 
         if (
             batch_idx < 10
@@ -231,15 +233,16 @@ def test_step(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> None
             dataloader_idx: Index of the current dataloader.
         """
         x = batch["image"]
+        batch_size = x.shape[0]
         # TODO: remove .to(...) once we have a real pixelwise regression dataset
         y = batch[self.target_key].to(torch.float)
         y_hat = self(x)
         if y_hat.ndim != y.ndim:
             y = y.unsqueeze(dim=1)
         loss = self.criterion(y_hat, y)
-        self.log("test_loss", loss)
+        self.log("test_loss", loss, batch_size=batch_size)
         self.test_metrics(y_hat, y)
-        self.log_dict(self.test_metrics)
+        self.log_dict(self.test_metrics, batch_size=batch_size)
 
     def predict_step(
         self, batch: Any, batch_idx: int, dataloader_idx: int = 0

torchgeo/trainers/segmentation.py

@@ -227,11 +227,12 @@ def training_step(
         """
         x = batch["image"]
         y = batch["mask"]
+        batch_size = x.shape[0]
         y_hat = self(x)
         loss: Tensor = self.criterion(y_hat, y)
-        self.log("train_loss", loss)
+        self.log("train_loss", loss, batch_size=batch_size)
         self.train_metrics(y_hat, y)
-        self.log_dict(self.train_metrics)
+        self.log_dict(self.train_metrics, batch_size=batch_size)
         return loss
 
     def validation_step(
@@ -246,11 +247,12 @@ def validation_step(
         """
         x = batch["image"]
         y = batch["mask"]
+        batch_size = x.shape[0]
         y_hat = self(x)
         loss = self.criterion(y_hat, y)
-        self.log("val_loss", loss)
+        self.log("val_loss", loss, batch_size=batch_size)
         self.val_metrics(y_hat, y)
-        self.log_dict(self.val_metrics)
+        self.log_dict(self.val_metrics, batch_size=batch_size)
 
         if (
             batch_idx < 10
@@ -289,11 +291,12 @@ def test_step(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> None
         """
         x = batch["image"]
         y = batch["mask"]
+        batch_size = x.shape[0]
         y_hat = self(x)
         loss = self.criterion(y_hat, y)
-        self.log("test_loss", loss)
+        self.log("test_loss", loss, batch_size=batch_size)
         self.test_metrics(y_hat, y)
-        self.log_dict(self.test_metrics)
+        self.log_dict(self.test_metrics, batch_size=batch_size)
 
     def predict_step(
         self, batch: Any, batch_idx: int, dataloader_idx: int = 0

torchgeo/trainers/simclr.py

@@ -222,6 +222,7 @@ def training_step(
             AssertionError: If channel dimensions are incorrect.
         """
         x = batch["image"]
+        batch_size = x.shape[0]
 
         in_channels: int = self.hparams["in_channels"]
         assert x.size(1) == in_channels or x.size(1) == 2 * in_channels
@@ -250,8 +251,8 @@ def training_step(
         output_std = torch.mean(output_std, dim=0)
         self.avg_output_std = 0.9 * self.avg_output_std + (1 - 0.9) * output_std.item()
 
-        self.log("train_ssl_std", self.avg_output_std)
-        self.log("train_loss", loss)
+        self.log("train_ssl_std", self.avg_output_std, batch_size=batch_size)
+        self.log("train_loss", loss, batch_size=batch_size)
 
         return loss