Automating PyTorch ARM Builds with Docker BuildX for Nvidia CUDA and Python >...

My Workflow

One of the more urgent needs in the community was to automate the building of ARM wheels of PyTorch to utilize PyTorch on Nvidia Devices (e.g. Nvidia Jetson Nano) with CUDA enabled. Nvidia currently offers this but only with Python 3.6 (while we are at Python 3.11 now and many packages require Python > 3.6)

Therefore, a way had to be found to start automating the building of these wheels. However, it's not straightforward on how we can get this to build seeing that we require CUDA (for GPU acceleration on edge devices) which is required for AI models to run smoothly.

The entire process above took me around ~11 full days, starting of with figuring out how to build the Dockerfile and finally automating the CI process.

Full Write Up & Source Code

To be able to keep this post within reading limits, I decided to publish the entire build on my personal blog: https://xaviergeerinck.com/post/iot/nvidia-building-pytorch.

The source code can be found on GitHub with a build of the resulting wheel and including GitHub Action Workflow.

Contributions

In any kind of project of this size, there are specific contributions that were made. In my project I believe to have made the following:

• Install CUDA on non-GPU devices
• Compile PyTorch with CUDA enabled on non-GPU devices
• Compile PyTorch for Python > 3.6
• Build for ARM with CI through Docker Buildx

Utilized Actions

Since I did not want to reinvent the wheel I reused some actions:

1. docker/setup-buildx-action

   • I cross-compile for ARM on AMD64 machines in the pipeline

2. docker/setup-qemu-action

   • Configure QEMU to be able to compile for ARM and install the QEMU static binaries

3. actions/checkout

   • Check out a repo

4. actions/cache

   • Allow us to cache the docker layers

5. actions/upload-artifact

   • Upload the output of a directory to GitHub artifacts

Workflow Outline

1. When a release is created trigger the action (or pushed to master at the moment - will change seeing the LONG compilation time)
2. Clone the repository
3. Setup Docker with Buildx
4. Run our container
5. Copy over the Built Wheel to an artifact on GitHub

Submission Category:

DIY Deployments, Interesting IoT

Yaml File or Link to Code

https://github.com/XavierGeerinck/Jetson-Linux-PyTorch

Note: I had to put the link as the recommended liquid syntax gave an error ({% https://github.com/XavierGeerinck/Jetson-Linux-PyTorch %})

Additional Resources / Info