An Introduction to PyTorch Lightning | by Harsh Maheshwari | Towards Data Science

Word on the street is that PyTorch lightning is a much better version of normal PyTorch. But what could it possibly have that it brought such consensus in our world? Well, it helps researchers scale the model with multi-GPU training, Fp-16 training, and TPU enabled training by changing just one argument.
YESSSS you read that right. By changing just one argument, you can do all of these amazing things. Remember how we used to write a multi-GPU training code and had to learn about different training architectures PyTorch supports and then implement them ourselves. Thankfully, no more of that! PyTorch Lightning will change the way we write code and will help us code all these tasks in just a second.
To cover all the changes I have structured this blog in the following way-:
Overall Comparison between PyTorch and PyTorch Lightning for the difference in implementation
Comparing each module like the model, loss function, etc. for PyTorch and PyTorch lightning
Multi GPU Training
TPU Training
FP16 Training
Early Stopping
LR Finder
The code chunks with the same color represent the implementation of the same module. For example, the model definition in both the framework is colored light green.
The first thing to notice, PyTorch lightning has incorporated train_dataloader, configure_optimizers, training_step in the class Net itself. Also, note that the whole training loop that we write in PyTorch transfers to just a few lines in PyTorch lightning.
In general, a deep learning code has the following components
Model
Data
Loss
Optimizer
Train and Test Loop
Logging
Let’ understand the difference between PyTorch and PyTorch lightning with regards to the above-mentioned components.
As we noticed above, the model architecture and definition is the same except that in Pytorch Lightning all the other function definitions are also in the same class.
There could be two ways to define the data loader in Pytorch Lightning.
You can define the train_dataloder and val_dataloader function within the Net class, as it was done earlier(in the first example)
You can define your own train_dataloader and val_dataloader as in PyTorch, to trainer.fit as shown below.
Using the above method, you can define data loaders for PyTorch as well as PyTorch lightning. The main difference between them is that the trainer.fit() in Pytorch Lightning takes all the data loaders as arguments.
trainer.fit(net, train_dataloader, val_dataloader)trainer.test(net, test_dataloader)
For n-class classification, we want to compute the cross-entropy loss. Cross-entropy is the same as NegativeLogLikelihood(log_softmax), so we will use that instead.
Let’s use the Adam Optimizer.
The self.parameters() passed here, include all the learnable parameters defined in the model.
In PyTorch, we have to
Define the training loop
Load the data
Pass the data through the model
Compute loss
Do zero_grad
Backpropagate the loss function.
However, in PyTorch lightning, we have to just
Define the training_step and validation_step, where we define how we want the data to pass through the model
Compute the loss
We don’t even have to specify model.train() and model.eval().
Tensorboard is one of the most common loggers used by us researchers. So as you notice above in PyTorch lightning, in the last line of function training_step and validation_step, self.log() is mentioned, this is used to log the training loss. This creates a folder named lightning_logs and saves all the logs and epochs there.
The model definition process is similar with two main differences. 1) All other functions are also defined with the model class itself for PyTorch lightning. 2) The nn.Module in Pytorch is overridden in PyTorch lightning by nn.LightningModule.
Data Loader can be defined in the same way. For PyTorch lightning, we have to pass train_loader, and val_loader at the time of train.fit()
Optimizer and loss can be defined the same way, but they need to be present as a function in the main class for PyTorch lightning.
The training and validation loop are pre-defined in PyTorch lightning. We have to define training_step and validation_step, i.e., given a data point/batch, how would we like to pass the data through that model.
A function for logging is pre-defined and can be directly called in Pytorch Lightning.
Now that we know the difference between these two, let’s understand why we started this comparison in the first place. To make the code scalable and reduce our engineering efforts let’s look at the features of interest.
We can do that using the code below.
trainer = Trainer(gpus=8, distributed_backend='dp')
You can define the number of GPUs you want to use for distributed training, and the backend you want to use. Here I have defined ‘dp’ which is Distributed Parallel. You can also define it as ‘ddp’, i.e. Distributed Data-Parallel.
We can do that using the code below.
trainer = Trainer(tpu_cores=[5])
This code means that the model will train on a TPU core with ID 5. We can also define how many cores we need to use by enabling multi-TPU training using a single argument.
This is my favourite. FP16 helps to speed up the training process without compromising much on performance.
trainer = Trainer(precision=16)
This is used to stop the training when you don’t see any further improvement in model performance while training the model. How do we keep check if the model performance is improving or not? We can use validation loss or accuracy for that purpose.
from pytorch_lightning.callbacks.early_stopping import EarlyStoppingdef validation_step(...):    self.log('val_loss', loss)trainer = Trainer(callbacks=[EarlyStopping(monitor='val_loss', patience=3)])
In the above example, the trainer will track the validation accuracy. If there is no improvement in performance from the past 3 epochs (value of patience), then it will stop training and hence prevent overfitting.
Learning rate is one of the most important hyperparameters, and it’s crucial to set the initial learning rate properly. Otherwise, it might happen that while training the model, the model converges to a local optimum that does not provide the best performance.
net = LightningMNISTClassifier()# finds learning rate automatically# sets hparams.lr or hparams.learning_rate to that learning ratetrainer = Trainer(auto_lr_find=True)trainer.tune(net)
To find optimal learning rate, you have to make the argument auto_lr_find True, and then tune the trainer (using trainer.tune()), this will help you find the learning rate. After that, you can call trainer.fit() for training the model.
Just by adding a few arguments, Pytorch Lightning helps us enjoy a lot of functionality, and each one of them is very crucial for us. Now that you have a basic idea, start practicing PyTorch lightning by writing your first neural network and exploring different functionality that the library offers.
If you want to learn cool applications of deep learning check out our blog on speech processing with deep learning here.
Become a Medium member to unlock and read many other stories on medium. Follow us on Medium for reading more such blog posts.