PyTorch Lightning: Deep Learning API for PyTorch

<div class="talk-meta">
<div>
<p class="talk-title">PyTorch Lightning: Deep Learning API for PyTorch</p>
<p class="talk-subtitle">Campbell Lab Meeting: Tech Talk</p>
</div>
<div class="talk-author">
<span><svg aria-hidden="true" role="img" viewBox="0 0 576 512" style="height:1em;width:1.12em;vertical-align:-0.125em;margin-left:auto;margin-right:auto;font-size:inherit;fill:#FFFFFF;overflow:visible;position:relative;"><path d="M336.2 64H47.8C21.4 64 0 85.4 0 111.8v288.4C0 426.6 21.4 448 47.8 448h288.4c26.4 0 47.8-21.4 47.8-47.8V111.8c0-26.4-21.4-47.8-47.8-47.8zm189.4 37.7L416 177.3v157.4l109.6 75.5c21.2 14.6 50.4-.3 50.4-25.8V127.5c0-25.4-29.1-40.4-50.4-25.8z"/></svg> Matt Warkentin</span>
<br/>
<br/>
<div>
<svg aria-hidden="true" role="img" viewBox="0 0 576 512" style="height:1em;width:1.12em;vertical-align:-0.125em;margin-left:auto;margin-right:auto;font-size:inherit;fill:#FFFFFF;overflow:visible;position:relative;"><path d="M280.37 148.26L96 300.11V464a16 16 0 0 0 16 16l112.06-.29a16 16 0 0 0 15.92-16V368a16 16 0 0 1 16-16h64a16 16 0 0 1 16 16v95.64a16 16 0 0 0 16 16.05L464 480a16 16 0 0 0 16-16V300L295.67 148.26a12.19 12.19 0 0 0-15.3 0zM571.6 251.47L488 182.56V44.05a12 12 0 0 0-12-12h-56a12 12 0 0 0-12 12v72.61L318.47 43a48 48 0 0 0-61 0L4.34 251.47a12 12 0 0 0-1.6 16.9l25.5 31A12 12 0 0 0 45.15 301l235.22-193.74a12.19 12.19 0 0 1 15.3 0L530.9 301a12 12 0 0 0 16.9-1.6l25.5-31a12 12 0 0 0-1.7-16.93z"/></svg>
Lunenfeld-Tanenbaum Research Institute, .b[.sinai-blue[Sinai] .sinai-orange[Health] .sinai-red[System]]
</div>
</div>
<div class="talk-date">
<span>
<svg aria-hidden="true" role="img" viewBox="0 0 448 512" style="height:1em;width:0.88em;vertical-align:-0.125em;margin-left:auto;margin-right:auto;font-size:inherit;fill:#FFFFFF;overflow:visible;position:relative;"><path d="M0 464c0 26.5 21.5 48 48 48h352c26.5 0 48-21.5 48-48V192H0v272zm320-196c0-6.6 5.4-12 12-12h40c6.6 0 12 5.4 12 12v40c0 6.6-5.4 12-12 12h-40c-6.6 0-12-5.4-12-12v-40zm0 128c0-6.6 5.4-12 12-12h40c6.6 0 12 5.4 12 12v40c0 6.6-5.4 12-12 12h-40c-6.6 0-12-5.4-12-12v-40zM192 268c0-6.6 5.4-12 12-12h40c6.6 0 12 5.4 12 12v40c0 6.6-5.4 12-12 12h-40c-6.6 0-12-5.4-12-12v-40zm0 128c0-6.6 5.4-12 12-12h40c6.6 0 12 5.4 12 12v40c0 6.6-5.4 12-12 12h-40c-6.6 0-12-5.4-12-12v-40zM64 268c0-6.6 5.4-12 12-12h40c6.6 0 12 5.4 12 12v40c0 6.6-5.4 12-12 12H76c-6.6 0-12-5.4-12-12v-40zm0 128c0-6.6 5.4-12 12-12h40c6.6 0 12 5.4 12 12v40c0 6.6-5.4 12-12 12H76c-6.6 0-12-5.4-12-12v-40zM400 64h-48V16c0-8.8-7.2-16-16-16h-32c-8.8 0-16 7.2-16 16v48H160V16c0-8.8-7.2-16-16-16h-32c-8.8 0-16 7.2-16 16v48H48C21.5 64 0 85.5 0 112v48h448v-48c0-26.5-21.5-48-48-48z"/></svg>
June 10th, 2022
</span>
</div>
</div>

---

# Background

- This presentation assumes some moderate to advanced familiarity with deep learning workflows and `PyTorch`

- Briefly: `PyTorch` is an open source machine learning framework that accelerates the path from research prototyping to production deployment.
    + https://pytorch.org/

- PyTorch, like Tensorflow, is a low-level computational library

- Often, as machine learning practitioners, we may wish to use higher-level libraries that simplify or abstract away technical details that slow down getting up and running
    + Quick victories and iteration are key to successful AI projects

---

# Problems with high-level libraries

- Debatably, Keras is to Tensorflow what PyTorch Lightning/Ignite/FastAI/Other are to PyTorch

- My primary issue with Keras was that it was too *hand-holdy*

- It works really well for "traditional" deep learning projects, but as soon as you needed to deviate from the anticipated workflow, there was a lot of friction

- In other words, Keras abstracted away technical debt at the cost of losing flexibility

- Why is **PyTorch Lightning** different?
    + PyTorch Lightning isn't an abstraction, it is PyTorch
    + Anything that can be achieved with vanilla PyTorch can be done with PTL
    + Your project can be as simple or complex as you desire; decide how much PTL code you want to overwrite...

---

# Every deep learning project...

```py
for i in range(epochs):
  for j in range(train_batches):
    # put model in train mode
    # assemble a batch
    # forward pass
    # compute loss
    # log metrics
    # zero gradients
    # backward step
    # optimizer step
  
  for j in range(valid_batches):
    # put model in eval mode
    # assemble a batch
    # forward pass
    # compute loss
    # log metrics
```

---

# The unique part of your project...

```py
for i in range(epochs):
  for j in range(train_batches):
    # put model in train mode
    # assemble a batch
*   # forward pass
    # compute loss
    # log metrics
    # zero gradients
    # backward step
    # optimizer step
  
  for j in range(valid_batches):
    # put model in eval mode
    # assemble a batch
*   # forward pass
    # compute loss
    # log metrics
```

---

# PyTorch Lightning API

- PTL provides a class called a `LightningModule` which wraps up the logic for training a deep learning model
    + Note: `LightningModule` is a `nn.Module` with extra features

- As a user, your model needs to subclass the `LightningModule` and define a small number of methods, while letting PTL handle the rest

- Requisite methods:
    + `init`
    + `forward`
    + `configure_optimizers`
    + `training_step`

- Optional, but common methods:
    + `validation_step`
    + `test_step`

---

# PyTorch Lightning API

class MySweetModel(pl.LightningModule):
    def __init__(self, ...):
        pass
    
    def forward(self, x):
        return y_hat
    
    def configure_optimizers(self):
        return optimizer
    
    def training_step(self, batch, batch_id):
        return train_loss
```
]

---

# PyTorch Lightning API

- While a `LightningModule` defines the logic for your deep learning model, you will also want to define a class that encapsulates the logic of your data

- Taken together, the data logic and model logic should be transportable and work in (almost) any environment

> A datamodule is a shareable, reusable class that encapsulates all the steps needed to process data

- Similar to `LightningModule`, your data module should subclass `LightningDataModule` and define some important methods

- Requisite methods:
    + `init`
    + `setup`
    + `train_dataloader`
    
- Optional: `prepare_data`, `val_dataloader`, `test_dataloader`

---

# PyTorch Lightning API

#### LightningDataModule
]

.right-column[
```py
class MySweetData(pl.LightningDataModule):
    def __init__(self, ...):
        pass
    
    def prepare_data(self):
        # Performed only once (e.g. download data)
    
    def setup(self, stage):
        # Train/validation/test splits
        # Define any transforms
        # Create PyTorch Dataset(s)
        
    def train_dataloader(self):
        return train_dl
```
]

---

# PyTorch Lightning API

#### .gray[LightningDataModule]

#### Trainer
]

.right-column[
> Once you’ve organized your PyTorch code into a LightningModule, the Trainer automates everything else.

- There are many arguments that can be passed to your trainer to configure the training process ([docs](https://pytorch-lightning.readthedocs.io/en/stable/common/trainer.html))

```py
trainer = pl.Trainer(
  min_epochs=10,
  max_epochs=100,
  callbacks=...,
  logger=...,
  profiler='simple'
)
```
]

---

# PyTorch Lightning API

#### .gray[LightningDataModule]

#### Trainer
]

.right-column[
- Create instances of your model and data module and pass them to the `fit()` method of the `trainer` to begin the fitting process

```py
model = MySweetModel()
datamodule = MySweetData()

trainer = pl.Trainer(...)

trainer.fit(model, datamodule)

# If test logic is defined
trainer.test(model, datamodule)
```
]

---

# PyTorch Lightning API

#### .gray[LightningDataModule]

#### .gray[Trainer]

#### Hooks
]

.right-column[
- Hooks are simply points along the training-validation-test loop where you can "hook in" and perform some task

- That task can be basically anything you want

- Pre-built callbacks are defined using these hooks, by performing some task (e.g. saving a model checkpoint) at the right point in the process

- [Pseudocode of hooks](https://pytorch-lightning.readthedocs.io/en/stable/common/lightning_module.html#hooks)

- [ModelHooks API](https://pytorch-lightning.readthedocs.io/en/stable/api/pytorch_lightning.core.hooks.ModelHooks.html#pytorch_lightning.core.hooks.ModelHooks)
]

---

# PyTorch Lightning API

#### .gray[LightningDataModule]

#### .gray[Trainer]

#### .gray[Hooks]

#### Callbacks
]

```py
cb_progress = RichProgressBar()

cb_earlystop = EarlyStopping()

cb_chkpt = ModelCheckpoint()

trainer = pl.Trainer(
  ...,
  callbacks=[cb_progress, cb_earlystop, cb_chkpt]
)

trainer.fit(model, datamodule)
```
]

---

# PyTorch Lightning API

#### .gray[LightningDataModule]

#### .gray[Trainer]

#### .gray[Hooks]

#### Callbacks
]

.right-column[
- Easily create custom callbacks by subclassing `Callbacks` and defining logic at the appropriate point in the training process

```py
class EmailMe(Callback):
  def __init__(self, address):
      self.address = address
    
  def on_validation_epoch_end(self, trainer, pl_module):
      # Some code to send me an email

cb_email = EmailMe('warkentin@lunenfeld.ca')

trainer = pl.Trainer(callbacks=[cb_email])
trainer.fit(model, datamodule)
```
]

---

# PyTorch Lightning API

#### .gray[LightningDataModule]

#### .gray[Trainer]

#### .gray[Hooks]

#### .gray[Callbacks]

#### Loggers
]

.right-column[
- Loggers are used to record something to disk for later inspection
    + Comet, CSV, MLFlow, Neptune, Tensorboard, Test Tube, Weights and Biases

```py
log_csv = CSVLogger()

log_tb = TensorBoardLogger()

trainer = pl.Trainer(
  ...,
  logger=[log_csv, log_tb]
)

trainer.fit(model, datamodule)
```
]

---

# PyTorch Lightning API

#### .gray[LightningDataModule]

#### .gray[Trainer]

#### .gray[Hooks]

#### .gray[Callbacks]

#### .gray[Loggers]

#### What else?
]

- Solid documentation/examples

- Hooks everywhere

- Callbacks
    + Early stopping
    + Model checkpoints
    + Model pruning/quantization
    + Progress bars
    + Custom? Inherit from `Callback`

- Loggers
    + CSV, Tensorboard, Weights and Biases, and more...
    
- Profiling
]

---

# PyTorch Lightning API

#### .gray[LightningDataModule]

#### .gray[Trainer]

#### .gray[Hooks]

#### .gray[Callbacks]

#### .gray[Loggers]

#### What else?
]

.right-column[
- Accelerators
    + Train on non-CPU resources without changing a single line of code
    + CPU, GPU, HPU, IPU, TPU...
    + GPU-acceleration for M1 Macs (coming very soon)
        
- Plugins for mixed precision training

- Training Strategies (e.g. DDP, Bagua, etc.)

- LightningCLI for command-line use

- First-class environment support for SLURM
    + Including auto-requeing
    + Or use `test-tube` to easily construct SLURM submission scripts and enable hyperparameter tuning/grid search
]

---

# PyTorch Lightning API

#### .gray[LightningDataModule]

#### .gray[Trainer]

#### .gray[Hooks]

#### .gray[Callbacks]

#### .gray[Loggers]

#### What else?
]

.right-column[
- TorchMetrics
    > TorchMetrics is a collection of 80+ PyTorch metrics implementations and an easy-to-use API to create custom metrics.

- Lightning Flash
    > Flash is a high-level deep learning framework for fast prototyping, baselining, finetuning and solving deep learning problems.

- Lightning Bolts
    > PyTorch Lightning Bolts, is our official collection of prebuilt models across many research domains.
]

---

# <svg aria-hidden="true" role="img" viewBox="0 0 512 512" style="height:1em;width:1em;vertical-align:-0.125em;margin-left:auto;margin-right:auto;font-size:inherit;fill:currentColor;overflow:visible;position:relative;"><path d="M504 256c0 136.997-111.043 248-248 248S8 392.997 8 256C8 119.083 119.043 8 256 8s248 111.083 248 248zM262.655 90c-54.497 0-89.255 22.957-116.549 63.758-3.536 5.286-2.353 12.415 2.715 16.258l34.699 26.31c5.205 3.947 12.621 3.008 16.665-2.122 17.864-22.658 30.113-35.797 57.303-35.797 20.429 0 45.698 13.148 45.698 32.958 0 14.976-12.363 22.667-32.534 33.976C247.128 238.528 216 254.941 216 296v4c0 6.627 5.373 12 12 12h56c6.627 0 12-5.373 12-12v-1.333c0-28.462 83.186-29.647 83.186-106.667 0-58.002-60.165-102-116.531-102zM256 338c-25.365 0-46 20.635-46 46 0 25.364 20.635 46 46 46s46-20.636 46-46c0-25.365-20.635-46-46-46z"/></svg> Helpful Resources

.pull-left[
- [PyTorch Lightning Documentation](https://pytorch-lightning.readthedocs.io/en/latest/?_ga=2.113270336.699856249.1638893730-1586744328.1636324376)

- [PyTorch Lightning GitHub](https://github.com/PyTorchLightning/pytorch-lightning)

- [TorchMetrics](https://torchmetrics.readthedocs.io/en/latest/)

- [PyTorch Lightning Flash](https://lightning-flash.readthedocs.io/en/latest/)

- [PyTorch Lightning Bolts](https://pytorch-lightning.readthedocs.io/en/latest/ecosystem/bolts.html)
]

.pull-right[
<img src="imgs/Pytorch_logo.png" width="100%" /><img src="imgs/Lightning_logo.png" width="100%" />
]

---

# Live Demonstration