Efficient PyTorch — Supercharging Training Pipeline
By: Eugene Khvedchenya
Tags: Catalyst, Deep Learning, Tips And Tricks, Pytorch

The final goal of every deep learning project is to bring value to the product. Of course, we want to have the best possible model. What is “best” — depends on the particular use case, and I will leave this discussion aside from this post. I want to talk about how to get the maximum from your train.py script.

In this post we’re going to cover the following tips:

High-level frameworks instead own-made train loops Monitor progress of the training with additional metrics Use TensorBoard Visualize predictions of the model Use Dict as return value for Dataset and Model Detect anomalies and address numerical instabilities

Disclaimer: In the next section, I will include a few source-code listings. Most of them are tailored for the Catalyst framework (version 20.08) and available in pytorch-toolbelt.

Don’t re-invent the wheel

Photo by Colter Olmstead on Unsplash

Advice 1 — Leverage high-level training frameworks from PyTorch ecosystem

PyTorch offers excellent flexibility and freedom in writing your training loop from scratch. In theory, this opens an endless possibility to write any training logic. In practice, you rarely will write exotic training loops for training CycleGAN, distilling BERT, or implementing 3D object detection from scratch.

Writing a full training loop from scratch is an excellent way to learn the fundamentals of PyTorch. Yet I highly recommend switching into high-level frameworks once you get some grasp. There are plenty of options: Catalyst, PyTorch-Lightning, Fast.AI, Ignite, and others. High-level libraries save your time by:

Offering well-tested training loops

Support of the configuration files

Support for multi-GPU and distributed training

Management of checkpoints/experiments

Automatic logging of the training progress

It takes some time to get the maximum from these high-level libraries. However, this one-time investment pays off in the long-run.

Pros

The training pipeline becomes smaller — the less code — the less chance of bugs.

Easier experiment management.

Simplified distributed and mixed-precision training

Cons

Another abstraction level — as usual, when using a high-level framework, we have to write code within design principles and paradigms of that particular framework.

Time investment — learning additional framework requires time.

Show me the metrics

Photo by Mikail McVerry on Unsplash

Advice 2 — Look at additional metrics during the training

Nearly every quick-start example project for classifying images in MNIST or CIFAR and even ImageNet have one thing in common — they report a minimal set of metrics during and after the training. Most often, it’s Top-1 and Top-5 accuracy, error rate, train/validation losses, and that’s it. Although these metrics are essential, it’s just the tip of an iceberg!

Modern image classification models have tens of millions of parameters. Do you want to evaluate it using only one scalar value?

CNN classification model that has the best Top-1 accuracy may not be the best one in terms of generalization. Depending on your domain and requirements, you may want to save models with the lowest false-positive/false-negative rates or the one with the highest average precision.

Let me give you some ideas, what data you can log during the training: