Reuse code in demos.BoringModel (#16242)

Author: carmocca

examples/app_components/python/pl_script.py

@@ -1,60 +1,5 @@
-import torch
-from torch.utils.data import DataLoader, Dataset
-
-from pytorch_lightning import LightningModule, Trainer
-
-
-class RandomDataset(Dataset):
-    def __init__(self, size: int, length: int):
-        self.len = length
-        self.data = torch.randn(length, size)
-
-    def __getitem__(self, index):
-        return self.data[index]
-
-    def __len__(self):
-        return self.len
-
-
-class BoringModel(LightningModule):
-    def __init__(self):
-        super().__init__()
-        self.layer = torch.nn.Linear(32, 2)
-
-    def forward(self, x):
-        return self.layer(x)
-
-    def loss(self, batch, prediction):
-        # An arbitrary loss to have a loss that updates the model weights during `Trainer.fit` calls
-        return torch.nn.functional.mse_loss(prediction, torch.ones_like(prediction))
-
-    def training_step(self, batch, batch_idx):
-        output = self(batch)
-        loss = self.loss(batch, output)
-        return {"loss": loss}
-
-    def validation_step(self, batch, batch_idx):
-        output = self(batch)
-        loss = self.loss(batch, output)
-        return {"x": loss}
-
-    def test_step(self, batch, batch_idx):
-        output = self(batch)
-        loss = self.loss(batch, output)
-        return {"y": loss}
-
-    def configure_optimizers(self):
-        optimizer = torch.optim.SGD(self.layer.parameters(), lr=0.1)
-        lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1)
-        return [optimizer], [lr_scheduler]
-
-    def train_dataloader(self):
-        return DataLoader(RandomDataset(32, 64))
-
-    val_dataloader = train_dataloader
-    test_dataloader = train_dataloader
-    predict_dataloader = train_dataloader
-
+from pytorch_lightning import Trainer
+from pytorch_lightning.demos.boring_classes import BoringModel
 
 if __name__ == "__main__":
     model = BoringModel()

examples/app_components/python/pytorch_lightning_script.py

@@ -83,34 +83,39 @@ def __init__(self) -> None:
         - subclass
         - modify the behavior for what you want
 
-        class TestModel(BaseTestModel):
-            def training_step(...):
-                # do your own thing
+        Example::
 
-        or:
+            class TestModel(BoringModel):
+                def training_step(self, ...):
+                    ...  # do your own thing
 
-        model = BaseTestModel()
-        model.training_epoch_end = None
+                training_epoch_end = None  # disable hook
+
+        or
+
+        Example::
+
+            model = BoringModel()
+            model.training_epoch_end = None  # disable hook
         """
         super().__init__()
         self.layer = torch.nn.Linear(32, 2)
 
     def forward(self, x: Tensor) -> Tensor:
         return self.layer(x)
 
-    def loss(self, batch: Tensor, preds: Tensor) -> Tensor:
+    def loss(self, preds: Tensor, labels: Optional[Tensor] = None) -> Tensor:
+        if labels is None:
+            labels = torch.ones_like(preds)
         # An arbitrary loss to have a loss that updates the model weights during `Trainer.fit` calls
-        return torch.nn.functional.mse_loss(preds, torch.ones_like(preds))
+        return torch.nn.functional.mse_loss(preds, labels)
 
-    def step(self, x: Tensor) -> Tensor:
-        x = self(x)
-        out = torch.nn.functional.mse_loss(x, torch.ones_like(x))
-        return out
+    def step(self, batch: Tensor) -> Tensor:
+        output = self(batch)
+        return self.loss(output)
 
     def training_step(self, batch: Tensor, batch_idx: int) -> STEP_OUTPUT:
-        output = self(batch)
-        loss = self.loss(batch, output)
-        return {"loss": loss}
+        return {"loss": self.step(batch)}
 
     def training_step_end(self, training_step_outputs: STEP_OUTPUT) -> STEP_OUTPUT:
         return training_step_outputs
@@ -120,18 +125,14 @@ def training_epoch_end(self, outputs: EPOCH_OUTPUT) -> None:
         torch.stack([x["loss"] for x in outputs]).mean()
 
     def validation_step(self, batch: Tensor, batch_idx: int) -> Optional[STEP_OUTPUT]:
-        output = self(batch)
-        loss = self.loss(batch, output)
-        return {"x": loss}
+        return {"x": self.step(batch)}
 
     def validation_epoch_end(self, outputs: Union[EPOCH_OUTPUT, List[EPOCH_OUTPUT]]) -> None:
         outputs = cast(List[Dict[str, Tensor]], outputs)
         torch.stack([x["x"] for x in outputs]).mean()
 
     def test_step(self, batch: Tensor, batch_idx: int) -> Optional[STEP_OUTPUT]:
-        output = self(batch)
-        loss = self.loss(batch, output)
-        return {"y": loss}
+        return {"y": self.step(batch)}
 
     def test_epoch_end(self, outputs: Union[EPOCH_OUTPUT, List[EPOCH_OUTPUT]]) -> None:
         outputs = cast(List[Dict[str, Tensor]], outputs)
@@ -194,8 +195,7 @@ def __init__(self) -> None:
     def training_step(self, batch: Tensor, batch_idx: int) -> STEP_OUTPUT:
         opt = self.optimizers()
         assert isinstance(opt, (Optimizer, LightningOptimizer))
-        output = self(batch)
-        loss = self.loss(batch, output)
+        loss = self.step(batch)
         opt.zero_grad()
         self.manual_backward(loss)
         opt.step()

src/pytorch_lightning/demos/boring_classes.py

@@ -39,19 +39,13 @@
 
 class IPUModel(BoringModel):
     def training_step(self, batch, batch_idx):
-        output = self(batch)
-        loss = self.loss(batch, output)
-        return loss
+        return self.step(batch)
 
     def validation_step(self, batch, batch_idx):
-        output = self(batch)
-        loss = self.loss(batch, output)
-        return loss
+        return self.step(batch)
 
     def test_step(self, batch, batch_idx):
-        output = self(batch)
-        loss = self.loss(batch, output)
-        return loss
+        return self.step(batch)
 
     def training_epoch_end(self, outputs) -> None:
         pass

tests/tests_pytorch/accelerators/test_ipu.py

@@ -144,8 +144,7 @@ def on_train_batch_start(self, batch, batch_idx):
         def training_step(self, batch, batch_idx):
             self.called["training_step"] += 1
             opt = self.optimizers()
-            output = self.layer(batch)
-            loss = self.loss(batch, output)
+            loss = self.step(batch)
 
             if self.should_update:
                 self.manual_backward(loss)

tests/tests_pytorch/accelerators/test_tpu.py

@@ -52,11 +52,6 @@ def __init__(self):
             self.layer = torch.nn.Linear(32, 2)
             self.backbone.has_been_used = False
 
-        def training_step(self, batch, batch_idx):
-            output = self(batch)
-            loss = self.loss(batch, output)
-            return {"loss": loss}
-
         def forward(self, x):
             self.backbone.has_been_used = True
             x = self.backbone(x)
@@ -101,11 +96,6 @@ def __init__(self):
             self.layer = None
             self.backbone.has_been_used = False
 
-        def training_step(self, batch, batch_idx):
-            output = self(batch)
-            loss = self.loss(batch, output)
-            return {"loss": loss}
-
         def forward(self, x):
             self.backbone.has_been_used = True
             x = self.backbone(x)

tests/tests_pytorch/callbacks/test_finetuning_callback.py

@@ -65,9 +65,7 @@ def __init__(
     def training_step(self, batch, batch_idx):
         if self.crash_on_epoch and self.trainer.current_epoch >= self.crash_on_epoch:
             raise Exception("SWA crash test")
-        output = self.forward(batch)
-        loss = self.loss(batch, output)
-        return {"loss": loss}
+        return super().training_step(batch, batch_idx)
 
     def train_dataloader(self):
         dset_cls = RandomIterableDataset if self.iterable_dataset else RandomDataset

tests/tests_pytorch/callbacks/test_stochastic_weight_avg.py

@@ -859,8 +859,7 @@ def test_checkpoint_repeated_strategy(tmpdir):
 
     class ExtendedBoringModel(BoringModel):
         def validation_step(self, batch, batch_idx):
-            output = self.layer(batch)
-            loss = self.loss(batch, output)
+            loss = self.step(batch)
             self.log("val_loss", loss)
 
     model = ExtendedBoringModel()
@@ -898,8 +897,7 @@ def test_checkpoint_repeated_strategy_extended(tmpdir):
 
     class ExtendedBoringModel(BoringModel):
         def validation_step(self, batch, batch_idx):
-            output = self.layer(batch)
-            loss = self.loss(batch, output)
+            loss = self.step(batch)
             self.log("val_loss", loss)
             return {"val_loss": loss}
 

tests/tests_pytorch/checkpointing/test_model_checkpoint.py

@@ -37,8 +37,7 @@ def configure_optimizers(self):
             return [optimizer], [lr_scheduler]
 
         def validation_step(self, batch, batch_idx):
-            output = self.layer(batch)
-            loss = self.loss(batch, output)
+            loss = self.step(batch)
             self.log("val_loss", loss, on_epoch=True, prog_bar=True)
 
     model = ExtendedBoringModel()

tests/tests_pytorch/checkpointing/test_trainer_checkpoint.py

@@ -85,8 +85,7 @@ def training_step(self, batch, batch_idx):
             assert isinstance(opt_2, LightningOptimizer)
 
             def closure(opt):
-                output = self.layer(batch)
-                loss = self.loss(batch, output)
+                loss = self.step(batch)
                 opt.zero_grad()
                 self.manual_backward(loss)
 
@@ -323,8 +322,7 @@ def test_lightning_optimizer_keeps_hooks(tmpdir):
 def test_params_groups_and_state_are_accessible(tmpdir):
     class TestModel(BoringModel):
         def training_step(self, batch, batch_idx, optimizer_idx):
-            output = self.layer(batch)
-            loss = self.loss(batch, output)
+            loss = self.step(batch)
             self.__loss = loss
             return loss