A library for developing and deploying PyTorch models using ONNX Runtime.
Installation • Training • Inference • Docs • License
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™
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 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.
-
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
Mixture of Experts layer implementation is available in the ort_moe folder.
Clone this repo
git clone https://github.com/pytorch/ort.gitBuild MoE
cd ort_moe
python setup.py install- Ubuntu 18.04, 20.04
- Python* 3.7, 3.8 or 3.9
pip install torch-ort-infer[openvino]- Run post installation configuration script
python -m torch_ort.configure
-
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
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.
from torch_ort import ORTModule
model = ORTModule(model)
# PyTorch training script followsimport 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
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
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.
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.
➕
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.
| Execution Providers |
|---|
OpenVINO  |
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 followsSupported 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.
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)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.
Please refer to our contributing guide for more information on how to contribute!
This project has an MIT license, as found in the LICENSE file.