Generally, an external script should be used to perform the re-encoding.

Additionally, we provide APIs such as the utility class, VideoClips, that simplifies the task of enumerating all possible clips of fixed size in a list of video files by creating an index of all clips in a set of videos. It also allows you to specify a fixed frame-rate for the videos. An example of the API is provided below:
from torchvision.datasets.video_utils import VideoClips

class MyVideoDataset(object):
    def __init__(self, video_paths):
        self.video_clips = VideoClips(video_paths,
                                      clip_length_in_frames=16,
                                      frames_between_clips=1,
                                      frame_rate=15)

    def __getitem__(self, idx):
        video, audio, info, video_idx = self.video_clips.get_clip(idx)
        return video, audio

    def __len__(self):
        return self.video_clips.num_clips()


return self.video_clips.num_clips()
```
Most of the user-facing API is in Python, similar to PyTorch, which makes it easily extensible. Plus, the underlying implementation is fast — torchvision decodes as little as possible from the video on-the-fly in order to return a clip from the video.
Check out the torchvision 0.4 release notes here for more details.
We look forward to continuing our collaboration with the community and hearing your feedback as we further improve and expand the PyTorch deep learning platform.
We’d like to thank the entire PyTorch team and the community for all of the contributions to this work!


layout: blog_detail
title: 'How to Train State-Of-The-Art Models Using TorchVision’s Latest Primitives'
author: Vasilis Vryniotis
featured-img: 'assets/images/fx-image2.png'


A few weeks ago, TorchVision v0.11 was released packed with numerous new primitives, models and training recipe improvements which allowed achieving state-of-the-art (SOTA) results. The project was dubbed “TorchVision with Batteries Included” and aimed to modernize our library. We wanted to enable researchers to reproduce papers and conduct research more easily by using common building blocks. Moreover, we aspired to provide the necessary tools to Applied ML practitioners to train their models on their own data using the same SOTA techniques as in research. Finally, we wanted to refresh our pre-trained weights and offer better off-the-shelf models to our users, hoping that they would build better applications.


Though there is still much work to be done, we wanted to share with you some exciting results from the above work. We will showcase how one can use the new tools included in TorchVision to achieve state-of-the-art results on a highly competitive and well-studied architecture such as ResNet50 [1]. We will share the exact recipe used to improve our baseline by over 4.7 accuracy points to reach a final top-1 accuracy of 80.9% and share the journey for deriving the new training process. Moreover, we will show that this recipe generalizes well to other model variants and families. We hope that the above will influence future research for developing stronger generalizable training methodologies and will inspire the community to adopt and contribute to our efforts.
The Results
Using our new training recipe found on ResNet50, we’ve refreshed the pre-trained weights of the following models:

Model
Accuracy@1
Accuracy@5


|----------|:--------:|:----------:|
| ResNet50    | 80.858 | 95.434|
|----------|:--------:|:----------:|
| ResNet101    | 81.886 | 95.780|
|----------|:--------:|:----------:|
| ResNet152    | 82.284 | 96.002|
|----------|:--------:|:----------:|
| ResNeXt50-32x4d    | 81.198 | 95.340|
Note that the accuracy of all models except RetNet50 can be further improved by adjusting their training parameters slightly, but our focus was to have a single robust recipe which performs well for all.
UPDATE: We have refreshed the majority of popular classification models of TorchVision, you can find the details on this blog post.
There are currently two ways to use the latest weights of the model.
Using the Multi-pretrained weight API


Using the Multi-pretrained weight API
We are currently working on a new prototype mechanism which will extend the model builder methods of TorchVision to support multiple weights. Along with the weights, we store useful meta-data (such as the labels, the accuracy, links to recipe etc) and the preprocessing transforms necessary for using the models. Example:
```python
  from PIL import Image
  from torchvision import prototype as P
  img = Image.open("test/assets/encode_jpeg/grace_hopper_517x606.jpg")

  # Initialize model
  weights = P.models.ResNet50_Weights.IMAGENET1K_V2
  model = P.models.resnet50(weights=weights)
  model.eval()
# Initialize inference transforms
  preprocess = weights.transforms()

  # Apply inference preprocessing transforms
  batch = preprocess(img).unsqueeze(0)


batch = preprocess(img).unsqueeze(0)
  prediction = model(batch).squeeze(0).softmax(0)

  # Make predictions
  label = prediction.argmax().item()
  score = prediction[label].item()

  # Use meta to get the labels
  category_name = weights.meta['categories'][label]
  print(f"{category_name}: {100 * score}%")

## Using the legacy API

Those who don’t want to use a prototype API have the option of accessing the new weights via the legacy API using the following approach:

```python
  from torchvision.models import resnet

  # Overwrite the URL of the previous weights
  resnet.model_urls["resnet50"] = "https://download.pytorch.org/models/resnet50-11ad3fa6.pth"

  # Initialize the model using the legacy API
  model = resnet.resnet50(pretrained=True)

  # TODO: Apply preprocessing + call the model
  # ...

The Training Recipe


...
```
The Training Recipe
Our goal was to use the newly introduced primitives of TorchVision to derive a new strong training recipe which achieves state-of-the-art results for the vanilla ResNet50 architecture when trained from scratch on ImageNet with no additional external data. Though by using architecture specific tricks [2] one could further improve the accuracy, we’ve decided not to include them so that the recipe can be used in other architectures. Our recipe heavily focuses on simplicity and builds upon work by FAIR [3], [4], [5], [6], [7]. Our findings align with the parallel study of Wightman et al. [7], who also report major accuracy improvements by focusing on the training recipes.


Without further ado, here are the main parameters of our recipe:
  # Optimizer & LR scheme
  ngpus=8,
  batch_size=128,  # per GPU

  epochs=600,
  opt='sgd',
  momentum=0.9,

  lr=0.5,
  lr_scheduler='cosineannealinglr',
  lr_warmup_epochs=5,
  lr_warmup_method='linear',
  lr_warmup_decay=0.01,

  # Regularization and Augmentation
  weight_decay=2e-05,
  norm_weight_decay=0.0,

  label_smoothing=0.1,
  mixup_alpha=0.2,
  cutmix_alpha=1.0,
  auto_augment='ta_wide',
  random_erase=0.1,

  ra_sampler=True,
  ra_reps=4,

  # EMA configuration
  model_ema=True,
  model_ema_steps=32,
  model_ema_decay=0.99998,

  # Resizing
  interpolation='bilinear',
  val_resize_size=232,
  val_crop_size=224,
  train_crop_size=176,

Using our standard training reference script, we can train a ResNet50 using the following command:
```
torchrun --nproc_per_node=8 train.py --model resnet50 --batch-size 128 --lr 0.5 \


--lr-scheduler cosineannealinglr --lr-warmup-epochs 5 --lr-warmup-method linear \
--auto-augment ta_wide --epochs 600 --random-erase 0.1 --weight-decay 0.00002 \
--norm-weight-decay 0.0 --label-smoothing 0.1 --mixup-alpha 0.2 --cutmix-alpha 1.0 \
--train-crop-size 176 --model-ema --val-resize-size 232 --ra-sampler --ra-reps 4
```
Methodology
There are a few principles we kept in mind during our explorations:

Training is a stochastic process and the validation metric we try to optimize is a random variable. This is due to the random weight initialization scheme employed and the existence of random effects during the training process. This means that we can’t do a single run to assess the effect of a recipe change. The standard practice is doing multiple runs (usually 3 to 5) and studying the summarization stats (such as mean, std, median, max, etc).


There is usually a significant interaction between different parameters, especially for techniques that focus on Regularization and reducing overfitting. Thus changing the value of one can have effects on the optimal configurations of others. To account for that one can either adopt a greedy search approach (which often leads to suboptimal results but tractable experiments) or apply grid search (which leads to better results but is computationally expensive). In this work, we used a mixture of both.
Techniques that are non-deterministic or introduce noise usually require longer training cycles to improve model performance. To keep things tractable, we initially used short training cycles (small number of epochs) to decide which paths can be eliminated early and which should be explored using longer training.


There is a risk of overfitting the validation dataset [8] because of the repeated experiments. To mitigate some of the risk, we apply only training optimizations that provide a significant accuracy improvements and use K-fold cross validation to verify optimizations done on the validation set. Moreover we confirm that our recipe ingredients generalize well on other models for which we didn’t optimize the hyper-parameters.

Break down of key accuracy improvements


Break down of key accuracy improvements
As discussed in earlier blogposts, training models is not a journey of monotonically increasing accuracies and the process involves a lot of backtracking. To quantify the effect of each optimization, below we attempt to show-case an idealized linear journey of deriving the final recipe starting from the original recipe of TorchVision. We would like to clarify that this is an oversimplification of the actual path we followed and thus it should be taken with a grain of salt.

In the table below, we provide a summary of the performance of stacked incremental improvements on top of Baseline. Unless denoted otherwise, we report the model with best Acc@1 out of 3 runs:


Accuracy@1
Accuracy@5
Incremental Diff
Absolute Diff

ResNet50 Baseline
76.130
92.862
0.000
0.000

----------
:--------:
:----------:
:---------
:--------:

+ LR optimizations
76.494
93.198
0.364
0.364

----------
:--------:
:----------:
:---------
:--------:

+ TrivialAugment
76.806
93.272
0.312
0.676

----------
:--------:
:----------:
:---------
:--------:

+ Long Training
78.606
94.052
1.800
2.476

----------
:--------:
:----------:
:---------
:--------:

+ Random Erasing
78.796
94.094
0.190
2.666

----------
:--------:
:----------:
:---------
:--------:

+ Label Smoothing
79.114
94.374
0.318
2.984

----------
:--------:
:----------:
:---------
:--------:

+ Mixup
79.232
94.536
0.118
3.102

----------
:--------:
:----------:
:---------
:--------:

+ Cutmix
79.510
94.642
0.278
3.380

----------
:--------:
:----------:
:---------
:--------:


+ Weight Decay tuning
80.036
94.746
0.526
3.906

+ FixRes mitigations
80.196
94.672
0.160
4.066

----------
:--------:
:----------:
:---------
:--------:

+ EMA
80.450
94.908
0.254
4.320

----------
:--------:
:----------:
:---------
:--------:

+ Inference Resize tuning *
80.674
95.166
0.224
4.544

----------
:--------:
:----------:
:---------
:--------:

+ Repeated Augmentation **
80.858
95.434
0.184
4.728

*The tuning of the inference size was done on top of the last model. See below for details.
** Community contribution done after the release of the article. See below for details.
Baseline
Our baseline is the previously released ResNet50 model of TorchVision. It was trained with the following recipe:
```python
  # Optimizer & LR scheme
  ngpus=8,
  batch_size=32,  # per GPU
epochs=90,
  opt='sgd',
  momentum=0.9,
lr=0.1,
  lr_scheduler='steplr',
  lr_step_size=30,
  lr_gamma=0.1,


lr_step_size=30,
  lr_gamma=0.1,
# Regularization
  weight_decay=1e-4,
# Resizing
  interpolation='bilinear',
  val_resize_size=256,
  val_crop_size=224,
  train_crop_size=224,
```
Most of the above parameters are the defaults on our training scripts. We will start building on top of this baseline by introducing optimizations until we gradually arrive at the final recipe.
LR optimizations


LR optimizations
There are a few parameter updates we can apply to improve both the accuracy and the speed of our training. This can be achieved by increasing the batch size and tuning the LR. Another common method is to apply warmup and gradually increase our learning rate. This is beneficial especially when we use very high learning rates and helps with the stability of the training in the early epochs. Finally, another optimization is to apply Cosine Schedule to adjust our LR during the epochs. A big advantage of cosine is that there are no hyper-parameters to optimize, which cuts down our search space.
Here are the additional optimizations applied on top of the baseline recipe. Note that we’ve run multiple experiments to determine the optimal configuration of the parameters:
  batch_size=128,  # per GPU

  lr=0.5,
  lr_scheduler='cosineannealinglr',
  lr_warmup_epochs=5,
  lr_warmup_method='linear',
  lr_warmup_decay=0.01,


lr_warmup_decay=0.01,

The above optimizations increase our top-1 Accuracy by 0.364 points comparing to the baseline. Note that in order to combine the different LR strategies we use the newly introduced [SequentialLR](https://pytorch.org/docs/stable/generated/torch.optim.lr_scheduler.SequentialLR.html#torch.optim.lr_scheduler.SequentialLR) scheduler.

## TrivialAugment

The original model was trained using basic augmentation transforms such as Random resized crops and horizontal flips. An easy way to improve our accuracy is to apply more complex “Automatic-Augmentation” techniques. The one that performed best for us is TrivialAugment [[9]](https://arxiv.org/abs/2103.10158), which is extremely simple and can be considered “parameter free”, which means it can help us cut down our search space further.

Here is the update applied on top of the previous step:

auto_augment='ta_wide',
```
The use of TrivialAugment increased our top-1 Accuracy by 0.312 points compared to the previous step.


Long Training
Longer training cycles are beneficial when our recipe contains ingredients that behave randomly. More specifically as we start adding more and more techniques that introduce noise, increasing the number of epochs becomes crucial. Note that at early stages of our exploration, we used relatively short cycles of roughly 200 epochs which was later increased to 400 as we started narrowing down most of the parameters and finally increased to 600 epochs at the final versions of the recipe.
Below we see the update applied on top of the earlier steps:
epochs=600,


epochs=600,

This further increases our top-1 Accuracy by 1.8 points on top of the previous step. This is the biggest increase we will observe in this iterative process. It’s worth noting that the effect of this single optimization is overstated and somehow misleading. Just increasing the number of epochs on top of the old baseline won’t yield such significant improvements. Nevertheless the combination of the LR optimizations with strong Augmentation strategies helps the model benefit from longer cycles. It’s also worth mentioning that the reason we introduce the lengthy training cycles so early in the process is because in the next steps we will introduce techniques that require significantly more epochs to provide good results.
Random Erasing


Random Erasing
Another data augmentation technique known to help the classification accuracy is Random Erasing [10], [11]. Often paired with Automatic Augmentation methods, it usually yields additional improvements in accuracy due to its regularization effect. In our experiments we tuned only the probability of applying the method via a grid search and found that it’s beneficial to keep its probability at low levels, typically around 10%.
Here is the extra parameter introduced on top of the previous:
random_erase=0.1,

Applying Random Erasing increases our Acc@1 by further 0.190 points.
Label Smoothing


Label Smoothing
A good technique to reduce overfitting is to stop the model from becoming overconfident. This can be achieved by softening the ground truth using Label Smoothing [12]. There is a single parameter which controls the degree of smoothing (the higher the stronger) that we need to specify. Though optimizing it via grid search is possible, we found that values around 0.05-0.15 yield similar results, so to avoid overfitting it we used the same value as on the paper that introduced it.
Below we can find the extra config added on this step:
label_smoothing=0.1,

We use PyTorch’s newly introduced CrossEntropyLoss label_smoothing parameter and that increases our accuracy by an additional 0.318 points.
Mixup and Cutmix


Mixup and Cutmix
Two data augmentation techniques often used to produce SOTA results are Mixup and Cutmix [13], [14]. They both provide strong regularization effects by softening not only the labels but also the images. In our setup we found it beneficial to apply one of them randomly with equal probability. Each is parameterized with a hyperparameter alpha, which controls the shape of the Beta distribution from which the smoothing probability is sampled. We did a very limited grid search, focusing primarily on common values proposed on the papers.
Below you will find the optimal values for the alpha parameters of the two techniques:
mixup_alpha=0.2,
cutmix_alpha=1.0,

Applying mixup increases our accuracy by 0.118 points and combining it with cutmix improves it by additional 0.278 points.
Weight Decay tuning


Weight Decay tuning
Our standard recipe uses L2 regularization to reduce overfitting. The Weight Decay parameter controls the degree of the regularization (the larger the stronger) and is applied universally to all learned parameters of the model by default. In this recipe, we apply two optimizations to the standard approach. First we perform grid search to tune the parameter of weight decay and second we disable weight decay for the parameters of the normalization layers.
Below you can find the optimal configuration of weight decay for our recipe:
weight_decay=2e-05,
norm_weight_decay=0.0,

The above update improves our accuracy by a further 0.526 points, providing additional experimental evidence for a known fact that tuning weight decay has significant effects on the performance of the model. Our approach for separating the Normalization parameters from the rest was inspired by ClassyVision’s approach.
FixRes mitigations


FixRes mitigations
An important property identified early in our experiments is the fact that the models performed significantly better if the resolution used during validation was increased from the 224x224 of training. This effect is studied in detail on the FixRes paper [5] and two mitigations are proposed: a) one could try to reduce the training resolution so that the accuracy on the validation resolution is maximized or b) one could fine-tune the model on a two-phase training so that it adjusts on the target resolution. Since we didn’t want to introduce a 2-phase training, we went for option a). This means that we reduced the train crop size from 224 and used grid search to find the one that maximizes the validation on resolution of 224x224.
Below you can see the optimal value used on our recipe:
val_crop_size=224,
train_crop_size=176,

The above optimization improved our accuracy by an additional 0.160 points and sped up our training by 10%.


It’s worth noting that the FixRes effect still persists, meaning that the model continues to perform better on validation when we increase the resolution. Moreover, further reducing the training crop-size actually hurts the accuracy. This intuitively makes sense because one can only reduce the resolution so much before critical details start disappearing from the picture. Finally, we should note that the above FixRes mitigation seems to benefit models with similar depth to ResNet50. Deeper variants with larger receptive fields seem to be slightly negatively affected (typically by 0.1-0.2 points). Hence we consider this part of the recipe optional. Below we visualize the performance of the best available checkpoints (with the full recipe) for models trained with 176 and 224 resolution:


Exponential Moving Average (EMA)
EMA is a technique that allows one to push the accuracy of a model without increasing its complexity or inference time. It performs an exponential moving average on the model weights and this leads to increased accuracy and more stable models. The averaging happens every few iterations and its decay parameter was tuned via grid search.
Below you can see the optimal values for our recipe:
model_ema=True,
model_ema_steps=32,
model_ema_decay=0.99998,


model_ema_steps=32,
model_ema_decay=0.99998,
```
The use of EMA increases our accuracy by 0.254 points comparing to the previous step. Note that TorchVision’s EMA implementation is build on top of PyTorch’s AveragedModel class with the key difference being that it averages not only the model parameters but also its buffers. Moreover, we have adopted tricks from Pycls which allow us to parameterize the decay in a way that doesn’t depend on the number of epochs.
Inference Resize tuning


Inference Resize tuning
Unlike all other steps of the process which involved training models with different parameters, this optimization was done on top of the final model. During inference, the image is resized to a specific resolution and then a central 224x224 crop is taken from it. The original recipe used a resize size of 256, which caused a similar discrepancy as the one described on the FixRes paper [5]. By bringing this resize value closer to the target inference resolution, one can improve the accuracy. To select the value we run a short grid search between interval [224, 256] with step of 8. To avoid overfitting, the value was selected using half of the validation set and confirmed using the other half.
Below you can see the optimal value used on our recipe:
val_resize_size=232,


val_resize_size=232,

The above is an optimization which improved our accuracy by 0.224 points. It’s worth noting that the optimal value for ResNet50 works also best for ResNet101, ResNet152 and ResNeXt50, which hints that it generalizes across models:

[UPDATE] Repeated Augmentation


[UPDATE] Repeated Augmentation
Repeated Augmentation [15], [16] is another technique which can improve the overall accuracy and has been used by other strong recipes such as those at [6], [7]. Tal Ben-Nun, a community contributor, has further improved upon our original recipe by proposing training the model with 4 repetitions. His contribution came after the release of this article.
Below you can see the optimal value used on our recipe:
ra_sampler=True,
ra_reps=4,

The above is the final optimization which improved our accuracy by 0.184 points.
Optimizations that were tested but not adopted


Optimizations that were tested but not adopted
During the early stages of our research, we experimented with additional techniques, configurations and optimizations. Since our target was to keep our recipe as simple as possible, we decided not to include anything that didn’t provide a significant improvement. Here are a few approaches that we took but didn’t make it to our final recipe:

Optimizers: Using more complex optimizers such as Adam, RMSProp or SGD with Nesterov momentum didn’t provide significantly better results than vanilla SGD with momentum.
LR Schedulers: We tried different LR Scheduler schemes such as StepLR and Exponential. Though the latter tends to work better with EMA, it often requires additional hyper-parameters such as defining the minimum LR to work well. Instead, we just use cosine annealing decaying the LR up to zero and choose the checkpoint with the highest accuracy.


Automatic Augmentations: We’ve tried different augmentation strategies such as AutoAugment and RandAugment. None of these outperformed the simpler parameter-free TrivialAugment.
Interpolation: Using bicubic or nearest interpolation didn’t provide significantly better results than bilinear.
Normalization layers: Using Sync Batch Norm didn’t yield significantly better results than using the regular Batch Norm.

Acknowledgements
We would like to thank Piotr Dollar, Mannat Singh and Hugo Touvron for providing their insights and feedback during the development of the recipe and for their previous research work on which our recipe is based on. Their support was invaluable for achieving the above result. Moreover, we would like to thank Prabhat Roy, Kai Zhang, Yiwen Song, Joel Schlosser, Ilqar Ramazanli, Francisco Massa, Mannat Singh, Xiaoliang Dai, Samuel Gabriel, Allen Goodman and Tal Ben-Nun for their contributions to the Batteries Included project.
References


References

Kaiming He, Xiangyu Zhang, Shaoqing Ren, Jian Sun. “Deep Residual Learning for Image Recognition”.
Tong He, Zhi Zhang, Hang Zhang, Zhongyue Zhang, Junyuan Xie, Mu Li. “Bag of Tricks for Image Classification with Convolutional Neural Networks”
Piotr Dollár, Mannat Singh, Ross Girshick. “Fast and Accurate Model Scaling”
Tete Xiao, Mannat Singh, Eric Mintun, Trevor Darrell, Piotr Dollár, Ross Girshick. “Early Convolutions Help Transformers See Better”
Hugo Touvron, Andrea Vedaldi, Matthijs Douze, Hervé Jégou. “Fixing the train-test resolution discrepancy
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, Hervé Jégou. “Training data-efficient image transformers & distillation through attention”
Ross Wightman, Hugo Touvron, Hervé Jégou. “ResNet strikes back: An improved training procedure in timm”
Benjamin Recht, Rebecca Roelofs, Ludwig Schmidt, Vaishaal Shankar. “Do ImageNet Classifiers Generalize to ImageNet?”


Samuel G. Müller, Frank Hutter. “TrivialAugment: Tuning-free Yet State-of-the-Art Data Augmentation”
Zhun Zhong, Liang Zheng, Guoliang Kang, Shaozi Li, Yi Yang. “Random Erasing Data Augmentation”
Terrance DeVries, Graham W. Taylor. “Improved Regularization of Convolutional Neural Networks with Cutout”
Christian Szegedy, Vincent Vanhoucke, Sergey Ioffe, Jon Shlens, Zbigniew Wojna. “Rethinking the Inception Architecture for Computer Vision”
Hongyi Zhang, Moustapha Cisse, Yann N. Dauphin, David Lopez-Paz. “mixup: Beyond Empirical Risk Minimization”
Sangdoo Yun, Dongyoon Han, Seong Joon Oh, Sanghyuk Chun, Junsuk Choe, Youngjoon Yoo. “CutMix: Regularization Strategy to Train Strong Classifiers with Localizable Features”
Elad Hoffer, Tal Ben-Nun, Itay Hubara, Niv Giladi, Torsten Hoefler, Daniel Soudry. “Augment your batch: better training with larger batches”


Maxim Berman, Hervé Jégou, Andrea Vedaldi, Iasonas Kokkinos, Matthijs Douze. “Multigrain: a unified image embedding for classes and instances”


layout: blog_detail
title: 'Updates & Improvements to PyTorch Tutorials'
author: Team PyTorch

PyTorch.org provides researchers and developers with documentation, installation instructions, latest news, community projects, tutorials, and more. Today, we are introducing usability and content improvements including tutorials in additional categories, a new recipe format for quickly referencing common topics, sorting using tags, and an updated homepage.
Let’s take a look at them in detail.
TUTORIALS HOME PAGE UPDATE
The tutorials home page now provides clear actions that developers can take. For new PyTorch users, there is an easy-to-discover button to take them directly to “A 60 Minute Blitz”. Right next to it, there is a button to view all recipes which are designed to teach specific features quickly with examples.


In addition to the existing left navigation bar, tutorials can now be quickly filtered by multi-select tags. Let’s say you want to view all tutorials related to “Production” and “Quantization”. You can select the “Production” and “Quantization” filters as shown in the image shown below:

The following additional resources can also be found at the bottom of the Tutorials homepage:
* PyTorch Cheat Sheet
* PyTorch Examples
* Tutorial on GitHub
PYTORCH RECIPES
Recipes are new bite-sized, actionable examples designed to teach researchers and developers how to use specific PyTorch features. Some notable new recipes include:
* Loading Data in PyTorch


Model Interpretability Using Captum
How to Use TensorBoard

View the full recipes here.
LEARNING PYTORCH
This section includes tutorials designed for users new to PyTorch. Based on community feedback, we have made updates to the current Deep Learning with PyTorch: A 60 Minute Blitz tutorial, one of our most popular tutorials for beginners. Upon completion, one can understand what PyTorch and neural networks are, and be able to build and train a simple image classification network. Updates include adding explanations to clarify output meanings and linking back to where users can read more in the docs, cleaning up confusing syntax errors, and reconstructing and explaining new concepts for easier readability.


DEPLOYING MODELS IN PRODUCTION
This section includes tutorials for developers looking to take their PyTorch models to production. The tutorials include:
* Deploying PyTorch in Python via a REST API with Flask
* Introduction to TorchScript
* Loading a TorchScript Model in C++
* Exporting a Model from PyTorch to ONNX and Running it using ONNX Runtime
FRONTEND APIS
PyTorch provides a number of frontend API features that can help developers to code, debug, and validate their models more efficiently. This section includes tutorials that teach what these features are and how to use them. Some tutorials to highlight:


Introduction to Named Tensors in PyTorch
Using the PyTorch C++ Frontend
Extending TorchScript with Custom C++ Operators
Extending TorchScript with Custom C++ Classes
Autograd in C++ Frontend

MODEL OPTIMIZATION
Deep learning models often consume large amounts of memory, power, and compute due to their complexity. This section provides tutorials for model optimization:
* Pruning
* Dynamic Quantization on BERT


Static Quantization with Eager Mode in PyTorch

PARALLEL AND DISTRIBUTED TRAINING
PyTorch provides features that can accelerate performance in research and production such as native support for asynchronous execution of collective operations and peer-to-peer communication that is accessible from Python and C++. This section includes tutorials on parallel and distributed training:
* Single-Machine Model Parallel Best Practices
* Getting started with Distributed Data Parallel
* Getting started with Distributed RPC Framework
* Implementing a Parameter Server Using Distributed RPC Framework


Making these improvements are just the first step of improving PyTorch.org for the community. Please submit your suggestions here.
Cheers,
Team PyTorch


layout: blog_detail
title: "PyTorch 2.0: Our next generation release that is faster, more Pythonic and Dynamic as ever"

We are excited to announce the release of PyTorch® 2.0 which we highlighted during the PyTorch Conference on 12/2/22! PyTorch 2.0 offers the same eager-mode development and user experience, while fundamentally changing and supercharging how PyTorch operates at compiler level under the hood with faster performance and support for Dynamic Shapes and Distributed.


This next-generation release includes a Stable version of Accelerated Transformers (formerly called Better Transformers); Beta includes torch.compile as the main API for PyTorch 2.0,  the scaled_dot_product_attention function as part of torch.nn.functional, the MPS backend, functorch APIs in the torch.func module; and other Beta/Prototype improvements across various inferences, performance and training optimization features on GPUs and CPUs. For a comprehensive introduction and technical overview of torch.compile, please visit the 2.0 Get Started page.
Along with 2.0, we are also releasing a series of beta updates to the PyTorch domain libraries, including those that are in-tree, and separate libraries including TorchAudio, TorchVision, and TorchText. An update for TorchX is also being released as it moves to community supported mode. More details can be found in this library blog.


This release is composed of over 4,541 commits and 428 contributors since 1.13.1. We want to sincerely thank our dedicated community for your contributions. As always, we encourage you to try these out and report any issues as we improve 2.0 and the overall 2-series this year.
Summary:
* torch.compile is the main API for PyTorch 2.0, which wraps your model and returns a compiled model. It is a fully additive (and optional) feature and hence 2.0 is 100% backward compatible by definition.
* As an underpinning technology of torch.compile, TorchInductor with Nvidia and AMD GPUs will rely on OpenAI Triton deep learning compiler to generate performant code and hide low level hardware details. OpenAI Triton-generated kernels achieve performance that's on par with hand-written kernels and specialized cuda libraries such as cublas.


Accelerated Transformers introduce high-performance support for training and inference using a custom kernel architecture for scaled dot product attention (SPDA). The API is integrated with torch.compile() and model developers may also use the scaled dot product attention kernels directly by calling the new scaled_dot_product_attention() operator.
Metal Performance Shaders (MPS) backend provides GPU accelerated PyTorch training on Mac platforms with added support for Top 60 most used ops, bringing coverage to over 300 operators.
Amazon AWS optimizes the PyTorch CPU inference on AWS Graviton3 based C7g instances. PyTorch 2.0 improves inference performance on Graviton compared to the previous releases, including improvements for Resnet50 and Bert.


New prototype features and technologies across TensorParallel, DTensor, 2D parallel, TorchDynamo, AOTAutograd, PrimTorch and TorchInductor.

Stable

Beta

Prototype

Performance Improvements

Accelerated PT 2 Transformers

torch.compile

DTensor

CUDA support for 11.7 & 11.8 (deprecating CUDA 11.6)

PyTorch MPS Backend

TensorParallel


Python 3.8 (deprecating Python 3.7)

Scaled dot product attention

2D Parallel

AWS Graviton3

functorch

Torch.compile (dynamic=True)

Dispatchable Collectives

Torch.set_default & torch.device


X86 quantization backend

GNN inference and training performance

*To see a full list of public 2.0, 1.13 and 1.12 feature submissions click here.
Stable Features
[Stable] Accelerated PyTorch 2 Transformers


[Stable] Accelerated PyTorch 2 Transformers
The PyTorch 2.0 release includes a new high-performance implementation of the PyTorch Transformer API. In releasing Accelerated PT2 Transformers, our goal is to make training and deployment of state-of-the-art Transformer models affordable across the industry. This release introduces high-performance support for training and inference using a custom kernel architecture for scaled dot product attention (SPDA), extending the inference “fastpath” architecture, previously known as "Better Transformer."
Similar to the “fastpath” architecture, custom kernels are fully integrated into the PyTorch Transformer API – thus, using the native Transformer and MultiHeadAttention API will enable users to:

transparently see significant speed improvements;
support many more use cases including models using Cross-Attention, Transformer Decoders, and for training models; and


continue to use fastpath inference for fixed and variable sequence length Transformer Encoder and Self Attention use cases.

To take full advantage of different hardware models and Transformer use cases, multiple SDPA custom kernels are supported (see below), with custom kernel selection logic that will pick the highest-performance kernel for a given model and hardware type. In addition to the existing Transformer API, model developers may also use the scaled dot product attention kernels directly by calling the new scaled_dot_product_attention() operator. Accelerated PyTorch 2 Transformers are integrated with torch.compile() .  To use your model while benefiting from the additional acceleration of PT2-compilation (for inference or training), pre-process the model with model = torch.compile(model).


We have achieved major speedups for training transformer models and in particular large language models with Accelerated PyTorch 2 Transformers using a combination of custom kernels and torch.compile().

Figure: Using scaled dot product attention with custom kernels and torch.compile delivers significant speedups for training large language models, such as for nanoGPT shown here.
Beta Features
[Beta] torch.compile
torch.compile is the main API for PyTorch 2.0, which wraps your model and returns a compiled model. It is a fully additive (and optional) feature and hence 2.0 is 100% backward compatible by definition.
Underpinning torch.compile are new technologies – TorchDynamo, AOTAutograd, PrimTorch and TorchInductor:


TorchDynamo captures PyTorch programs safely using Python Frame Evaluation Hooks and is a significant innovation that was a result of 5 years of our R&D into safe graph capture.
AOTAutograd overloads PyTorch’s autograd engine as a tracing autodiff for generating ahead-of-time backward traces.
PrimTorch canonicalizes ~2000+ PyTorch operators down to a closed set of ~250 primitive operators that developers can target to build a complete PyTorch backend. This substantially lowers the barrier of writing a PyTorch feature or backend.
TorchInductor is a deep learning compiler that generates fast code for multiple accelerators and backends. For NVIDIA and AMD GPUs, it uses OpenAI Triton as a key building block. For intel CPUs, we generate C++ code using multithreading, vectorized instructions and offloading appropriate operations to mkldnn when possible.


With all the new technologies, torch.compile is able to work 93% of time across 165 open-source models and runs 20% faster on average at float32 precision and 36% faster on average at AMP precision.
For more information, please refer to https://pytorch.org/get-started/pytorch-2.0/ and for TorchInductor CPU with Intel here.
[Beta] PyTorch MPS Backend
MPS backend provides GPU-accelerated PyTorch training on Mac platforms. This release brings improved correctness, stability, and operator coverage.


MPS backend now includes support for the Top 60 most used ops, along with the most frequently requested operations by the community, bringing coverage to over 300 operators. The major focus of the release was to enable full OpInfo-based forward and gradient mode testing to address silent correctness issues. These changes have resulted in wider adoption of MPS backend by 3rd party networks such as Stable Diffusion, YoloV5, WhisperAI, along with increased coverage for Torchbench networks and Basic tutorials. We encourage developers to update to the latest macOS release to see the best performance and stability on the MPS backend.
Links

MPS Backend
Developer information
Accelerated PyTorch training on Mac


Metal, Metal Performance Shaders & Metal Performance Shaders Graph

[Beta] Scaled dot product attention 2.0
We are thrilled to announce the release of PyTorch 2.0, which introduces a powerful scaled dot product attention function as part of torch.nn.functional. This function includes multiple implementations that can be seamlessly applied depending on the input and hardware in use.
In previous versions of PyTorch, you had to rely on third-party implementations and install separate packages to take advantage of memory-optimized algorithms like FlashAttention. With PyTorch 2.0, all these implementations are readily available by default.


These implementations include FlashAttention from HazyResearch, Memory-Efficient Attention from the xFormers project, and a native C++ implementation that is ideal for non-CUDA devices or when high-precision is required.
PyTorch 2.0 will automatically select the optimal implementation for your use case, but you can also toggle them individually for finer-grained control. Additionally, the scaled dot product attention function can be used to build common transformer architecture components.
Learn more with the documentation and this tutorial.
[Beta] functorch -> torch.func


[Beta] functorch -> torch.func
Inspired by Google JAX, functorch is a library that offers composable vmap (vectorization) and autodiff transforms. It enables advanced autodiff use cases that would otherwise be tricky to express in PyTorch. Examples include:
* model ensembling
* efficiently computing jacobians and hessians
* computing per-sample-gradients (or other per-sample quantities)


We’re excited to announce that, as the final step of upstreaming and integrating functorch into PyTorch, the functorch APIs are now available in the torch.func module. Our function transform APIs are identical to before, but we have changed how the interaction with NN modules work. Please see the docs and the migration guide for more details.
Furthermore, we have added support for torch.autograd.Function: one is now able to apply function transformations (e.g. vmap, grad, jvp) over torch.autograd.Function.
[Beta] Dispatchable Collectives


[Beta] Dispatchable Collectives
Dispatchable collectives is an improvement to the existing init_process_group() API which changes backend to an optional argument. For users, the main advantage of this feature is that it will allow them to write code that can run on both GPU and CPU machines without having to change the backend specification. The dispatchability feature will also make it easier for users to support both GPU and CPU collectives, as they will no longer need to specify the backend manually (e.g. “NCCL” or “GLOO”). Existing backend specifications by users will be honored and will not require change.
Usage example:
import torch.distributed.dist
…
# old
dist.init_process_group(backend=”nccl”, ...)
dist.all_reduce(...) # with CUDA tensors works
dist.all_reduce(...) # with CPU tensors does not work

# new
dist.init_process_group(...) # backend is optional
dist.all_reduce(...) # with CUDA tensors works
dist.all_reduce(...) # with CPU tensors works


dist.all_reduce(...) # with CPU tensors works
```
Learn more here.
[Beta] torch.set_default_device and torch.device as context manager
torch.set_default_device allows users to change the default device that factory functions in PyTorch allocate on. For example, if you torch.set_default_device(‘cuda’), a call to torch.empty(2) will allocate on CUDA (rather than on CPU). You can also use torch.device as a context manager to change the default device on a local basis. This resolves a long standing feature request from PyTorch’s initial release for a way to do this.
Learn more here.
[Beta] "X86" as the new default quantization backend for x86 CPU


The new X86 quantization backend, which utilizes FBGEMM and oneDNN kernel libraries, replaces FBGEMM as the default quantization backend for x86 CPU platforms and offers improved int8 inference performance compared to the original FBGEMM backend, leveraging the strengths of both libraries, with 1.3X – 2X inference performance speedup measured on 40+ deep learning models. The new backend is functionally compatible with the original FBGEMM backend.
Table: Geomean Speedup of X86 Quantization Backend vs. FBGEMM Backend

1 core/instance

2 cores/instance

4 cores/instance

1 socket (32 cores)/instance

Intel(R) Xeon(R) Platinum 8358 CPU @ 2.60GHz

1.76X

1.80X

2.04X

1.34X


1.34X

By default, users on x86 platforms will utilize the x86 quantization backend and their PyTorch programs will remain unchanged when using the default backend. Alternatively, users have the option to specify "X86" as the quantization backend explicitly. Example code is shown below:
import torch
from torch.ao.quantization import get_default_qconfig_mappingfrom torch.quantization.quantize_fx
import prepare_fx, convert_fx

# get default configuration
qconfig_mapping = get_default_qconfig_mapping()

# or explicitly specify the backend
# qengine = 'x86'
# torch.backends.quantized.engine = qengine
# qconfig_mapping = get_default_qconfig_mapping(qengine)

# construct fp32 model
model_fp32 = ...

# prepare
prepared_model = prepare_fx(model_fp32, qconfig_mapping, example_inputs=x)

# calibrate
...

# convert
quantized_model = convert_fx(prepared_model)


quantized_model = convert_fx(prepared_model)
```
Find more information: https://github.com/pytorch/pytorch/issues/83888 and https://www.intel.com/content/www/us/en/developer/articles/technical/accelerate-pytorch-int8-inf-with-new-x86-backend.html.
[Beta] GNN inference and training optimization on CPU
PyTorch 2.0 includes several critical optimizations to improve GNN inference and training performance on CPU. Before 2.0, GNN models of PyG suffers from low efficiency on CPU due to lack of performance tuning for several critical kernels (scatter/gather, etc) and the lack of GNN-related sparse matrix multiplication ops. To be specific, optimizations include:
* scatter_reduce: performance hotspot in Message Passing when the edge index is stored in Coordinate format (COO).


gather: backward of scatter_reduce, specially tuned for the GNN compute when the index is an expanded tensor.
torch.sparse.mm with reduce flag: performance hotspot in Message Passing when the edge index is stored in Compressed Sparse Row (CSR). Supported reduce flag of: sum, mean, amax, amin.

On PyG benchmarks/examples, OGB benchmarks, a 1.12x - 4.07x performance speedup is measured (1.13.1 compared with 2.0) for single node inference and training.

Model-Dataset

Option

Speedup Ratio

GCN-Reddit (inference)

512-2-64-dense

1.22x

1024-3-128-dense

1.25x

512-2-64-sparse

1.31x

1024-3-128-sparse

1.68x


1.68x

512-2-64-dense

1.22x

GraphSage-ogbn-products (inference)

1024-3-128-dense

1.15x

512-2-64-sparse

1.20x

1024-3-128-sparse

1.33x

full-batch-sparse

4.07x

GCN-PROTEINS (training)

3-32

1.67x

GCN-REDDIT-BINARY (training)

3-32

1.67x

GCN-Reddit (training)

512-2-64-dense

1.20x

1024-3-128-dense

1.12x

Learn more: PyG CPU Performance Optimization.


[Beta] Accelerating inference on CPU with PyTorch by leveraging oneDNN Graph
oneDNN Graph API extends oneDNN with a flexible graph API to maximize the optimization opportunity for generating efficient code on AI hardware.
* It automatically identifies the graph partitions to be accelerated via fusion.
* The fusion patterns focus on fusing compute-intensive operations such as convolution, matmul and their neighbor operations for both inference and training use cases.


Although work is ongoing to integrate oneDNN Graph with TorchDynamo as well, its integration with the PyTorch JIT Fuser attained beta status in PyTorch 2.0 for Float32 & BFloat16 inference (on machines that support AVX512_BF16 ISA).

From a developer’s/researcher’s perspective, the usage is quite simple & intuitive, with the only change in code being an API invocation:
* Leverage oneDNN Graph, with JIT-tracing, a model is profiled with an example input.
* The context manager with torch.jit.fuser(“fuser3”): can also be used instead of invoking torch.jit.enable_onednn_fusion(True).


For accelerating BFloat16 inference, we rely on eager-mode AMP (Automatic Mixed Precision) support in PyTorch & disable JIT mode’s AMP, as both of them are currently divergent:

# Assuming we have a model of the name 'model'

example_input = torch.rand(1, 3, 224, 224)

# enable oneDNN Graph
torch.jit.enable_onednn_fusion(True)
# Disable AMP for JIT
torch._C._jit_set_autocast_mode(False)
with torch.no_grad(), torch.cpu.amp.autocast():
    model = torch.jit.trace(model, (example_input))
    model = torch.jit.freeze(model)
    # 2 warm-ups (2 for tracing/scripting with an example, 3 without an example)
    model(example_input)
    model(example_input)

    # speedup would be observed in subsequent runs.
    model(example_input)

Learn more here.
Prototype Features
Distributed API


Prototype Features
Distributed API
[Prototype] DTensor


PyTorch DistributedTensor (DTensor) is a prototyping effort with distributed tensor primitives to allow easier distributed computation authoring in the SPMD (Single Program Multiple Devices) paradigm. The primitives are simple but powerful when used to express tensor distributions with both sharded and replicated parallelism strategies. PyTorch DTensor empowered PyTorch Tensor Parallelism along with other advanced parallelism explorations. In addition, it also offers a uniform way to save/load state_dict for distributed checkpointing purposes, even when there’re complex tensor distribution strategies such as combining tensor parallelism with parameter sharding in FSDP. More details can be found in this RFC and the DTensor examples notebook.


[Prototype] TensorParallel
We now support DTensor based Tensor Parallel which users can distribute their model parameters across different GPU devices. We also support Pairwise Parallel which shards two concatenated linear layers in a col-wise and row-wise style separately so that only one collective(all-reduce/reduce-scatter) is needed in the end. More details can be found in this example.
[Prototype] 2D Parallel
We implemented the integration of the aforementioned TP with FullyShardedDataParallel(FSDP) as 2D parallel to further scale large model training. More details can be found in this slide and code example.
[Prototype] torch.compile(dynamic=True)


[Prototype] torch.compile(dynamic=True)
Experimental support for PT2 compilation with dynamic shapes is available in this release. Inference compilation with inductor for simple models is supported, but there are a lot of limitations:

Training available in a future release (This is partially fixed in nightlies!)
Minifier available in a future release.
It is easy to end up in a situation where the dimension you wanted to be dynamic gets specialized anyway. Some of these issues are fixed in nightlies, others are not.
We do not appropriately propagate Inductor guards to the top-level, this is tracked at #96296.
Data-dependent operations like nonzero still require a graph break.
Dynamic does not work with non-standard modes like reduce-overhead or max-autotune.


There are many bugs in Inductor compilation. To track known bugs, check the dynamic shapes label on the PyTorch issue tracker.

For the latest and greatest news about dynamic shapes support on master, check out our status reports.
Highlights/Performance Improvements
Deprecation of Cuda 11.6 and Python 3.7 support for PyTorch 2.0
If you are still using or depending on CUDA 11.6 or Python 3.7 builds, we strongly recommend moving to at least CUDA 11.7 and Python 3.8, as it would be the minimum versions required for PyTorch 2.0. For more detail, please refer to the Release Compatibility Matrix for PyTorch releases.
Python 3.11 support on Anaconda Platform


Python 3.11 support on Anaconda Platform
Due to lack of Python 3.11 support for packages that PyTorch depends on, including NumPy, SciPy, SymPy, Pillow and others on the Anaconda platform. We will not be releasing Conda binaries compiled with Python 3.11 for PyTorch Release 2.0. The Pip packages with Python 3.11 support will be released, hence if you intend to use PyTorch 2.0 with Python 3.11 please use our Pip packages. Please note: Conda packages with Python 3.11 support will be made available on our nightly channel. Also we are planning on releasing Conda Python 3.11 binaries as part of future release once Anaconda provides these key dependencies. More information and instructions on how to download the Pip packages can be found here.
Optimized PyTorch Inference with AWS Graviton processors


The optimizations focused on three key areas: GEMM kernels, bfloat16 support, primitive caching and the memory allocator. For aarch64 platforms, PyTorch supports Arm Compute Library (ACL) GEMM kernels via Mkldnn(OneDNN) backend. The ACL library provides Neon/SVE GEMM kernels for fp32 and bfloat16 formats. The bfloat16 support on c7g allows efficient deployment of bfloat16 trained, AMP (Automatic Mixed Precision) trained, or even the standard fp32 trained models. The standard fp32 models leverage bfloat16 kernels via OneDNN fast math mode, without any model quantization. Next we implemented primitive caching for conv, matmul and inner product operators. More information on the updated PyTorch user guide with the upcoming 2.0 release improvements and TorchBench benchmark details can be found here.


layout: blog_detail
title: "Case Study: PathAI Uses PyTorch to Improve Patient Outcomes with AI-powered Pathology"
author: Logan Kilpatrick - Sr. Technology Advocate, Harshith Padigela - ML Engineer, Syed Ashar Javed - ML Technical Lead, Robert Egger - Biomedical Data Scientist
featured-img: "/assets/images/2022-7-15-PathAI-Uses-PyTorch-to-Improve-Patient-Outcomes-with-AI-powered-Pathology-1.png"

PathAI is the leading provider of AI-powered technology tools and services for pathology (the study of disease). Our platform was built to enable substantial improvements to the accuracy of diagnosis and the measurement of therapeutic efficacy for complex diseases, leveraging modern approaches in machine learning like image segmentation, graph neural networks, and multiple instance learning.


Traditional manual pathology is prone to subjectivity and observer variability that can negatively affect diagnoses and drug development trials. Before we dive into how we use PyTorch to improve our diagnosis workflow, let us first lay out the traditional analog Pathology workflow without machine learning.
How Traditional Biopharma Works
There are many avenues that biopharma companies take to discover novel therapeutics or diagnostics. One of those avenues relies heavily on the analysis of pathology slides to answer a variety of questions: how does a particular cellular communication pathway work? Can a specific disease state be linked to the presence or lack of a particular protein? Why did a particular drug in a clinical trial work for some patients but not others? Might there be an association between patient outcomes and a novel biomarker?


To help answer these questions, biopharma companies rely on expert pathologists to analyze slides and help evaluate the questions they might have.
As you might imagine, it takes an expert board certified pathologist to make accurate interpretations and diagnosis. In one study, a single biopsy result was given to 36 different pathologists and the outcome was 18 different diagnoses varying in severity from no treatment to aggressive treatment necessary. Pathologists also often solicit feedback from colleagues in difficult edge cases. Given the complexity of the problem, even with expert training and collaboration, pathologists can still have a hard time making a correct diagnosis. This potential variance can be the difference between a drug being approved and it failing the clinical trial.
How PathAI utilizes machine learning to power drug development


PathAI develops machine learning models which provide insights for drug development R&D, for powering clinical trials, and for making diagnoses. To this end, PathAI leverages PyTorch for slide level inference using a variety of methods including graph neural networks (GNN) as well as multiple instance learning. In this context, “slides” refers to full size scanned images of glass slides, which are pieces of glass with a thin slice of tissue between them, stained to show various cell formations. PyTorch enables our teams using these different methodologies to share a common framework which is robust enough to work in all the conditions we need. PyTorch’s high level, imperative, and pythonic syntax allows us to prototype models quickly and then take those models to scale once we have the results we want.
Multi-instance learning on gigabyte images


One of the uniquely challenging aspects of applying ML to pathology is the immense size of the images. These digital slides can often be 100,000 x 100,000 pixels or more in resolution and gigabytes in size. Loading the full image in GPU memory and applying traditional computer vision algorithms on them is an almost impossible task. It also takes both a considerable amount of time and resources to have a full slide image (100k x 100k) annotated, especially when annotators need to be domain experts (board-certified pathologists). We often build models to predict image-level labels, like the presence of cancer, on a patient slide which covers a few thousand pixels in the whole image. The cancerous area is sometimes a tiny fraction of the entire slide, which makes the ML problem similar to finding a needle in a haystack. On the other hand, some problems like the prediction of certain histological biomarkers require an aggregation of information from the whole slide which is again hard due to the size of the images. All these factors add significant algorithmic, computational, and logistical complexity when applying ML techniques to pathology problems.


Breaking down the image into smaller patches, learning patch representations, and then pooling those representations to predict an image-level label is one way to solve this problem as is depicted in the image below. One popular method for doing this is called Multiple Instance Learning (MIL). Each patch is considered an ‘instance’ and a set of patches forms a ‘bag’. The individual patch representations are pooled together to predict a final bag-level label. Algorithmically, the individual patch instances in the bag do not require labels and hence allow us to learn bag-level labels in a weakly-supervised way. They also use permutation invariant pooling functions which make the prediction independent of the order of patches and allows for an efficient aggregation of information. Typically, attention based pooling functions are used which not only allow for efficient aggregation but also provide attention values for each patch in the bag. These values indicate the importance of the corresponding patch in the prediction and can be visualized to better understand the model predictions. This element of interpretability can be very important to drive adoption of these models in the real world and we use variations like Additive MIL models to enable such spatial explainability. Computationally, MIL models circumvent the problem of applying neural networks to large image sizes since patch representations are obtained independently of the size of the image.


At PathAI, we use custom MIL models based on deep nets to predict image-level labels. The overview of this process is as follows:

Select patches from a slide using different sampling approaches.
Construct a bag of patches based on random sampling or heuristic rules.
Generate patch representations for each instance based on pre-trained models or large-scale representation learning models.
Apply permutation invariant pooling functions to get the final slide-level score.

Now that we have walked through some of the high-level details around MIL in PyTorch, let’s look at some code to see how simple it is to go from ideation to code in production with PyTorch. We begin by defining a sampler, transformations, and our MIL dataset:
```python
Create a bag sampler which randomly samples patches from a slide


bag_sampler = RandomBagSampler(bag_size=12)
Setup the transformations
crop_transform = FlipRotateCenterCrop(use_flips=True)
Create the dataset which loads patches for each bag
train_dataset = MILDataset(
  bag_sampler=bag_sampler,
  samples_loader=sample_loader,
  transform=crop_transform,
)

After we have defined our sampler and dataset, we need to define the model we will actually train with said dataset. PyTorch’s familiar model definition syntax makes this easy to do while also allowing us to create bespoke models at the same time.

```python
classifier = DefaultPooledClassifier(hidden_dims=[256, 256], input_dims=1024, output_dims=1)

pooling = DefaultAttentionModule(
  input_dims=1024,
  hidden_dims=[256, 256],
  output_activation=StableSoftmax()
)

# Define the model which is a composition of the featurizer, pooling module and a classifier
model = DefaultMILGraph(featurizer=ShuffleNetV2(), classifier=classifier, pooling = pooling)


```
Since these models are trained end-to-end, they offer a powerful way to go directly from a gigapixel whole slide image to a single label. Due to their wide applicability to different biological problems, two aspects of their implementation and deployment are important:

Configurable control over each part of the pipeline including the data loaders, the modular parts of the model, and their interaction with each other.
Ability to rapidly iterate through the ideate-implement-experiment-productionize loop.


PyTorch has various advantages when it comes to MIL modeling. It offers an intuitive way to create dynamic computational graphs with flexible control flow which is great for rapid research experimentation. The map-style datasets, configurable sampler and batch-samplers allow us to customize how we construct bags of patches, enabling faster experimentation. Since MIL models are IO heavy, data parallelism and pythonic data loaders make the task very efficient and user friendly. Lastly, the object-oriented nature of PyTorch enables building of reusable modules which aid in the rapid experimentation, maintainable implementation and ease of building compositional components of the pipeline.
Exploring spatial tissue organization with GNNs in PyTorch


In both healthy and diseased tissue, the spatial arrangement and structure of cells can oftentimes be as important as the cells themselves. For example, when assessing lung cancers, pathologists try to look at the overall grouping and structure of tumor cells (do they form solid sheets? Or do they occur in smaller, localized clusters?) to determine if the cancer belongs to specific subtypes which can have vastly different prognosis. Such spatial relationships between cells and other tissue structures can be modeled using graphs to capture tissue topology and cellular composition at the same time. Graph Neural Networks (GNNs) allow learning spatial patterns within these graphs that relate to other clinical variables, for example overexpression of genes in certain cancers.


In late 2020, when PathAI started using GNNs on tissue samples, PyTorch had the best and most mature support for GNN functionality via the PyG package. This made PyTorch the natural choice for our team given that GNN models were something that we knew would be an important ML concept we wanted to explore.


One of the main value-adds of GNN’s in the context of tissue samples is that the graph itself can uncover spatial relationships that would otherwise be very difficult to find by visual inspection alone. In our recent AACR publication, we showed that by using GNNs, we can better understand the way the presence of immune cell aggregates (specifically tertiary lymphoid structures, or TLS) in the tumor microenvironment can influence patient prognosis. In this case, the GNN approach was used to predict expression of genes associated with the presence of TLS, and identify histological features beyond the TLS region itself that are relevant to TLS. Such insights into gene expression are difficult to identify from tissue sample images when unassisted by ML models.


One of the most promising GNN variations we have had success with is self attention graph pooling. Let’s take a look at how we define our Self Attention Graph Pooling (SAGPool) model using PyTorch and PyG:
class SAGPool(torch.nn.Module):
  def __init__(self, ...):
    super().__init__()
    self.conv1 = GraphConv(in_features, hidden_features, aggr='mean')
    self.convs = torch.nn.ModuleList()
    self.pools = torch.nn.ModuleList()
    self.convs.extend([GraphConv(hidden_features, hidden_features, aggr='mean') for i in range(num_layers - 1)])
    self.pools.extend([SAGPooling(hidden_features, ratio, GNN=GraphConv, min_score=min_score) for i in range((num_layers) // 2)])
    self.jump = JumpingKnowledge(mode='cat')
    self.lin1 = Linear(num_layers * hidden_features, hidden_features)
    self.lin2 = Linear(hidden_features, out_features)
    self.out_activation = out_activation
    self.dropout = dropout


self.dropout = dropout
```
In the above code, we begin by defining a single convolutional graph layer and then add two module list layers which allow us to pass in a variable number of layers. We then take our empty module list and append a variable number of GraphConv layers followed by a variable number of SAGPooling layers. We finish up our SAGPool definition by adding a JumpingKnowledge Layer, two linear layers, our activation function, and our dropout value. PyTorch’s intuitive syntax allows us to abstract away the complexity of working with state of the art methods like SAG Poolings while also maintaining the common approach to model development we are familiar with.


Models like our SAG Pool one described above are just one example of how GNNs with PyTorch are allowing us to explore new and novel ideas. We also recently explored multimodal CNN - GNN hybrid models which ended up being 20% more accurate than traditional Pathologist consensus scores. These innovations and interplay between traditional CNNs and GNNs are again enabled by the short research to production model development loop.
Improving Patient Outcomes


Improving Patient Outcomes
In order to achieve our mission of improving patient outcomes with AI-powered pathology, PathAI needs to rely on an ML development framework that (1) facilitates quick iteration and easy extension (i.e. Model configuration as code) during initial phases of development and exploration (2) scales model training and inference to massive images (3) easily and robustly serves models for production uses of our products (in clinical trials and beyond). As we’ve demonstrated, PyTorch offers us all of these capabilities and more. We are incredibly excited about the future of PyTorch and cannot wait to see what other impactful challenges we can solve using the framework.


Torch Distributed Elastic
Makes distributed PyTorch fault-tolerant and elastic.
Get Started
Usage
^^^^^

Quickstart

Train script

Examples

Documentation
API
^^^

torchrun (Elastic Launch)

Elastic Agent

Multiprocessing

Error Propagation

Rendezvous

Expiration Timers

Metrics

Events

Advanced
^^^^^^^^

Customization

Plugins
^^^^^^^

TorchElastic Kubernetes


torch.overrides
This module exposes various helper functions for the
"torch_function" protocol. See Extending torch for more detail on
the "torch_function" protocol.
Functions
torch.overrides.get_ignored_functions()
Return public functions that cannot be overridden by
   "torch_function".
Returns:
      A tuple of functions that are publicly available in the torch
      API but cannot be overridden with "torch_function". Mostly
      this is because none of the arguments of these functions are
      tensors or tensor-likes.
Return type:
      Set[Callable]
-[ Examples ]-

torch.Tensor.as_subclass in torch.overrides.get_ignored_functions()
   True
torch.add in torch.overrides.get_ignored_functions()
   False

torch.overrides.get_overridable_functions()
List functions that are overridable via torch_function
Returns:
      A dictionary that maps namespaces that contain overridable


functions to functions in that namespace that can be overridden.
Return type:
      Dict[Any, List[Callable]]
torch.overrides.resolve_name(f)
Get a human readable string name for a function passed to
   torch_function
Parameters:
      f (Callable) -- Function to resolve the name of.
Returns:
      Name of the function; if eval'ed it should give back the input
      function.
Return type:
      str
torch.overrides.get_testing_overrides()
Return a dict containing dummy overrides for all overridable
   functions
Returns:
      A dictionary that maps overridable functions in the PyTorch API
      to lambda functions that have the same signature as the real
      function and unconditionally return -1. These lambda functions
      are useful for testing API coverage for a type that defines
      "torch_function".
Return type:
      Dict[Callable, Callable]
-[ Examples ]-

import inspect


-[ Examples ]-

import inspect
my_add = torch.overrides.get_testing_overrides()[torch.add]
inspect.signature(my_add)

torch.overrides.handle_torch_function(public_api, relevant_args, args, *kwargs)
Implement a function with checks for "torch_function"
   overrides.
See torch::autograd::handle_torch_function for the equivalent of
   this function in the C++ implementation.
Parameters:
      * public_api (function) -- Function exposed by the public
        torch API originally called like "public_api(args, *kwargs)"
        on which arguments are now being checked.
  * **relevant_args** (*iterable*) -- Iterable of arguments to
    check for __torch_function__ methods.

  * **args** (*tuple*) -- Arbitrary positional arguments
    originally passed into "public_api".

  * **kwargs** (*tuple*) -- Arbitrary keyword arguments originally
    passed into "public_api".

Returns:


passed into "public_api".
Returns:
      Result from calling "implementation" or an "torch_function"
      method, as appropriate.
Return type:
      object
:raises TypeError : if no implementation is found.:
-[ Example ]-

def func(a):
   ...     if has_torch_function_unary(a):
   ...         return handle_torch_function(func, (a,), a)
   ...     return a + 0

torch.overrides.has_torch_function()
Check for torch_function implementations in the elements of an
   iterable or if a torch_function mode is enabled.  Considers
   exact "Tensor" s and "Parameter" s non-dispatchable.  Use this to
   guard a call to "handle_torch_function()"; don't use it to test if
   something is Tensor-like, use "is_tensor_like()" instead. :param
   relevant_args: Iterable or arguments to check for
   torch_function methods. :type relevant_args: iterable
Returns:
      True if any of the elements of relevant_args have
      torch_function implementations, False otherwise.


Return type:
      bool
See also:
 "torch.is_tensor_like"
    Checks if something is a Tensor-like, including an exact
    "Tensor".

torch.overrides.is_tensor_like(inp)
Returns "True" if the passed-in input is a Tensor-like.
Currently, this occurs whenever there's a "torch_function"
   attribute on the type of the input.
-[ Examples ]-
A subclass of tensor is generally a Tensor-like.

class SubTensor(torch.Tensor): ...
is_tensor_like(SubTensor([0]))
   True

Built-in or user types aren't usually Tensor-like.

is_tensor_like(6)
   False
is_tensor_like(None)
   False
class NotATensor: ...
is_tensor_like(NotATensor())
   False

But, they can be made Tensor-like by implementing
   torch_function.

class TensorLike:
   ...     @classmethod
   ...     def torch_function(cls, func, types, args, kwargs):
   ...         return -1
is_tensor_like(TensorLike())
   True


is_tensor_like(TensorLike())
   True

torch.overrides.is_tensor_method_or_property(func)
Returns True if the function passed in is a handler for a method or
   property belonging to "torch.Tensor", as passed into
   "torch_function".
Note:
 For properties, their "__get__" method must be passed in.

This may be needed, in particular, for the following reasons:

Methods/properties sometimes don't contain a module slot.

They require that the first passed-in argument is an instance of
      "torch.Tensor".

-[ Examples ]-

is_tensor_method_or_property(torch.Tensor.add)
   True
is_tensor_method_or_property(torch.add)
   False

Return type:
      bool
torch.overrides.wrap_torch_function(dispatcher)
Wraps a given function with "torch_function" -related
   functionality.
Parameters:
      dispatcher (Callable) -- A callable that returns an
      iterable of Tensor-likes passed into the function.
Note:


Note:
 This decorator may reduce the performance of your code.
 Generally, it's enough to express your code as a series of
 functions that, themselves, support __torch_function__. If you
 find yourself in the rare situation where this is not the case,
 e.g. if you're wrapping a low-level library and you also need it
 to work for Tensor-likes, then this function is available.

-[ Examples ]-

def dispatcher(a): # Must have the same signature as func
   ...     return (a,)
@torch.overrides.wrap_torch_function(dispatcher)
def func(a): # This will make func dispatchable by torch_function
   ...     return a + 0


Quantization Accuracy Debugging
This document provides high level strategies for improving
quantization accuracy. If a quantized model has error compared to the
original model, we can categorize the error into:

data insensitive error - caused by intrinsic model quantization
   error, large portion of input data has large error

data sensitive error - caused by outlier input data, small
   portion of input data has large error

implementation error - quantized kernel is not matching
   reference implementation

Data insensitive error
General tips

For PTQ, ensure that the data you are calibrating with is
   representative of your dataset. For example, for a classification
   problem a general guideline is to have multiple samples in every
   category, and the overall number of samples should be at least 100.
   There is no penalty for calibrating with more data other than
   calibration time.


calibration time.

If your model has Conv-BN or Linear-BN patterns, consider fusing
   them. If you are using FX graph mode quantization, this is done
   automatically by the workflow. If you are using Eager mode
   quantization, you can do this manually with the
   "torch.ao.quantization.fuse_modules" API.

Increase the precision of dtype of the problematic ops. Usually,
   fp32 will have the highest accuracy, followed by fp16, followed by
   dynamically quantized int8, followed by statically quantized int8.

Note: this is trading off performance for accuracy.

Note: availability of kernels per dtype per op can vary by
      backend.

Note: dtype conversions add an additional performance cost. For
      example, "fp32_op -> quant -> int8_op -> dequant -> fp32_op ->
      quant -> int8_op -> dequant" will have a performance penalty
      compared to "fp32_op -> fp32_op -> quant -> int8_op -> int8_op
      -> dequant" because of a higher number of required dtype
      conversions.


conversions.

If you are using PTQ, consider using QAT to recover some of the
   accuracy loss from quantization.

Int8 quantization tips

If you are using per-tensor weight quantization, consider using
   per-channel weight quantization.

If you are doing inference on fbgemm, ensure that you set the
   reduce_range argument to False if your CPU is Cooperlake or
   newer, and to True otherwise.

Audit the input activation distribution variation across different
   samples. If this variation is high, the layer may be suitable for
   dynamic quantization but not static quantization.

Data sensitive error
If you are using static quantization and a small portion of your input
data is resulting in high quantization error, you can try:

Adjust your calibration dataset to make it more representative of
   your inference dataset.

Manually inspect (using Numeric Suite) which layers have high


quantization error. For these layers, consider leaving them in
   floating point or adjusting the observer settings to choose a
   better scale and zero_point.
Implementation error
If you are using PyTorch quantization with your own backend you may
see differences between the reference implementation of an operation
(such as "dequant -> op_fp32 -> quant") and the quantized
implementation (such as op_int8) of the op on the target hardware.
This could mean one of two things:

the differences (usually small) are expected due to specific
   behavior of the target kernel on the target hardware compared to
   fp32/cpu. An example of this is accumulating in an integer dtype.
   Unless the kernel guarantees bitwise equivalency with the reference
   implementation, this is expected.

the kernel on the target hardware has an accuracy issue. In this
   case, reach out to the kernel developer.

Numerical Debugging Tooling (prototype)
Warning:


Warning:
Numerical debugging tooling is early prototype and subject to
  change.

torch.ao.ns._numeric_suite Eager mode numeric suite

torch.ao.ns._numeric_suite_fx FX numeric suite


JIT Utils - torch.utils.jit


Distributed Optimizers
Warning:
Distributed optimizer is not currently supported when using CUDA
  tensors
"torch.distributed.optim" exposes DistributedOptimizer, which takes a
list of remote parameters ("RRef") and runs the optimizer locally on
the workers where the parameters live.  The distributed optimizer can
use any of the local optimizer Base class to apply the gradients on
each worker.
class torch.distributed.optim.DistributedOptimizer(optimizer_class, params_rref, args, *kwargs)
DistributedOptimizer takes remote references to parameters
   scattered across workers and applies the given optimizer locally
   for each parameter.
This class uses "get_gradients()" in order to retrieve the
   gradients for specific parameters.
Concurrent calls to "step()", either from the same or different
   clients, will be serialized on each worker -- as each worker's
   optimizer can only work on one set of gradients at a time. However,


there is no guarantee that the full forward-backward-optimizer
   sequence will execute for one client at a time. This means that the
   gradients being applied may not correspond to the latest forward
   pass executed on a given worker. Also, there is no guaranteed
   ordering across workers.
DistributedOptimizer creates the local optimizer with TorchScript
   enabled by default, so that optimizer updates are not blocked by
   the Python Global Interpreter Lock (GIL) in the case of
   multithreaded training (e.g. Distributed Model Parallel). This
   feature is currently enabled for most optimizers. You can also
   follow the recipe in PyTorch tutorials to enable TorchScript
   support for your own custom optimizers.
Parameters:
      * optimizer_class (optim.Optimizer) -- the class of
        optimizer to instantiate on each worker.
  * **params_rref** (*list**[**RRef**]*) -- list of RRefs to local
    or remote parameters to optimize.


or remote parameters to optimize.
  * **args** -- arguments to pass to the optimizer constructor on
    each worker.

  * **kwargs** -- arguments to pass to the optimizer constructor
    on each worker.

Example::
      >>> import torch.distributed.autograd as dist_autograd
      >>> import torch.distributed.rpc as rpc
      >>> from torch import optim
      >>> from torch.distributed.optim import DistributedOptimizer
      >>>
      >>> with dist_autograd.context() as context_id:
      >>>   # Forward pass.
      >>>   rref1 = rpc.remote("worker1", torch.add, args=(torch.ones(2), 3))
      >>>   rref2 = rpc.remote("worker1", torch.add, args=(torch.ones(2), 1))
      >>>   loss = rref1.to_here() + rref2.to_here()
      >>>
      >>>   # Backward pass.
      >>>   dist_autograd.backward(context_id, [loss.sum()])
      >>>
      >>>   # Optimizer.
      >>>   dist_optim = DistributedOptimizer(
      >>>      optim.SGD,
      >>>      [rref1, rref2],
      >>>      lr=0.05,


 lr=0.05,
  >>>   )
  >>>   dist_optim.step(context_id)

step(context_id)
  Performs a single optimization step.

  This will call "torch.optim.Optimizer.step()" on each worker
  containing parameters to be optimized, and will block until all
  workers return. The provided "context_id" will be used to
  retrieve the corresponding "context" that contains the gradients
  that should be applied to the parameters.

  Parameters:
     **context_id** -- the autograd context id for which we should
     run the optimizer step.

class torch.distributed.optim.PostLocalSGDOptimizer(optim, averager)
Wraps an arbitrary "torch.optim.Optimizer" and runs post-local SGD,
   This optimizer runs local optimizer at every step. After the warm-
   up stage, it averages parameters periodically afer the local
   optimizer is applied.
Parameters:
      * optim (Optimizer) -- The local optimizer.


averager (ModelAverager) -- A model averager instance to
        run post-localSGD algorithm.

Example:
  >>> import torch
  >>> import torch.distributed as dist
  >>> import torch.distributed.algorithms.model_averaging.averagers as averagers
  >>> import torch.nn as nn
  >>> from torch.distributed.optim import PostLocalSGDOptimizer
  >>> from torch.distributed.algorithms.ddp_comm_hooks.post_localSGD_hook import (
  >>>   PostLocalSGDState,
  >>>   post_localSGD_hook,
  >>> )
  >>>
  >>> model = nn.parallel.DistributedDataParallel(
  >>>    module, device_ids=[rank], output_device=rank
  >>> )
  >>>
  >>> # Register a post-localSGD communication hook.
  >>> state = PostLocalSGDState(process_group=None, subgroup=None, start_localSGD_iter=100)
  >>> model.register_comm_hook(state, post_localSGD_hook)
  >>>
  >>> # Create a post-localSGD optimizer that wraps a local optimizer.


Note that warmup_steps used in PostLocalSGDOptimizer must be the same as
  >>> # ``start_localSGD_iter`` used in ``PostLocalSGDState``.
  >>> local_optim = torch.optim.SGD(params=model.parameters(), lr=0.01)
  >>> opt = PostLocalSGDOptimizer(
  >>>     optim=local_optim,
  >>>     averager=averagers.PeriodicModelAverager(period=4, warmup_steps=100)
  >>> )
  >>>
  >>> # In the first 100 steps, DDP runs global gradient averaging at every step.
  >>> # After 100 steps, DDP runs gradient averaging within each subgroup (intra-node by default),
  >>> # and post-localSGD optimizer runs global model averaging every 4 steps after applying the local optimizer.
  >>> for step in range(0, 200):
  >>>    opt.zero_grad()
  >>>    loss = loss_fn(output, labels)
  >>>    loss.backward()
  >>>    opt.step()

load_state_dict(state_dict)
  This is the same as "torch.optim.Optimizer" "load_state_dict()",


but also restores model averager's step value to the one saved
      in the provided "state_dict".
  If there is no ""step"" entry in "state_dict", it will raise a
  warning and initialize the model averager's step to 0.

state_dict()
  This is the same as "torch.optim.Optimizer" "state_dict()", but
  adds an extra entry to record model averager's step to the
  checkpoint to ensure reload does not cause unnecessary warm up
  again.

step()
  Performs a single optimization step (parameter update).

class torch.distributed.optim.ZeroRedundancyOptimizer(params, optimizer_class, process_group=None, parameters_as_bucket_view=False, overlap_with_ddp=False, **defaults)
This class wraps an arbitrary "optim.Optimizer" and shards its
   states across ranks in the group as described by ZeRO. The local
   optimizer instance in each rank is only responsible for updating
   approximately "1 / world_size" parameters and hence only needs to


keep "1 / world_size" optimizer states. After parameters are
   updated locally, each rank will broadcast its parameters to all
   other peers to keep all model replicas in the same state.
   "ZeroRedundancyOptimizer" can be used in conjunction with
   "torch.nn.parallel.DistributedDataParallel" to reduce per-rank peak
   memory consumption.
"ZeroRedundancyOptimizer" uses a sorted-greedy algorithm to pack a
   number of parameters at each rank. Each parameter belongs to a
   single rank and is not divided among ranks. The partition is
   arbitrary and might not match the the parameter registration or
   usage order.
Parameters:
      params ("Iterable") -- an "Iterable" of "torch.Tensor" s or
      "dict" s giving all parameters, which will be sharded across
      ranks.
Keyword Arguments:
      * optimizer_class ("torch.nn.Optimizer") -- the class of the
        local optimizer.
  * **process_group** ("ProcessGroup", optional) --


"torch.distributed" "ProcessGroup" (default:
        "dist.group.WORLD" initialized by
        "torch.distributed.init_process_group()").
  * **parameters_as_bucket_view** (*bool**, **optional*) -- if
    "True", parameters are packed into buckets to speed up
    communication, and "param.data" fields point to bucket views
    at different offsets; if "False", each individual parameter is
    communicated separately, and each "params.data" stays intact
    (default: "False").

  * **overlap_with_ddp** (*bool**, **optional*) -- if "True",
    "step()" is overlapped with "DistributedDataParallel" 's
    gradient synchronization; this requires (1) either a
    functional optimizer for the "optimizer_class" argument or one
    with a functional equivalent and (2) registering a DDP
    communication hook constructed from one of the functions in
    "ddp_zero_hook.py"; parameters are packed into buckets


matching those in "DistributedDataParallel", meaning that the
        "parameters_as_bucket_view" argument is ignored. If "False",
        "step()" runs disjointly after the backward pass (per normal).
        (default: "False")
  * ****defaults** -- any trailing arguments, which are forwarded
    to the local optimizer.

Example:
  >>> import torch.nn as nn
  >>> from torch.distributed.optim import ZeroRedundancyOptimizer
  >>> from torch.nn.parallel import DistributedDataParallel as DDP
  >>> model = nn.Sequential(*[nn.Linear(2000, 2000).to(rank) for _ in range(20)])
  >>> ddp = DDP(model, device_ids=[rank])
  >>> opt = ZeroRedundancyOptimizer(
  >>>     ddp.parameters(),
  >>>     optimizer_class=torch.optim.Adam,
  >>>     lr=0.01
  >>> )
  >>> ddp(inputs).sum().backward()
  >>> opt.step()

Warning:
 Currently, "ZeroRedundancyOptimizer" requires that all of the
 passed-in parameters are the same dense type.

Warning:


Warning:
 If you pass "overlap_with_ddp=True", be wary of the following:
 Given the way that overlapping "DistributedDataParallel" with
 "ZeroRedundancyOptimizer" is currently implemented, the first two
 or three training iterations do not perform parameter updates in
 the optimizer step, depending on if "static_graph=False" or
 "static_graph=True", respectively. This is because it needs
 information about the gradient bucketing strategy used by
 "DistributedDataParallel", which is not finalized until the
 second forward pass if "static_graph=False" or until the third
 forward pass if "static_graph=True". To adjust for this, one
 option is to prepend dummy inputs.

Warning:
 ZeroRedundancyOptimizer is experimental and subject to change.

add_param_group(param_group)
  Add a parameter group to the "Optimizer" 's "param_groups".

  This can be useful when fine tuning a pre-trained network, as


frozen layers can be made trainable and added to the "Optimizer"
      as training progresses.
  Parameters:
     **param_group** (*dict*) -- specifies the parameters to be
     optimized and group-specific optimization options.

  Warning:

    This method handles updating the shards on all partitions but
    needs to be called on all ranks. Calling this on a subset of
    the ranks will cause the training to hang because
    communication primitives are called depending on the managed
    parameters and expect all the ranks to participate on the same
    set of parameters.

consolidate_state_dict(to=0)
  Consolidate a list of "state_dict" s (one per rank) on the
  target rank.

  Parameters:
     **to** (*int*) -- the rank that receives the optimizer states
     (default: 0).

  Raises:
     **RuntimeError** -- if "overlap_with_ddp=True" and this
     method is     called before this "ZeroRedundancyOptimizer"


instance     has been fully initialized, which happens once
         "DistributedDataParallel" gradient buckets have been
         rebuilt.
  Warning:

    This needs to be called on all ranks.

join_hook(**kwargs)
  Returns the ZeRO join hook, which enables training on uneven
  inputs by shadowing the collective communications in the
  optimizer step.

  Gradients must be properly set before this hook is called.

  Parameters:
     **kwargs** (*dict*) -- a "dict" containing any keyword
     arguments to modify the behavior of the join hook at run
     time; all "Joinable" instances sharing the same join context
     manager are forwarded the same value for "kwargs".

  This hook does not support any keyword arguments; i.e. "kwargs"
  is unused.

load_state_dict(state_dict)
  Load the state pertaining to the given rank from the input
  "state_dict", updating the local optimizer as needed.

  Parameters:


Parameters:
         state_dict (dict) -- optimizer state; should be an
         object returned from a call to "state_dict()".
  Raises:
     **RuntimeError** -- if "overlap_with_ddp=True" and this
     method is     called before this "ZeroRedundancyOptimizer"
     instance     has been fully initialized, which happens once
     "DistributedDataParallel" gradient buckets have been
     rebuilt.

state_dict()
  Returns the last global optimizer state known to this rank.

  Raises:
     **RuntimeError** -- if "overlap_with_ddp=True" and this
     method is     called before this "ZeroRedundancyOptimizer"
     instance     has been fully initialized, which happens once
     "DistributedDataParallel" gradient buckets have been
     rebuilt; or if this method is called without a preceding call
     to "consolidate_state_dict()".

  Return type:
     *Dict*[str, *Any*]

step(closure=None, **kwargs)


step(closure=None, **kwargs)
  Performs a single optimizer step and syncs parameters across all
  ranks.

  Parameters:
     **closure** (*Callable*) -- a closure that re-evaluates the
     model and returns the loss; optional for most optimizers.

  Returns:
     Optional loss depending on the underlying local optimizer.

  Return type:
     *Optional*[float]


Distributed Autograd Design
This note will present the detailed design for distributed autograd
and walk through the internals of the same. Make sure you're familiar
with Autograd mechanics and the Distributed RPC Framework before
proceeding.
Background
Let's say you have two nodes and a very simple model partitioned
across two nodes. This can be implemented using
"torch.distributed.rpc" as follows:
import torch
   import torch.distributed.rpc as rpc
def my_add(t1, t2):
     return torch.add(t1, t2)
# On worker 0:
   t1 = torch.rand((3, 3), requires_grad=True)
   t2 = torch.rand((3, 3), requires_grad=True)
# Perform some computation remotely.
   t3 = rpc.rpc_sync("worker1", my_add, args=(t1, t2))
# Perform some computation locally based on remote result.
   t4 = torch.rand((3, 3), requires_grad=True)
   t5 = torch.mul(t3, t4)
# Compute some loss.
   loss = t5.sum()
The main motivation behind distributed autograd is to enable running a


backward pass on such distributed models with the "loss" that we've
computed and record appropriate gradients for all tensors that require
gradients.
Autograd recording during the forward pass
PyTorch builds the autograd graph during the forward pass and this
graph is used to execute the backward pass. For more details see How
autograd encodes the history.
For distributed autograd, we need to keep track of all RPCs during the
forward pass to ensure the backward pass is executed appropriately.
For this purpose, we attach "send" and "recv" functions to the
autograd graph when we perform an RPC.

The "send" function is attached to the source of the RPC and its
  output edges point to the autograd function for the input tensors of
  the RPC. The input for this function during the backward pass is
  received from the destination as the output of the appropriate
  "recv" function.

The "recv" function is attached to the destination of the RPC and


its inputs are retrieved from operators executed on the destination
  using the input tensors. The output gradients of this function are
  sent to the source node to the appropriate "send" function during
  the backward pass.

Each "send-recv" pair is assigned a globally unique
  "autograd_message_id" to uniquely identify the pair. This is useful
  to look up the corresponding function on a remote node during the
  backward pass.

For RRef, whenever we call "torch.distributed.rpc.RRef.to_here()" we
  attach an appropriate "send-recv" pair for the tensors involved.

As an example, this is what the autograd graph for our example above
would look like (t5.sum() excluded for simplicity):
[image]
Distributed Autograd Context
Each forward and backward pass that uses distributed autograd is
assigned a unique "torch.distributed.autograd.context" and this
context has a globally unique "autograd_context_id". This context is
created on each node as needed.


created on each node as needed.
This context serves the following purpose:

Multiple nodes running distributed backward passes might accumulate
   gradients on the same tensor and as a result the ".grad" field of
   the tensor would have gradients from a variety of distributed
   backward passes before we have the opportunity to run the
   optimizer. This is similar to calling "torch.autograd.backward()"
   multiple times locally. In order to provide a way of separating out
   the gradients for each backward pass, the gradients are accumulated
   in the "torch.distributed.autograd.context" for each backward pass.

During the forward pass we store the "send" and "recv" functions
   for each autograd pass in this context. This ensures we hold
   references to the appropriate nodes in the autograd graph to keep
   it alive. In addition to this, it is easy to look up the
   appropriate "send" and "recv" functions during the backward pass.


In general we also use this context to store some metadata for each
   distributed autograd pass.

From the user's perspective the autograd context is setup as follows:
import torch.distributed.autograd as dist_autograd
   with dist_autograd.context() as context_id:
     loss = model.forward()
     dist_autograd.backward(context_id, loss)
It is important to note that your model's forward pass must be invoked
within the distributed autograd context manager, as a valid context is
needed in order to ensure that all "send" and "recv" functions are
stored properly to run the backward pass across all participating
nodes.
Distributed Backward Pass
In this section we outline the challenge of computing dependencies
accurately during a distributed backward pass and describe a couple of
algorithms (with tradeoffs) on how we can execute a distributed
backward pass.
Computing dependencies
Consider the following piece of code being run on a single machine


import torch
   a = torch.rand((3, 3), requires_grad=True)
   b = torch.rand((3, 3), requires_grad=True)
   c = torch.rand((3, 3), requires_grad=True)
   d = a + b
   e = b * c
   d.sum.().backward()
This is what the autograd graph for the code above would look like:
[image]
The first step the autograd engine performs as part of the backward
pass is computing the number of dependencies for each node in the
autograd graph. This helps the autograd engine know when a node in the
graph is ready for execution. The numbers in brackets for "add(1)" and
"mul(0)" denote the number of dependencies. As you can see, this means
during the backward pass the "add" node needs 1 input and the "mul"
node doesn't need any inputs (in other words doesn't need to be
executed). The local autograd engine computes these dependencies by
traversing the graph from the root nodes ("d" in this case).
The fact that certain nodes in the autograd graph might not be
executed in the backward pass poses a challenge for distributed


autograd. Consider this piece of code which uses RPC.
import torch
   import torch.distributed.rpc as rpc
a = torch.rand((3, 3), requires_grad=True)
   b = torch.rand((3, 3), requires_grad=True)
   c = torch.rand((3, 3), requires_grad=True)
d = rpc.rpc_sync("worker1", torch.add, args=(a, b))
   e = rpc.rpc_sync("worker1", torch.mul, args=(b, c))
   loss = d.sum()
The associated autograd graph for the code above would be:
[image]
Computing dependencies of this distributed autograd graph is much more
challenging and requires some overhead (either in terms of computation
or network communication).
For performance sensitive applications we can avoid a lot of overhead
by assuming every "send" and "recv" function are valid as part of the
backward pass (most applications don't perform RPCs that aren't used).
This simplifies the distributed autograd algorithm and is much more
efficient, but at the cost that the application needs to be aware of


the limitations. This algorithm is called the FAST mode algorithm and
is described in detail below.
In the general case it might not be necessary that every "send" and
"recv" function is valid as part of the backward pass. To address
this, we have proposed a SMART mode algorithm which is described in a
later section. Please note that currently, only the FAST mode
algorithm is implemented.
FAST mode algorithm
The key assumption of this algorithm is that each "send" function has
a dependency of 1 when we run a backward pass. In other words, we
assume we'll receive a gradient over RPC from another node.
The algorithm is as follows:

We start from the worker which has the roots for the backward pass
   (all roots must be local).

Lookup all the "send" functions for the current Distributed
   Autograd Context.

Compute dependencies locally starting from the provided roots and
   all the "send" functions we retrieved.


all the "send" functions we retrieved.

After computing dependencies, kick off the local autograd engine
   with the provided roots.

When the autograd engine executes the "recv" function, the "recv"
   function sends the input gradients via RPC to the appropriate
   worker. Each "recv" function knows the destination worker id since
   it is recorded as part of the forward pass. The "recv" function
   also sends over the "autograd_context_id" and "autograd_message_id"
   to the remote host.

When this request is received on the remote host, we use the
   "autograd_context_id" and "autograd_message_id" to look up the
   appropriate "send" function.

If this is the first time a worker has received a request for the
   given "autograd_context_id", it will compute dependencies locally
   as described in points 1-3 above.

The "send" function retrieved in 6. is then enqueued for execution
   on the local autograd engine for that worker.


on the local autograd engine for that worker.

Finally, instead of accumulating the gradients on the ".grad" field
   of the Tensor, we accumulate the gradients separately per
   Distributed Autograd Context. The gradients are stored in a
   "Dict[Tensor, Tensor]", which is basically a map from Tensor to its
   associated gradient and this map can be retrieved using the
   "get_gradients()" API.

As an example the complete code with distributed autograd would be as
follows:
import torch
   import torch.distributed.autograd as dist_autograd
   import torch.distributed.rpc as rpc
def my_add(t1, t2):
     return torch.add(t1, t2)
# On worker 0:
# Setup the autograd context. Computations that take
   # part in the distributed backward pass must be within
   # the distributed autograd context manager.
   with dist_autograd.context() as context_id:
     t1 = torch.rand((3, 3), requires_grad=True)
     t2 = torch.rand((3, 3), requires_grad=True)
 # Perform some computation remotely.