March 30, 2020

fastai and PyTorch Lightning are democratising AI

Dimitris Poulopoulos

How the two frameworks make deep learning accessible to practitioners while remaining highly hackable for researchers.

Image for post
Image by Francis Ray from Pixabay

Deep learning frameworks are hugely responsible for the breakthrough research pushing the field forward. Megaprojects such as TensorFlow and PyTorch, language-specific endeavours like DL4J and experimental ventures like Swift for TensorFlow offer a wide variety of options to practitioners and researchers.

Yet, deep learning is still too hard for software engineers with no math background, while frameworks like TensorFlow and PyTorch are complicated for data science researchers with no solid engineering background.

Image for post

To this end, libraries like Keras, fastai and PyTorch Lightning offer higher abstractions on well-established codebases. Keras mostly uses TensorFlow for its backend, while fastai and PyTorch Lightning are built on PyTorch. In this story, we examine the latter two, what they offer and what we get with the new versions; fastai 2.0 and PyTorch Lightning 0.7.1.

Learning Rate is my weekly newsletter for those who are curious about the world of AI and MLOps. You’ll hear from me every Friday with updates and thoughts on the latest AI news, research, repos and books. Subscribe here!

fastai

fastai is a deep learning library developed mainly by Jeremy Howard and Sylvain Gugger and maintained by hundreds of contributors. It enables practitioners to achieve state-of-the-art results effortlessly and facilitates deep experimentation for researchers. The framework was introduced in 2017 to support the associated course and the second version is expected in July 2020.

Introduction

In this story, we mostly talk about the upcoming, second version of the fastai library. With that in mind, fastai is designed in a layered approach. It provides practitioners with high-level components that deliver state-of-the-art results quickly and researchers with low-level blocks that can be composed in various ways, enabling experimentation with new deep learning approaches [1].

The higher-level APIs are useful to beginners or AI practitioners who mainly want to apply pre-existing deep learning techniques. fastai supports four application domains by default: computer vision, natural language processing, tabular data or time-series and recommender systems via collaborative filtering. Furthermore, fastai is an opinionated framework offering customizable models with sensible defaults. Pre-defined hyper-parameters follow best practices and are informed by relevant state-of-the-art research where available. This design leads to clear, easy to follow code, tainted with fewer bugs if any.

On the other hand, lower-level APIs support mid and high-level functionality with optimized primitives. Those primitives are mainly developed on top of popular Python libraries, such as PyTorch, numpy, pandas or PIL. However, in order to be hackable, fastai interoperates seamlessly with each of the underlying libraries. For instance, fastai makes it easy to inject its components into PyTorch modules or use original PyTorch functionality with fastai.

In a nutshell, the goal of fastai is to have a flat learning curve for beginners but to be flexible enough for researchers.

Code example

To better understand the design choices for each library we implement a state-of-the-art classifier for the MNIST dataset. After all, MNIST is arguably the “Hello world” example for deep learning.

Image for post

The framework’s high-level APIs permits us to solve the MNIST challenge in more or less five lines of code. Despite this being a toy example, it highlights the possibilities of the layered architecture. Furthermore, several key elements of training deep learning models are hidden from us; the optimal learning rate, batch size and other hyper-parameters are set automatically, while the loss function is inferred by the data. Transfer learning is handled without sweat, with the ResNet architecture automatically adjusted to the problem at hand.

It took five lines of code and less than five minutes of GPU training to achieve > 98% accuracy on the MNIST dataset. Practitioners facing challenges supported by default in fastai (e.g. vision, text, tabular data or collaborative filtering) may have similar experiences. In any other case, much can be achieved with just a few tweaks.

New features in version 2.0

fastai version 2.0 is a complete rewrite of the first version. There are many breaking changes in the codebase and you can walk through most of them in the associated YouTube series by Jeremy Howard.

However, the most notable modifications are summarized below:

You can get a small taste of the new APIs in the second part of the previous year’s course, while more examples are expected when the 2020’s course is available for streaming.

PyTorch Lightning

PyTorch Lightning is a lightweight PyTorch wrapper for researchers, developed by William Falcon. However, its fast adoption by the research teams working for major organizations (e.g. FAIR) secured funding from venture capitalists and a full-time team is formed as we speak. PyTorch Lightning is more like a style-guide which helps you decouple the science code from the engineering.

Introduction

PyTorch Lightning is a flexible, light-weight wrapper on PyTorch, that sets a standard on how to structure your deep learning code. This way, it handles most of the engineering work, leaving you to focus on the science. This approach leads to less boilerplate code, thus, fewer worries and bugs.

Being a thin wrapper means that researchers or production teams do not need to learn the perks of another framework. At the same time, it offers great flexibility to try out cutting edge ideas while abstracting the engineering details.

The main modules of PyTorch Lightning are the LightningModule and the Trainer. All science logic goes into the LightningModule; data preparation, optimizer initialization, the training loop and the forward pass. On the other hand, all the engineering effort is hidden away into the Trainer module; fitting, automatic logging, checkpointing, early stopping and much more.

As its creator states, the goals of PyTorch lightning is the promotion of best practices in the deep learning community and the facilitation of research reproducibility.

Code example

Let us now see how to solve MNIST using PyTorch Lightning. We somehow need to use the same pre-trained model to avoid cheating. To this end, continuing from the previous example, we will grab the model from the learner (i.e. fastai learner) before training. We do that because the learner has already adjusted the model to the MNIST challenge. If we just grab a new ResNet-34 model it will be tailored to solve the ImageNet challenge, meaning different number of outputs for the last layer. One thing we need to change is to tell the first convolution to expect a one-channel input.

We are now ready to implement an MNISTSolver.

To fit the model we just instantiate the Trainer and we are ready to go.

PyTorch Lightning offers logging to TensoBoard — among others — by default. Just run the next magic command into Jupyter.

%reload_ext tensorboard
%tensorboard --logdir lightning_logs/

Image for post
Image for post

New features

The new stable version for PyTorch Lightning is 0.7.1. There we can find exciting new features, summarised below:

Visit the project’s GitHub page for more information and the documentation page to dive deeper.

Conclusion

In this story, we saw how fastai and PyTorch Lightning make it easy for practitioners and researchers to implement deep learning solutions.

fastai offers a higher level of abstraction that leads you to state-of-the-art results quicker. On the other hand, PyTorch Lightning will have a flatter learning curve to those familiar with PyTorch, offers TPU support and integration with many loggers.

In the end, what you choose to go with depends on the challenges you face. In any case, you cannot go wrong!

My name is Dimitris Poulopoulos and I’m a machine learning researcher at BigDataStack and PhD(c) at the University of Piraeus, Greece. I have worked on designing and implementing AI and software solutions for major clients such as the European Commission, Eurostat, IMF, the European Central Bank, OECD, and IKEA. If you are interested in reading more posts about Machine Learning, Deep Learning and Data Science, follow me on Medium, LinkedIn or @james2pl on twitter.

References

[1] Howard, Jeremy, and Sylvain Gugger. “fastai: A Layered API for Deep Learning.” Information 11.2 (2020): 108.