A library for developing and deploying PyTorch models using ONNX Runtime.

Installation • Training • Inference • Docs • License

Introduction

A library for accelerating PyTorch models using ONNX Runtime:

torch-ort to train PyTorch models faster with ONNX Runtime
moe to scale large models and improve their quality
torch-ort-infer to perform inference on PyTorch models with ONNX Runtime and Intel® OpenVINO™

🚀 Installation

Install for training

Pre-requisites

You need a machine with at least one NVIDIA or AMD GPU to run ONNX Runtime for PyTorch.

You can install and run torch-ort in your local environment, or with Docker.

Install in a local Python environment

Install CUDA
Install CuDNN
Install torch-ort
- pip install torch-ort
Run post-installation script for ORTModule
- python -m torch_ort.configure

Get install instructions for other combinations in the Get Started Easily section at https://www.onnxruntime.ai/ under the Optimize Training tab.

Verify your installation

Clone this repo
- git clone git@github.com:pytorch/ort.git
Install extra dependencies
- pip install wget pandas sklearn transformers
Run a test training script
- python ./ort/tests/bert_for_sequence_classification.py

Install Mixture Of Experts

Mixture of Experts layer implementation is available in the ort_moe folder.

Clone this repo

git clone https://github.com/pytorch/ort.git

Build MoE

cd ort_moe
python setup.py install

Install for Inference

Prerequisites

Ubuntu 18.04, 20.04
Python* 3.7, 3.8 or 3.9

Install in a local Python environment

pip install torch-ort-infer[openvino]
Run post installation configuration script python -m torch_ort.configure

Verify your installation

Clone this repo
- git clone git@github.com:pytorch/ort.git
Install extra dependencies
- pip install wget pandas sklearn transformers
Run a test script
- python ./torch_ort_inference/tests/bert_for_sequence_classification.py

📈 Training

The torch-ort library accelerates training of large transformer PyTorch models to reduce the training time and GPU cost with a few lines of code change. It is built on top of highly successful and proven technologies of ONNX Runtime and ONNX format and includes the ONNX Runtime Optimizer and Data Sampler.

Add ONNX Runtime for PyTorch to your PyTorch training script

from torch_ort import ORTModule
model = ORTModule(model)
# PyTorch training script follows

Usage of FusedAdam and FP16 Optimizer (Optional)

import torch
from torch_ort.optim import FusedAdam
class NeuralNet(torch.nn.Module):
    ...
# Only supports GPU Currently.
device = "cuda"
model = NeuralNet(...).to(device)
ort_fused_adam_optimizer = FusedAdam(
    model.parameters(), lr=1e-3, betas=(0.9, 0.999), weight_decay=0.01, eps=1e-8
)

# To use FP16_Optimizer, Add these lines : 
from torch_ort.optim import FP16_Optimizer
ort_fused_adam_optimizer = FP16_Optimizer(ort_fused_adam_optimizer)


loss = model(...).sum()
loss.backward()
ort_fused_adam_optimizer.step()
ort_fused_adam_optimizer.zero_grad()

For detailed documentation see FusedAdam

For a full working example see FusedAdam Test Example

FP16_Optimizer is a simple wrapper to replace inefficient FP16_Optimizer function calls implemented by libraries for example Apex, DeepSpeed, Megatron-LM.

For detailed documentation see FP16 Optimizer

Usage of LoadBalancingDistributedSampler

import torch
from torch.utils.data import DataLoader 
from torch_ort.utils.data import LoadBalancingDistributedSampler
class MyDataset(torch.utils.data.Dataset):
   ...
   
def collate_fn(data): 
    ...
    return samples, label_list 
samples = [...] 
labels = [...] 
dataset = MyDataset(samples, labels) 
data_sampler = sampler.LoadBalancingDistributedSampler( 
    dataset, complexity_fn=complexity_fn, world_size=2, rank=0, shuffle=False 
) 
train_dataloader = DataLoader(dataset, batch_size=2, sampler=data_sampler, collate_fn=collate_fn) 
for batched_data, batched_label in train_dataloader: 
    optimizer.zero_grad() 
    loss = loss_fn(model(batched_data) , batched_labels) 
    loss.backward() 
    optimizer.step()

For detailed documentation see LoadBalancingDistributedSampler

For a full working example see LoadBalancingDistributedSampler Test Example

Samples

To see torch-ort in action, see https://github.com/microsoft/onnxruntime-training-examples, which shows you how to train the most popular HuggingFace models.

🤓 Mixture of Experts

To run MoE, add the layer to your model as described in the tutorial: ort_moe/docs/tutorials/moe_tutorial.py

For more details, see ort_moe/docs/moe.md

Note: ONNX Runtime is not required to run the MoE layer. It is integrated in standalone PyTorch.

🎯 Inference

➕

ONNX Runtime for PyTorch supports PyTorch model inference using ONNX Runtime and Intel® OpenVINO™.

It is available via the torch-ort-infer python package. This package enables OpenVINO™ Execution Provider for ONNX Runtime by default for accelerating inference on various Intel® CPUs, Intel® integrated GPUs, and Intel® Movidius™ Vision Processing Units - referred to as VPU.

Supported Execution Providers

Execution Providers
OpenVINO

Provider Options

Users can configure different options for a given Execution Provider to run inference. As an example, OpenVINO™ Execution Provider options can be configured as shown below:

from torch_ort import ORTInferenceModule, OpenVINOProviderOptions
provider_options = OpenVINOProviderOptions(backend = "GPU", precision = "FP16")
model = ORTInferenceModule(model, provider_options = provider_options)

# PyTorch inference script follows

List of Provider Options

Supported backend-precision combinations:

Backend	Precision
CPU	FP32
GPU	FP32
GPU	FP16
MYRIAD	FP16

If no provider options are specified by user, OpenVINO™ Execution Provider is enabled with following options by default:

backend = "CPU"
precision = "FP32"

For more details on APIs, see usage.md.

Code Sample

Below is an example of how you can leverage OpenVINO™ integration with Torch-ORT in a simple NLP usecase.

A pretrained BERT model fine-tuned on the CoLA dataset from HuggingFace model hub is used to predict grammar correctness on a given input text.

from transformers 
import AutoTokenizer, AutoModelForSequenceClassification
import numpy as np
from torch_ort import ORTInferenceModule
tokenizer = AutoTokenizer.from_pretrained(
            "textattack/bert-base-uncased-CoLA")
model = AutoModelForSequenceClassification.from_pretrained(
        "textattack/bert-base-uncased-CoLA")
# Wrap model in ORTInferenceModule to prepare the model for inference using OpenVINO Execution Provider on CPU
model = ORTInferenceModule(model)
text = "Replace me any text by you'd like ."
encoded_input = tokenizer(text, return_tensors='pt')
output = model(**encoded_input)
# Post processing
logits = output.logits
logits = logits.detach().cpu().numpy()
# predictions
pred = np.argmax(logits, axis=1).flatten()
print("Grammar correctness label (0=unacceptable, 1=acceptable)")
print(pred)

Samples

To see OpenVINO™ integration with Torch-ORT in action, see demos, which shows you how to run inference on some of the most popular Deep Learning models.