如何配置Ubuntu的PyTorch环境

Ubuntu 配置 PyTorch 环境

一 准备与检查

• 更新系统并安装基础工具：sudo apt update && sudo apt install -y build-essential cmake git wget python3 python3-pip
• 查看显卡与推荐驱动：ubuntu-drivers devices；如需自动安装推荐驱动：sudo ubuntu-drivers autoinstall，重启后执行 nvidia-smi 验证驱动与 CUDA 运行时版本。

二 创建虚拟环境

• 使用 Conda（推荐）：
  • 安装 Miniconda：wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && bash Miniconda3-latest-Linux-x86_64.sh && source ~/.bashrc
  • 创建环境：conda create -n pytorch_env python=3.8 或 python=3.10
  • 激活环境：conda activate pytorch_env
• 使用 venv（系统自带）：
  • python3 -m venv ~/venvs/pytorch && source ~/venvs/pytorch/bin/activate。

三 安装 PyTorch

• GPU 版本（使用 Conda，自动携带 cudatoolkit，避免与系统 CUDA 冲突）：
  • conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia
• GPU 版本（使用 Pip，指定 CUDA 版本）：
  • pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
• CPU 版本：
  • pip install torch torchvision torchaudio
• 版本匹配提示：torchvision 与 torch 版本需对应，安装前可查阅官方安装指引或版本对应说明。

四 验证与常见问题

• 基础验证：
  • python - <<‘PY’ import torch, torchvision, torchaudio print(“torch:”, torch.version, “torchvision:”, torchvision.version, “torchaudio:”, torchaudio.version) print(“cuda available:”, torch.cuda.is_available()) if torch.cuda.is_available(): print(“device count:”, torch.cuda.device_count(), “current device:”, torch.cuda.current_device()) PY
• 若使用系统级 CUDA/cuDNN（非 Conda 自带 cudatoolkit），需配置环境变量（示例路径，按实际修改）：
  • echo ‘export PATH=/usr/local/cuda/bin:$PATH’ >> ~/.bashrc
  • echo ‘export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH’ >> ~/.bashrc
  • source ~/.bashrc
• 常见问题速查：
  • torch.cuda.is_available() 为 False：优先确认是否在 Conda 环境中安装了带 cudatoolkit 的 PyTorch；若使用系统 CUDA，核对驱动与 CUDA/cuDNN 版本匹配并正确设置环境变量；必要时重装匹配版本的驱动与库。