update examoles with BoringModel

Author: Borda

examples/resume-lightning-training.py

@@ -2,31 +2,22 @@
 This example demonstrates how to resume training of a model using the `download_model` function.
 """
 
-import torch.utils.data as data
-import torchvision as tv
 from lightning import Trainer
+from lightning.pytorch.demos.boring_classes import BoringModel
 from litmodels import download_model
-from sample_model import LitAutoEncoder
 
 # Define the model name - this should be unique to your model
 # The format is <organization>/<teamspace>/<model-name>:<model-version>
 MY_MODEL_NAME = "jirka/kaggle/lit-auto-encoder-callback:latest"
 
 
 if __name__ == "__main__":
-    dataset = tv.datasets.MNIST(".", download=True, transform=tv.transforms.ToTensor())
-    train, val = data.random_split(dataset, [55000, 5000])
-
     model_path = download_model(name=MY_MODEL_NAME, download_dir="my_models")
     print(f"model: {model_path}")
     # autoencoder = LitAutoEncoder.load_from_checkpoint(checkpoint_path=model_path)
 
-    trainer = Trainer(
-        max_epochs=4,
-    )
+    trainer = Trainer(max_epochs=4)
     trainer.fit(
-        LitAutoEncoder(),
-        data.DataLoader(train, batch_size=256),
-        data.DataLoader(val, batch_size=256),
+        BoringModel(),
         ckpt_path=model_path,
     )

examples/sample_model.py

@@ -2,29 +2,18 @@
 This example demonstrates how to train a model and upload it to the cloud using the `upload_model` function.
 """
 
-import torch.utils.data as data
-import torchvision as tv
 from lightning import Trainer
+from lightning.pytorch.demos.boring_classes import BoringModel
 from litmodels import upload_model
-from sample_model import LitAutoEncoder
 
 # Define the model name - this should be unique to your model
 # The format is <organization>/<teamspace>/<model-name>
 MY_MODEL_NAME = "jirka/kaggle/lit-auto-encoder-simple"
 
 
 if __name__ == "__main__":
-    dataset = tv.datasets.MNIST(".", download=True, transform=tv.transforms.ToTensor())
-    train, val = data.random_split(dataset, [55000, 5000])
-
-    autoencoder = LitAutoEncoder()
-
     trainer = Trainer(max_epochs=2)
-    trainer.fit(
-        autoencoder,
-        data.DataLoader(train, batch_size=256),
-        data.DataLoader(val, batch_size=256),
-    )
+    trainer.fit(BoringModel())
     checkpoint_path = getattr(trainer.checkpoint_callback, "best_model_path")
     print(f"best: {checkpoint_path}")
     upload_model(model=checkpoint_path, name=MY_MODEL_NAME)

examples/train-model-and-simple-save.py

@@ -2,29 +2,18 @@
 Train a model with a Lightning callback that uploads the best model to the cloud after each epoch.
 """
 
-import torch.utils.data as data
-import torchvision as tv
 from lightning import Trainer
-from litmodels.integrations import LitModelCheckpoint
-from sample_model import LitAutoEncoder
+from lightning.pytorch.demos.boring_classes import BoringModel
+from litmodels.integrations import LightningModelCheckpoint
 
 # Define the model name - this should be unique to your model
 # The format is <organization>/<teamspace>/<model-name>
 MY_MODEL_NAME = "lightning-ai/jirka/lit-auto-encoder-callback"
 
 
 if __name__ == "__main__":
-    dataset = tv.datasets.MNIST(".", download=True, transform=tv.transforms.ToTensor())
-    train, val = data.random_split(dataset, [55000, 5000])
-
-    autoencoder = LitAutoEncoder()
-
     trainer = Trainer(
         max_epochs=2,
-        callbacks=LitModelCheckpoint(model_name=MY_MODEL_NAME),
-    )
-    trainer.fit(
-        autoencoder,
-        data.DataLoader(train, batch_size=256),
-        data.DataLoader(val, batch_size=256),
+        callbacks=LightningModelCheckpoint(model_name=MY_MODEL_NAME),
     )
+    trainer.fit(BoringModel())

examples/train-model-with-lightning-callback.py

@@ -7,58 +7,24 @@
 
 """
 
-import os
-
-from lightning import LightningModule, Trainer
+from lightning import Trainer
+from lightning.pytorch.demos.boring_classes import BoringModel
 from litlogger import LightningLogger
-from psutil import cpu_count
-from torch import nn, optim
-from torch.utils.data import DataLoader
-from torchvision.datasets import MNIST
-from torchvision.transforms import ToTensor
-
-
-class LitAutoEncoder(LightningModule):
-    def __init__(self, lr=1e-3, inp_size=28):
-        super().__init__()
 
-        self.encoder = nn.Sequential(nn.Linear(inp_size * inp_size, 64), nn.ReLU(), nn.Linear(64, 3))
-        self.decoder = nn.Sequential(nn.Linear(3, 64), nn.ReLU(), nn.Linear(64, inp_size * inp_size))
-        self.lr = lr
-        self.save_hyperparameters()
 
+class DemoModel(BoringModel):
     def training_step(self, batch, batch_idx):
-        x, y = batch
-        x = x.view(x.size(0), -1)
-        z = self.encoder(x)
-        x_hat = self.decoder(z)
-        loss = nn.functional.mse_loss(x_hat, x)
-        # log metrics
-        self.log("train_loss", loss)
-        return loss
-
-    def configure_optimizers(self):
-        return optim.Adam(self.parameters(), lr=self.lr)
+        output = super().training_step(batch, batch_idx)
+        self.log("train_loss", output["loss"])
+        return output
 
 
 if __name__ == "__main__":
-    # init the autoencoder
-    autoencoder = LitAutoEncoder(lr=1e-3, inp_size=28)
-
-    # setup data
-    train_loader = DataLoader(
-        dataset=MNIST(os.getcwd(), download=True, transform=ToTensor()),
-        batch_size=32,
-        shuffle=True,
-        num_workers=cpu_count(),
-        persistent_workers=True,
-    )
-
     # configure the logger
     lit_logger = LightningLogger(log_model=True)
 
     # pass logger to the Trainer
     trainer = Trainer(max_epochs=5, logger=lit_logger)
 
     # train the model
-    trainer.fit(model=autoencoder, train_dataloaders=train_loader)
+    trainer.fit(model=DemoModel())