环境信息
- 硬件: Jetson Orin Nano
- 系统: Ubuntu 20.04 + JetPack
- 目标: 在Docker容器中使用宿主机的CUDA/cuDNN/PyTorch环境
问题背景
Jetson设备预装了优化的CUDA、cuDNN和PyTorch环境，但在Docker容器中默认无法访问这些GPU资源。本文介绍如何让容器继承宿主机的完整PyTorch环境。
解决方案
1. 检查宿主机环境
# 验证宿主机CUDA环境
nvcc --version
find /usr/lib -name "*cudnn*" | head -5
ls /home/jetson/torch-*.whl
2. 创建支持GPU的容器
docker run -it \
--net=host \
--privileged \
--device=/dev/nvidia0 \
--device=/dev/nvidiactl \
--device=/dev/nvidia-modeset \
--device=/dev/nvhost-ctrl \
--device=/dev/nvhost-ctrl-gpu \
--device=/dev/nvhost-prof-gpu \
--device=/dev/nvmap \
--device=/dev/nvhost-gpu \
--device=/dev/nvhost-as-gpu \
-v /dev:/dev \
-v /usr/local/cuda-11.4:/usr/local/cuda-11.4 \
-v /home/jetson:/home/jetson \
-v /usr/lib/aarch64-linux-gnu/tegra:/usr/lib/aarch64-linux-gnu/tegra \
-v /usr/lib/aarch64-linux-gnu/libcudnn.so.8:/usr/lib/aarch64-linux-gnu/libcudnn.so.8:ro \
-v /usr/lib/aarch64-linux-gnu/libcudnn.so.8.6.0:/usr/lib/aarch64-linux-gnu/libcudnn.so.8.6.0:ro \
-v /usr/lib/aarch64-linux-gnu/libcudnn_cnn_infer.so.8:/usr/lib/aarch64-linux-gnu/libcudnn_cnn_infer.so.8:ro \
-v /usr/lib/aarch64-linux-gnu/libcudnn_ops_infer.so.8:/usr/lib/aarch64-linux-gnu/libcudnn_ops_infer.so.8:ro \
-e DISPLAY=$DISPLAY \
-e CUDA_ROOT=/usr/local/cuda \
-e PATH=/usr/local/cuda-11.4/bin:$PATH \
-e LD_LIBRARY_PATH=/usr/local/cuda-11.4/lib64:/usr/lib/aarch64-linux-gnu/tegra:/usr/lib/aarch64-linux-gnu:$LD_LIBRARY_PATH \
--env="QT_X11_NO_MITSHM=1" \
-v /tmp/.X11-unix:/tmp/.X11-unix \
your_name:xxx
3. 安装依赖库
# 安装数学计算库
apt-get update
apt-get install -y libopenblas-base libopenblas-dev libblas3 liblapack3
# 更新库缓存
ldconfig
4. 安装PyTorch
# 使用宿主机的PyTorch wheel文件
pip3 install /home/jetson/torch-2.1.0a0+41361538.nv23.06-cp38-cp38-linux_aarch64.whl
5. 验证环境
python3 -c "
import torch
print(f'PyTorch: {torch.__version__}')
print(f'CUDA: {torch.cuda.is_available()}')
print(f'Device: {torch.cuda.get_device_name(0) if torch.cuda.is_available() else \"CPU\"}')
# GPU计算测试
if torch.cuda.is_available():
x = torch.randn(1000, 1000).cuda()
y = torch.matmul(x, x)
print(f'GPU test: ✅ {y.shape} on {y.device}')
"
关键要点
设备映射
- 映射所有GPU相关设备文件（
/dev/nvidia*,/dev/nvhost-*） - 使用
--privileged确保设备访问权限
库路径挂载
/usr/lib/aarch64-linux-gnu/tegra- Tegra GPU驱动/usr/local/cuda-11.4- CUDA工具包/usr/lib- 系统库（包含cuDNN）
环境变量
LD_LIBRARY_PATH必须包含所有GPU库路径CUDA_ROOT指向CUDA安装目录
依赖处理
- PyTorch需要OpenBLAS等数学库
- 容器中需要重新安装这些依赖
保存环境
# 配置完成后保存为镜像
docker commit <container_id> jetson-pytorch:ready
解决torchvision环境
pip3 install /home/jetson/vision-release-0.16/dist/torchvision-0.16.2-cp38-cp38-linux_aarch64.whl
同时，注意删除yolov5中requirements的torchvison依赖，这个傻逼软件启动会自动下载错误版本
然后运行这个下载yolov5的依赖
pip install -r requirements.txt --no-deps
常见问题
Q: torch.cuda.is_available() 返回 False？ A: 检查设备映射和LD_LIBRARY_PATH是否包含Tegra库路径。
Q: libopenblas.so.0 找不到？ A: 安装 apt-get install libopenblas-base。
Q: GPU性能如何？ A: Jetson Orin Nano提供20 TOPS算力，支持实时AI推理。
结果
成功配置后可获得：
- ✅ PyTorch 2.1.0（Jetson优化版）
- ✅ CUDA 11.4 + cuDNN 8.6
- ✅ 完整GPU加速支持
- ✅ 20 TOPS NPU算力可用
通过以上配置，Docker容器可完美继承Jetson宿主机的PyTorch环境，实现GPU加速的机器学习应用。
评论区