Introduction
Currently, the PyTorch official website does not provide the PyTorch binary file that supports GPUs in the Arm architecture. Therefore, you need to obtain PyTorch and related TorchVision and TorchAudio by compiling the source code. This section describes how to compile and generate a PyTorch installation package and the mapping TorchVision and TorchAudio installation packages.
In this section, nvidia/cuda:12.4.0-devel-ubuntu20.04 is used as the basic image, and Python 3.9 and PyTorch 2.4.1 are used. Table 1 lists the mapping between CUDA, PyTorch, TorchVision, TorchAudio, and Python.
PyTorch |
TorchVision |
TorchAudio |
CUDA |
Python |
|---|---|---|---|---|
2.5.0 |
0.20.0 |
2.5.0 |
11.8/12.1/12.4/12.6 |
3.8 to 3.11 |
2.4.1 |
0.19.1 |
2.4.1 |
11.8/12.1/12.4/12.6 |
3.8 to 3.11 |
2.4.0 |
0.19.0 |
2.4.0 |
11.8/12.1/12.4/12.6 |
3.8 to 3.11 |
2.3.1 |
0.18.1 |
2.3.1 |
11.8/12.1 |
3.8 to 3.11 |
2.3.0 |
0.18.0 |
2.3.0 |
11.8/12.1 |
3.8 to 3.11 |
2.2.2 |
0.17.2 |
2.2.2 |
11.8/12.1 |
3.8 to 3.11 |
2.2.1 |
0.17.1 |
2.2.1 |
11.8/12.1 |
3.8 to 3.11 |
2.2.0 |
0.17.0 |
2.2.0 |
11.8/12.1 |
3.8 to 3.11 |
2.1.2 |
0.16.2 |
2.1.2 |
11.8/12.1 |
3.8 to 3.11 |
2.1.1 |
0.16.1 |
2.1.1 |
11.8/12.1 |
3.8 to 3.11 |
2.1.0 |
0.16.0 |
2.1.0 |
11.8/12.1 |
3.8 to 3.11 |
2.0.1 |
0.15.2 |
2.0.2 |
11.7/11.8 |
3.8 to 3.11 |
2.0.0 |
0.15.0 |
2.0.0 |
11.7/11.8 |
3.8 to 3.11 |
1.13.1 |
0.14.1 |
0.13.1 |
11.6/11.7 |
3.7 to 3.10 |
Prerequisites
- The compilation server is in the Arm architecture and is equipped with GPU hardware. Docker has been installed and the network connection is normal.
- NVIDIA Container Toolkit has been installed on the compilation server.
- A CUDA image (nvidia/cuda:12.6.0-devel-ubuntu20.04 for example) has been pulled using docker pull. After the image is pulled, access the container.