What is PyTorch Lightning and how is it different than PyTorch?

PyTorch Lightning separates science code (lightning) from engineering code, allowing researchers to focus more on research. In addition, PyTorch Lightning allows you to easily scale your models by selecting how many GPUs to run your code on (distributed training), the precision level or even custom backward pass.

PyTorch Lightning can be seen as a “style-guide” that organises your PyTorch code, taking care of the engineering codes for you. PyTorch Lightning is organised in three categories:

  1. Research Code (Lightning Module)

  2. Engineering Code (No longer needed as it is handled by the Trainer)

  3. Non-essential research code (for example, logging. This goes into Callbacks)

What’s the benefit of PyTorch Lightning?

Given that PyTorch Lightning separates and takes care of the engineering code for you, this means that PyTorch Lightning automates things such as data batching, data iterations, cluster GPU training, TPU training, specific precision training, debugging, logging, etc. Essentially, all the boiler plates code that you needed to write for your deep learning project, PyTorch Lightning has extracted away most of it, giving you more time to focus on the essential part, the research codes.

How to convert PyTorch to LightningModule?

If you already have codes in PyTorch and you want to use PyTorch Lightning, you would need to convert / restructure your code into LightningModule. This code restructuring doesn’t require you to change your original code but to simply organise your code into different class methods that comes with the LightningModule. For example, in LightningModule, there’s a train_dataloader class method where you can simply copy and paste your PyTorch dataloader code into it. Similarly, there’s a training_step class method in LightningModule where you can copy and paste the codes inside your PyTorch training code, excluding the loops over epochs and batches as these are taken care for you by PyTorch Lightning. More on this with examples in the next post!

Should I start with PyTorch Lightning?

If you are new to PyTorch or deep learning, you should start with PyTorch first. Since PyTorch Lightning can be seen as more of a framework to organise PyTorch codes, it means that you still need to know the basics of PyTorch.



Data Scientist

Leave a Reply