Linux如何搭建PyTorch开发环境

Linux搭建PyTorch开发环境

一 准备与规划

• 明确硬件与驱动：是否有 NVIDIA GPU，以及 CUDA 版本。使用命令查看驱动与CUDA运行时信息：nvidia-smi（右上角显示 CUDA Version）。如需安装或升级 NVIDIA 驱动，请先关闭图形界面，禁用 nouveau，再运行官方驱动安装包，完成后用 cat /proc/driver/nvidia/version 验证。若不使用GPU，可直接安装 CPU 版本。
• 选择工具链：优先使用 Anaconda/Miniconda 或 venv 做环境隔离；Python 建议使用 3.8–3.11 的稳定版本。
• 安装常用依赖：build-essential cmake git wget unzip yasm pkg-config libopenblas-dev liblapack-dev libjpeg-dev libpng-dev 等，便于后续扩展与多媒体数据处理。

二 安装与配置环境

• 方案A Conda（推荐）
  1. 安装 Miniconda（示例）：wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && bash Miniconda3-latest-Linux-x86_64.sh，按提示完成安装。
  2. 创建并激活环境：conda create -n pytorch python=3.10 -y && conda activate pytorch。
  3. 安装 PyTorch（GPU）：conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia；CPU 版：conda install pytorch torchvision torchaudio cpuonly -c pytorch。
• 方案B venv + pip
  1. 创建虚拟环境：python3 -m venv ~/venvs/pytorch && source ~/venvs/pytorch/bin/activate。
  2. 升级 pip：pip install -U pip。
  3. 安装 PyTorch（GPU）：pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121；CPU 版：pip install torch torchvision torchaudio。
• 国内镜像（可选，加速下载）
  • Conda 添加清华源（示例）：
    conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main/
conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge/
conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/pytorch/
conda config --set show_channel_urls yes
  • Pip 临时使用镜像：pip install -i https://pypi.tuna.tsinghua.edu.cn/simple torch torchvision torchaudio。

三 验证安装与常见问题

• 基础验证：
  python - <<'PY' import torch, torchvision, torchaudio print("torch:", torch.__version__, "torchvision:", torchvision.__version__, "torchaudio:", torchaudio.__version__) print("cuda available:", torch.cuda.is_available()) print("device count:", torch.cuda.device_count()) if torch.cuda.is_available(): print("current device:", torch.cuda.current_device(), torch.cuda.get_device_name()) PY
  预期：输出版本号，cuda available: True（GPU版），并显示设备数量与名称。
• 常见问题
  • torch.cuda.is_available() 为 False：检查 nvidia-smi 是否正常、驱动与 CUDA 版本匹配、是否误装了 CPU 版本、Conda 是否混用 pip 导致冲突。
  • 导入报错（如 NumPy 版本）：升级依赖 pip install -U numpy 后重试。
  • 驱动安装导致登录循环：安装 NVIDIA 驱动时可使用 -no-opengl-files 参数规避。

四 开发工具与后续配置

• IDE 与 Notebook：安装 VS Code 或 PyCharm，选择已创建的 conda/venv 解释器；如需交互式开发，安装 Jupyter：pip install jupyter 或 conda install jupyter，然后 jupyter notebook 启动。
• 常用科学计算库：pip install numpy pandas matplotlib scikit-learn；深度学习训练常用：pip install tqdm tensorboard。
• 环境导出与复现：
  • Conda：conda env export > environment.yml；复现：conda env create -f environment.yml。
  • Pip：pip freeze > requirements.txt；复现：pip install -r requirements.txt。
• 多版本 CUDA 与驱动建议：优先使用 Conda 提供的 cudatoolkit 与 pytorch-cuda 元包管理 CUDA 运行时，减少与系统 CUDA 的耦合；如需系统级 CUDA/cuDNN，请严格匹配版本并正确设置环境变量（如 CUDA_HOME, PATH, LD_LIBRARY_PATH）。