Update PTL examples for v1.6 release (#146)

* update pytorch lightning examples for next release
* remove strategy
* use torchmetrics
* improvements
* use estimated stepping batches
* use auto devices
* fix examples
* Apply suggestions from code review
* limit 1 GPU

Co-authored-by: Jirka Borovec <Borda@users.noreply.github.com>
Co-authored-by: Jirka <jirka.borovec@seznam.cz>

Author: rohitgr7

lightning_examples/augmentation_kornia/augmentation.py

@@ -11,14 +11,13 @@
 from kornia import image_to_tensor, tensor_to_image
 from kornia.augmentation import ColorJitter, RandomChannelShuffle, RandomHorizontalFlip, RandomThinPlateSpline
 from pytorch_lightning import LightningModule, Trainer
+from pytorch_lightning.callbacks.progress import TQDMProgressBar
 from pytorch_lightning.loggers import CSVLogger
 from torch import Tensor
 from torch.nn import functional as F
 from torch.utils.data import DataLoader
 from torchvision.datasets import CIFAR10
 
-AVAIL_GPUS = min(1, torch.cuda.device_count())
-
 # %% [markdown]
 # ## Define Data Augmentations module
 #
@@ -106,10 +105,11 @@ def __init__(self):
 
         self.transform = DataAugmentation()  # per batch augmentation_kornia
 
-        self.accuracy = torchmetrics.Accuracy()
+        self.train_accuracy = torchmetrics.Accuracy()
+        self.val_accuracy = torchmetrics.Accuracy()
 
     def forward(self, x):
-        return F.softmax(self.model(x))
+        return self.model(x)
 
     def compute_loss(self, y_hat, y):
         return F.cross_entropy(y_hat, y)
@@ -127,21 +127,28 @@ def _to_vis(data):
         plt.figure(figsize=win_size)
         plt.imshow(_to_vis(imgs_aug))
 
+    def on_after_batch_transfer(self, batch, dataloader_idx):
+        x, y = batch
+        if self.trainer.training:
+            x = self.transform(x)  # => we perform GPU/Batched data augmentation
+        return x, y
+
     def training_step(self, batch, batch_idx):
         x, y = batch
-        x_aug = self.transform(x)  # => we perform GPU/Batched data augmentation
-        y_hat = self(x_aug)
+        y_hat = self(x)
         loss = self.compute_loss(y_hat, y)
+        self.train_accuracy.update(y_hat, y)
         self.log("train_loss", loss, prog_bar=False)
-        self.log("train_acc", self.accuracy(y_hat, y), prog_bar=False)
+        self.log("train_acc", self.train_accuracy, prog_bar=False)
         return loss
 
     def validation_step(self, batch, batch_idx):
         x, y = batch
         y_hat = self(x)
         loss = self.compute_loss(y_hat, y)
+        self.val_accuracy.update(y_hat, y)
         self.log("valid_loss", loss, prog_bar=False)
-        self.log("valid_acc", self.accuracy(y_hat, y), prog_bar=True)
+        self.log("valid_acc", self.val_accuracy, prog_bar=True)
 
     def configure_optimizers(self):
         optimizer = torch.optim.AdamW(self.model.parameters(), lr=1e-4)
@@ -158,7 +165,7 @@ def train_dataloader(self):
         return loader
 
     def val_dataloader(self):
-        dataset = CIFAR10(os.getcwd(), train=True, download=True, transform=self.preprocess)
+        dataset = CIFAR10(os.getcwd(), train=False, download=True, transform=self.preprocess)
         loader = DataLoader(dataset, batch_size=32)
         return loader
 
@@ -179,8 +186,9 @@ def val_dataloader(self):
 # %%
 # Initialize a trainer
 trainer = Trainer(
-    progress_bar_refresh_rate=20,
-    gpus=AVAIL_GPUS,
+    callbacks=[TQDMProgressBar(refresh_rate=20)],
+    accelerator="auto",
+    devices=1 if torch.cuda.is_available() else None,  # limiting got iPython runs
     max_epochs=10,
     logger=CSVLogger(save_dir="logs/", name="cifar10-resnet18"),
 )

lightning_examples/barlow-twins/barlow_twins.py

@@ -283,15 +283,12 @@ def shared_step(self, batch):
 
     def training_step(self, batch, batch_idx):
         loss = self.shared_step(batch)
-
-        self.log("train_loss", loss.item(), on_step=True, on_epoch=False)
+        self.log("train_loss", loss, on_step=True, on_epoch=False)
         return loss
 
     def validation_step(self, batch, batch_idx):
         loss = self.shared_step(batch)
-
         self.log("val_loss", loss, on_step=False, on_epoch=True)
-        return loss
 
     def configure_optimizers(self):
         optimizer = torch.optim.Adam(self.parameters(), lr=self.learning_rate)
@@ -329,7 +326,7 @@ def __init__(
         self.encoder_output_dim = encoder_output_dim
         self.num_classes = num_classes
 
-    def on_pretrain_routine_start(self, trainer: pl.Trainer, pl_module: pl.LightningModule) -> None:
+    def on_fit_start(self, trainer: pl.Trainer, pl_module: pl.LightningModule) -> None:
 
         # add linear_eval layer and optimizer
         pl_module.online_finetuner = nn.Linear(self.encoder_output_dim, self.num_classes).to(pl_module.device)
@@ -408,12 +405,12 @@ def on_validation_batch_end(
 )
 
 online_finetuner = OnlineFineTuner(encoder_output_dim=encoder_out_dim, num_classes=10)
-checkpoint_callback = ModelCheckpoint(every_n_val_epochs=100, save_top_k=-1, save_last=True)
+checkpoint_callback = ModelCheckpoint(every_n_epochs=100, save_top_k=-1, save_last=True)
 
 trainer = Trainer(
     max_epochs=max_epochs,
-    gpus=torch.cuda.device_count(),
-    precision=16 if torch.cuda.device_count() > 0 else 32,
+    accelerator="auto",
+    devices=1 if torch.cuda.is_available() else None,  # limiting got iPython runs
     callbacks=[online_finetuner, checkpoint_callback],
 )
 

lightning_examples/basic-gan/gan.py

@@ -1,6 +1,5 @@
 # %%
 import os
-from collections import OrderedDict
 
 import numpy as np
 import torch
@@ -9,12 +8,12 @@
 import torchvision
 import torchvision.transforms as transforms
 from pytorch_lightning import LightningDataModule, LightningModule, Trainer
+from pytorch_lightning.callbacks.progress import TQDMProgressBar
 from torch.utils.data import DataLoader, random_split
 from torchvision.datasets import MNIST
 
 PATH_DATASETS = os.environ.get("PATH_DATASETS", ".")
-AVAIL_GPUS = min(1, torch.cuda.device_count())
-BATCH_SIZE = 256 if AVAIL_GPUS else 64
+BATCH_SIZE = 256 if torch.cuda.is_available() else 64
 NUM_WORKERS = int(os.cpu_count() / 2)
 
 # %% [markdown]
@@ -205,9 +204,8 @@ def training_step(self, batch, batch_idx, optimizer_idx):
 
             # adversarial loss is binary cross-entropy
             g_loss = self.adversarial_loss(self.discriminator(self(z)), valid)
-            tqdm_dict = {"g_loss": g_loss}
-            output = OrderedDict({"loss": g_loss, "progress_bar": tqdm_dict, "log": tqdm_dict})
-            return output
+            self.log("g_loss", g_loss, prog_bar=True)
+            return g_loss
 
         # train discriminator
         if optimizer_idx == 1:
@@ -227,9 +225,8 @@ def training_step(self, batch, batch_idx, optimizer_idx):
 
             # discriminator loss is the average of these
             d_loss = (real_loss + fake_loss) / 2
-            tqdm_dict = {"d_loss": d_loss}
-            output = OrderedDict({"loss": d_loss, "progress_bar": tqdm_dict, "log": tqdm_dict})
-            return output
+            self.log("d_loss", d_loss, prog_bar=True)
+            return d_loss
 
     def configure_optimizers(self):
         lr = self.hparams.lr
@@ -240,7 +237,7 @@ def configure_optimizers(self):
         opt_d = torch.optim.Adam(self.discriminator.parameters(), lr=lr, betas=(b1, b2))
         return [opt_g, opt_d], []
 
-    def on_epoch_end(self):
+    def on_validation_epoch_end(self):
         z = self.validation_z.type_as(self.generator.model[0].weight)
 
         # log sampled images
@@ -252,7 +249,12 @@ def on_epoch_end(self):
 # %%
 dm = MNISTDataModule()
 model = GAN(*dm.size())
-trainer = Trainer(gpus=AVAIL_GPUS, max_epochs=5, progress_bar_refresh_rate=20)
+trainer = Trainer(
+    accelerator="auto",
+    devices=1 if torch.cuda.is_available() else None,  # limiting got iPython runs
+    max_epochs=5,
+    callbacks=[TQDMProgressBar(refresh_rate=20)],
+)
 trainer.fit(model, dm)
 
 # %%

lightning_examples/cifar10-baseline/baseline.py

@@ -13,6 +13,7 @@
 from pl_bolts.transforms.dataset_normalizations import cifar10_normalization
 from pytorch_lightning import LightningModule, Trainer, seed_everything
 from pytorch_lightning.callbacks import LearningRateMonitor
+from pytorch_lightning.callbacks.progress import TQDMProgressBar
 from pytorch_lightning.loggers import TensorBoardLogger
 from torch.optim.lr_scheduler import OneCycleLR
 from torch.optim.swa_utils import AveragedModel, update_bn
@@ -21,8 +22,7 @@
 seed_everything(7)
 
 PATH_DATASETS = os.environ.get("PATH_DATASETS", ".")
-AVAIL_GPUS = min(1, torch.cuda.device_count())
-BATCH_SIZE = 256 if AVAIL_GPUS else 64
+BATCH_SIZE = 256 if torch.cuda.is_available() else 64
 NUM_WORKERS = int(os.cpu_count() / 2)
 
 # %% [markdown]
@@ -137,14 +137,13 @@ def configure_optimizers(self):
 
 # %%
 model = LitResnet(lr=0.05)
-model.datamodule = cifar10_dm
 
 trainer = Trainer(
-    progress_bar_refresh_rate=10,
     max_epochs=30,
-    gpus=AVAIL_GPUS,
+    accelerator="auto",
+    devices=1 if torch.cuda.is_available() else None,  # limiting got iPython runs
     logger=TensorBoardLogger("lightning_logs/", name="resnet"),
-    callbacks=[LearningRateMonitor(logging_interval="step")],
+    callbacks=[LearningRateMonitor(logging_interval="step"), TQDMProgressBar(refresh_rate=10)],
 )
 
 trainer.fit(model, cifar10_dm)
@@ -189,18 +188,19 @@ def configure_optimizers(self):
         return optimizer
 
     def on_train_end(self):
-        update_bn(self.datamodule.train_dataloader(), self.swa_model, device=self.device)
+        update_bn(self.trainer.datamodule.train_dataloader(), self.swa_model, device=self.device)
 
 
 # %%
 swa_model = SWAResnet(model.model, lr=0.01)
 swa_model.datamodule = cifar10_dm
 
 swa_trainer = Trainer(
-    progress_bar_refresh_rate=20,
     max_epochs=20,
-    gpus=AVAIL_GPUS,
+    accelerator="auto",
+    devices=1 if torch.cuda.is_available() else None,  # limiting got iPython runs
     logger=TensorBoardLogger("lightning_logs/", name="swa_resnet"),
+    callbacks=[TQDMProgressBar(refresh_rate=20)],
 )
 
 swa_trainer.fit(swa_model, cifar10_dm)

lightning_examples/datamodules/datamodules.py

@@ -9,6 +9,7 @@
 import torch
 import torch.nn.functional as F
 from pytorch_lightning import LightningDataModule, LightningModule, Trainer
+from pytorch_lightning.callbacks.progress import TQDMProgressBar
 from torch import nn
 from torch.utils.data import DataLoader, random_split
 from torchmetrics.functional import accuracy
@@ -18,8 +19,7 @@
 from torchvision.datasets import CIFAR10, MNIST
 
 PATH_DATASETS = os.environ.get("PATH_DATASETS", ".")
-AVAIL_GPUS = min(1, torch.cuda.device_count())
-BATCH_SIZE = 256 if AVAIL_GPUS else 64
+BATCH_SIZE = 256 if torch.cuda.is_available() else 64
 
 # %% [markdown]
 # ### Defining the LitMNISTModel
@@ -84,7 +84,6 @@ def validation_step(self, batch, batch_idx):
         acc = accuracy(preds, y)
         self.log("val_loss", loss, prog_bar=True)
         self.log("val_acc", acc, prog_bar=True)
-        return loss
 
     def configure_optimizers(self):
         optimizer = torch.optim.Adam(self.parameters(), lr=self.learning_rate)
@@ -127,8 +126,9 @@ def test_dataloader(self):
 model = LitMNIST()
 trainer = Trainer(
     max_epochs=2,
-    gpus=AVAIL_GPUS,
-    progress_bar_refresh_rate=20,
+    accelerator="auto",
+    devices=1 if torch.cuda.is_available() else None,  # limiting got iPython runs
+    callbacks=[TQDMProgressBar(refresh_rate=20)],
 )
 trainer.fit(model)
 
@@ -252,15 +252,13 @@ def training_step(self, batch, batch_idx):
         return loss
 
     def validation_step(self, batch, batch_idx):
-
         x, y = batch
         logits = self(x)
         loss = F.nll_loss(logits, y)
         preds = torch.argmax(logits, dim=1)
         acc = accuracy(preds, y)
         self.log("val_loss", loss, prog_bar=True)
         self.log("val_acc", acc, prog_bar=True)
-        return loss
 
     def configure_optimizers(self):
         optimizer = torch.optim.Adam(self.parameters(), lr=self.learning_rate)
@@ -280,8 +278,9 @@ def configure_optimizers(self):
 # Init trainer
 trainer = Trainer(
     max_epochs=3,
-    progress_bar_refresh_rate=20,
-    gpus=AVAIL_GPUS,
+    callbacks=[TQDMProgressBar(refresh_rate=20)],
+    accelerator="auto",
+    devices=1 if torch.cuda.is_available() else None,  # limiting got iPython runs
 )
 # Pass the datamodule as arg to trainer.fit to override model hooks :)
 trainer.fit(model, dm)
@@ -343,9 +342,11 @@ def test_dataloader(self):
 # %%
 dm = CIFAR10DataModule()
 model = LitModel(*dm.size(), dm.num_classes, hidden_size=256)
+tqdm_progress_bar = TQDMProgressBar(refresh_rate=20)
 trainer = Trainer(
     max_epochs=5,
-    progress_bar_refresh_rate=20,
-    gpus=AVAIL_GPUS,
+    accelerator="auto",
+    devices=1 if torch.cuda.is_available() else None,  # limiting got iPython runs
+    callbacks=[tqdm_progress_bar],
 )
 trainer.fit(model, dm)