The main dependency is onnxruntime-training. It is only available on Linux. It is available from pypi for CPU. GPU versions are available Its installation replaces onnxruntime and includes onnxruntime and onnxruntime-training.

Installation of onnxruntime-training for GPU#

onnxruntime-training is only available on Linux. The CPU can be installed with the following instruction.

pip install onnxruntime-training --extra-index-url

Versions using GPU with CUDA or ROCm are available. Check to find a specific version. You can use it on Windows inside WSL (Windows Linux Subsystem) or compile it for CPU:

python tools\ci_build\ --skip_tests --build_dir .\build\Windows --config Release --build_shared_lib --build_wheel --numpy_version= --cmake_generator="Visual Studio 16 2019" --enable_training --enable_training_ops --enable_training_torch_interop

GPU versions work better on WSL, see Build onnxruntime on WSL (Windows Linux Subsystem).

Installation of onnxcustom#

pip install onnxcustom