共计 1568 个字符，预计需要花费 4 分钟才能阅读完成。
背景痛点：为什么 arm64 安装 PyTorch-GPU 这么难？
在 x86 架构上安装 PyTorch-GPU 版本只需要简单执行pip install torch torchvision，但在 arm64 设备（如 NVIDIA Jetson 系列）上会遇到几个典型问题：
- 官方源缺失：PyPI 官方仓库没有提供 arm64 架构的预编译 wheel 包，直接 pip 安装会提示
No matching distribution found - 驱动兼容性 ：Jetson 设备的 CUDA 驱动版本与 PyTorch 预编译包的 ABI 兼容性要求严格，版本不匹配会导致
libcudart.so加载失败 - 内存限制：Jetson 设备的共享内存较小，默认安装可能触发 OOM 错误
技术方案：conda 还是 pip？
方案对比
- conda：
- 优点：自动解决依赖关系
-
缺点：官方 conda 通道同样缺少 arm64 的 torch-gpu 包
-
pip + 预编译 wheel：
- 推荐使用 NVIDIA 官方提供的预编译包（下载页面）
- 需要手动指定
--extra-index-url指向 NVIDIA 的仓库
实战步骤
1. 环境检测
首先确认设备架构和 CUDA 版本：
# 查看 CPU 架构
uname -m # 应输出 aarch64
# 检查 CUDA 编译器版本
nvcc --version # 例如输出 CUDA 10.22. 安装 PyTorch-GPU
根据 CUDA 版本选择对应的 wheel（以 CUDA 10.2 为例）：
pip install --extra-index-url https://developer.download.nvidia.com/compute/redist/jp/v50 torch==1.9.03. 验证安装
运行 Python 验证脚本：
import torch
print(f"PyTorch 版本: {torch.__version__}")
print(f"CUDA 可用: {torch.cuda.is_available()}")
print(f"设备数量: {torch.cuda.device_count()}")
print(f"当前设备: {torch.cuda.current_device()}")
print(f"设备名称: {torch.cuda.get_device_name(0)}")避坑指南
错误 1：No CUDA runtime is found
现象 ：torch.cuda.is_available() 返回 False
解决方案：
1. 检查 /usr/local/cuda/bin 是否在 PATH 中
2. 验证 libcudart.so 的符号链接：
ls -l /usr/local/cuda/lib64/libcudart.so*错误 2：Jetson 内存不足
优化建议：
– 安装时添加 --no-cache-dir 选项
– 使用 swap 分区扩展虚拟内存
性能验证
对比 CPU 和 GPU 的矩阵运算速度：
import time
device = 'cuda' if torch.cuda.is_available() else 'cpu'
x = torch.randn(10000, 10000)
# CPU 计算
t_start = time.time()
x_cpu = x.to('cpu')
y_cpu = x_cpu @ x_cpu.T
print(f"CPU 耗时: {time.time() - t_start:.4f}s")
# GPU 计算
if device == 'cuda':
t_start = time.time()
x_gpu = x.to(device)
y_gpu = x_gpu @ x_gpu.T
print(f"GPU 耗时: {time.time() - t_start:.4f}s")思考题
如何为自定义的 arm64 设备交叉编译 PyTorch？需要考虑：
1. 目标设备的 ABI 兼容性
2. CUDA 工具链的版本匹配
3. 内存页锁定 (pinned memory) 的支持情况
欢迎在评论区分享你的经验！
正文完
