training Pytorch models with onnxruntime
Project description
Train PyTorch models with ONNX Runtime
ONNX Runtime for PyTorch accelerates PyTorch model training using ONNX Runtime.
It is available via the torch-ort python package.
This repository contains the source code for the package as well as instructions for running the package and samples demonstrating how to do so.
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.
Run in a Python environment
Default dependencies
By default, torch-ort depends on PyTorch 1.8.1, ONNX Runtime 1.8 and CUDA 10.2.
-
Install CUDA 10.2
-
Install CuDNN 7.6
-
Install torch-ort and dependencies
pip install ninjapip install torch-ort
Explicitly install for NVIDIA CUDA 10.2
-
Install CUDA 10.2
-
Install CuDNN 7.6
-
Install torch-ort and dependencies
pip install ninjapip install torch==1.8.1pip install --pre onnxruntime-training -f https://onnxruntimepackages.z14.web.core.windows.net/onnxruntime_nightly_cu102.htmlpip install torch-ort
Explicitly install for NVIDIA CUDA 11.1
-
Install CUDA 11.1
-
Install CuDNN 8.0
-
Install torch-ort and dependencies
pip install ninjapip install torch==1.8.1pip install --pre onnxruntime-training -f https://onnxruntimepackages.z14.web.core.windows.net/onnxruntime_nightly_cu111.htmlpip install torch-ort
Explicitly install for AMD ROCm 4.1
-
Install ROCm 4.1 base package (instructions)
-
Install ROCm 4.1 libraries (instructions)
-
Install ROCm 4.1 RCCL (instructions)
-
Install torch-ort and dependencies
pip install ninjapip install --pre torch -f https://download.pytorch.org/whl/nightly/rocm4.1/torch_nightly.htmlpip install --pre onnxruntime-training -f https://onnxruntimepackages.z14.web.core.windows.net/onnxruntime_nightly_rocm41.htmlpip install torch-ort
Run using Docker
The docker directory contains dockerfiles for the NVIDIA CUDA 11.1 configuration.
-
Build the docker image
docker build -f Dockerfile.ort-cu111-cudnn8-devel-ubuntu18.04 -t ort.cu111 . -
Run the docker container using the image you have just built
docker run -it --gpus all --name my-experiments ort.cu111:latest /bin/bash
Test your installation
- Clone this repo
git clone git@github.com:pytorch/ort.git
- Install extra dependencies
pip install wget pandas sklearn transformers
- Run the training script
python ./ort/tests/bert_for_sequence_classification.py
Add ONNX Runtime for PyTorch to your PyTorch training script
import onnxruntime
from torch_ort import ORTModule
model = ORTModule(model)
# PyTorch training script follows
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distributions
Built Distribution
File details
Details for the file torch_ort-0.0.10.dev20210512-py3-none-any.whl.
File metadata
- Download URL: torch_ort-0.0.10.dev20210512-py3-none-any.whl
- Upload date:
- Size: 3.5 kB
- Tags: Python 3
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/3.4.1 importlib_metadata/4.0.1 pkginfo/1.7.0 requests/2.22.0 requests-toolbelt/0.9.1 tqdm/4.60.0 CPython/3.8.5
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 | 9c0703503b22e69517074d72895d5d8db97a5460f039c95427f71a425a8e3edb |
|
| MD5 | 415b6198a619b68ebfc686bc8b3137d5 |
|
| BLAKE2b-256 | 3893168f674fb1ac1acf8a180926e8e531fbf32a49d41f24f6a5ba53bf7fef5a |
