目前，PyTorch官网尚未提供适用于Arm架构下支持GPU的PyTorch二进制文件。因此，需要通过源码编译的方式获取Pytorch及相关TorchVision和TorchAudio。该指导主要介绍如何编译并生成PyTorch安装包以及配套的TorchVision、TorchAudio安装包。
本次编译指导环境的Docker基础构建镜像以nvidia/cuda:12.4.0-devel-ubuntu20.04为基础镜像、Python 3.9版本和PyTorch 2.4.1版本为例，CUDA、PyTorch、TorchVision、TorchAudio以及Python的映射关系如表1所示。
PyTorch |
TorchVision |
TorchAudio |
CUDA |
Python |
|---|---|---|---|---|
2.5.0 |
0.20.0 |
2.5.0 |
11.8/12.1/12.4/12.6 |
>=3.8, <=3.11 |
2.4.1 |
0.19.1 |
2.4.1 |
11.8/12.1/12.4/12.6 |
>=3.8, <=3.11 |
2.4.0 |
0.19.0 |
2.4.0 |
11.8/12.1/12.4/12.6 |
>=3.8, <=3.11 |
2.3.1 |
0.18.1 |
2.3.1 |
11.8/12.1 |
>=3.8, <=3.11 |
2.3.0 |
0.18.0 |
2.3.0 |
11.8/12.1 |
>=3.8, <=3.11 |
2.2.2 |
0.17.2 |
2.2.2 |
11.8/12.1 |
>=3.8, <=3.11 |
2.2.1 |
0.17.1 |
2.2.1 |
11.8/12.1 |
>=3.8, <=3.11 |
2.2.0 |
0.17.0 |
2.2.0 |
11.8/12.1 |
>=3.8, <=3.11 |
2.1.2 |
0.16.2 |
2.1.2 |
11.8/12.1 |
>=3.8, <=3.11 |
2.1.1 |
0.16.1 |
2.1.1 |
11.8/12.1 |
>=3.8, <=3.11 |
2.1.0 |
0.16.0 |
2.1.0 |
11.8/12.1 |
>=3.8, <=3.11 |
2.0.1 |
0.15.2 |
2.0.2 |
11.7/11.8 |
>=3.8, <=3.11 |
2.0.0 |
0.15.0 |
2.0.0 |
11.7/11.8 |
>=3.8, <=3.11 |
1.13.1 |
0.14.1 |
0.13.1 |
11.6/11.7 |
>=3.7, <=3.10 |