在Ubuntu上部署PyTorch应用要注意什么

Ubuntu 部署 PyTorch 的关键注意事项

一 系统与驱动准备

• 选择 Ubuntu 18.04/20.04/22.04 LTS 等稳定版本，先更新系统并确认内核与发行版信息：cat /etc/os-release、uname -r。
• GPU 场景优先用官方工具安装驱动：sudo ubuntu-drivers autoinstall && sudo reboot，重启后用 nvidia-smi 验证。注意 nvidia-smi 显示的 CUDA Version 是驱动支持的最高运行时版本，并不等同于已安装的 CUDA Toolkit 版本；PyTorch 官方预编译包通常向下兼容，例如驱动支持 CUDA 12.0 也可选择 CUDA 11.8 构建的 PyTorch。
• 若 nvidia-smi 不可用，检查 Secure Boot、显卡直连/虚拟化、驱动签名与内核模块加载（lsmod | grep nvidia）。

二 Python 环境与安装方式

• 使用 Miniconda 管理隔离环境（更轻量、依赖解析更省内存）：conda create -n torch_env python=3.8-3.11 -y && conda activate torch_env。低内存（如 ≤4GB RAM）机器优先用 pip 安装 PyTorch，依赖解析更轻量；必要时增加 swap（如 fallocate -l 2G /swapfile && chmod 600 /swapfile && mkswap /swapfile && swapon /swapfile）。
• GPU 安装建议：
  • Conda（同时处理 CUDA/cuDNN 依赖，一致性更好）：conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia
  • Pip（网络与缓存友好）：pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
• CPU-only 场景：pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu。
• 避免混用 pip 与 conda 的二进制包；同一环境尽量统一安装渠道。安装后用 python -c "import torch; print(torch.__version__, torch.cuda.is_available())" 验证。

三 版本匹配与常见故障排查

• 版本匹配要点：
  • 驱动通过 nvidia-smi 查看支持的最高 CUDA Runtime；
  • 运行 nvcc --version 查看实际 CUDA Toolkit 版本（若系统单独安装了 Toolkit）；
  • PyTorch 的 CUDA 标签（如 cu118）需与驱动支持范围匹配。
• 典型问题与处理：
  • torch.cuda.is_available() 为 False：检查驱动加载、nvidia-smi 输出、PyTorch 是否为 GPU 版本、虚拟环境是否激活。
  • ImportError: libcublas.so.11 等库缺失：确认 CUDA 库路径（如 /usr/local/cuda/lib64）是否在 LD_LIBRARY_PATH，必要时添加：echo 'export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH' >> ~/.bashrc && source ~/.bashrc。
  • 依赖解析或安装阶段 MemoryError：低内存机器优先用 Miniconda + pip、减少并行任务、增加 swap。

四 性能与多 GPU 部署

• 内存分配优化：设置环境变量以缓解碎片与回收滞后，例如
  export PYTORCH_CUDA_ALLOC_CONF=garbage_collection_threshold:0.6,max_split_size_mb:128。
• 多 GPU 训练：
  • 安装 NCCL 通信库：sudo apt-get install libnccl2 libnccl-dev；
  • 设置通信与网络：export NCCL_DEBUG=INFO、export NCCL_SOCKET_IFNAME=eth0（按实际网卡名调整）；
  • 优先使用 DistributedDataParallel（DDP），并合理设置 batch_size、num_workers 与 pin_memory。
• 大模型与资源规划：确保 系统内存 ≥16GB（推荐 32GB+），并使用混合精度训练（torch.cuda.amp）降低显存占用、提升吞吐。

五 上线与运维实践

• 使用虚拟环境隔离不同项目与版本；为每次上线记录并锁定依赖（requirements.txt 或 environment.yml，如 pip freeze > requirements.txt）。
• 采用 Docker 封装运行环境，减少“本机能跑、服务器不行”的环境漂移；镜像中仅保留运行所需依赖。
• 日志与监控：输出 GPU 利用率/显存、关键指标与异常堆栈；必要时接入 Prometheus + Grafana 或云监控。
• 安全与维护：限制容器/进程权限，定期更新驱动与依赖，回滚方案与灰度发布策略就绪。