Can't install onnxruntime-training on JetPack 5.1.2 (Jetson Orin)

Autonomous Machines Jetson & Embedded Systems Jetson AGX Orin
, , ,
1

I have set up a virtual environment with Python3.8.10 inside a Docker container on the Jetson Orin Dev Kit. I want to run ONNX on device training within this venv. Therefore, I installed the onnx-runtime GPU version 1.17.0, which worked perfectly.

(py38) $ wget https://nvidia.box.com/shared/static/zostg6agm00fb6t5uisw51qi6kpcuwzd.whl -O onnxruntime_gpu-1.17.0-cp38-cp38-linux_aarch64.whl
(py38) $ pip install onnxruntime_gpu-1.17.0-cp38-cp38-linux_aarch64.whl

Additionally, for on-device training, the onnxruntime-training package is required. According to the getting started installation instructions by onnx-runtime, it can be installed as shown in the picture below.

image
image1744×551 67.4 KB

But when I run this command inside my virtual environment the following error occures

(py38) $ pip install -i https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ORT/pypi/simple/ onnxruntime-training
Looking in indexes: https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ORT/pypi/simple/
ERROR: Could not find a version that satisfies the requirement onnxruntime-training (from versions: none)
ERROR: No matching distribution found for onnxruntime-training

Any ideas on how I could get this to work?

Here are the Jetson informations:

$ sudo apt-cache show nvidia-jetpack
Package: nvidia-jetpack
Version: 5.1.2-b104
Architecture: arm64
Maintainer: NVIDIA Corporation
Installed-Size: 194
Depends: nvidia-jetpack-runtime (= 5.1.2-b104), nvidia-jetpack-dev (= 5.1.2-b104)
Homepage: http://developer.nvidia.com/jetson
Priority: standard
Section: metapackages
Filename: pool/main/n/nvidia-jetpack/nvidia-jetpack_5.1.2-b104_arm64.deb
Size: 29304
SHA256: fda2eed24747319ccd9fee9a8548c0e5dd52812363877ebe90e223b5a6e7e827
SHA1: 78c7d9e02490f96f8fbd5a091c8bef280b03ae84
MD5sum: 6be522b5542ab2af5dcf62837b34a5f0
Description: NVIDIA Jetpack Meta Package
Description-md5: ad1462289bdbc54909ae109d1d32c0a8
3

I managed to solve it. Now ONNX Runtime (ORT) on-device training works with GPU support on the Jetson Orin Dev Kit with JetPack v5.1.2. Unfortunately for ORT with training and GPU support, no ready-to-install Python wheel is available for an ARM architecture. Thus, it needs to be built from the source. Therefore the following steps are required:

  1. Make sure to have `cmake >= v3.26 installed how-to
  2. Make sure to have gcc, g++ v11.5 or newer installed how-to
  3. ORT training depends on Torch. Since JetPack v5.1.2 forces Torch v2.1.0 and ORT recommends using the latest Torch version, ORT v1.16.3 should be used. Thus, the next step is to clone the ORT repository and check out the tag v1.16.3.
git clone --recursive https://github.com/Microsoft/onnxruntime.git
cd onnxruntime
git checkout v1.16.3
  1. Build ORT and the Python wheel. You might need to adjust the CUDA paths and version:
./build.sh --config RelWithDebInfo --build_shared_lib --parallel --enable_training --build_wheel --use_cuda --cuda_home /usr/local/cuda --cudnn_home /usr/local/cuda --cuda_version=11.4 --skip_tests --compile_no_warning_as_error
  1. Install the Python wheel:
python -m pip install build/Linux/RelWithDebInfo/dist/onnxruntime_training-1.16.3*.whl

ORT recommends using ONNX v1.14.1 for ORT v1.16.x`.

4

Hi,

Sorry for the late update.

The official website doesn’t support the Jetson environment.
But we do provide the packages on the link below:
https://elinux.org/Jetson_Zoo#ONNX_Runtime

Your instructions to build ONNXRuntime from the source also look correct.
So you can use the package directly.

Thanks.

5

Hi,

are you sure that the packages from the link you provided support ONNX Runtime on-device training? I think I tried one of these packages and import onnxruntime.training failed with onnxruntime does not have an attribute training.

6

Hi,

The package is provided by the ONNXRuntime team for the Jetson user.
In general, it should enable most of the common APIs.

If you need further support, please raise an issue in the below repo:

Thanks.

8

This topic was automatically closed 14 days after the last reply. New replies are no longer allowed.