基于Ubuntu的深度学习训练环境搭建步骤是什么？

2026-08-25 02:49:42 分类：云计算

下面给出一个在 Ubuntu 上搭建深度学习训练环境的通用步骤（以 PyTorch + CUDA 为例，TensorFlow 类似）。你可以根据需要调整版本和框架。

一、系统准备

1. 确认硬件与驱动

  • 显卡：NVIDIA GPU（推荐）
  • 查看显卡信息： 
    nvidia-smi
  • 若未安装或版本过低，先安装/更新 NVIDIA 驱动： 
    sudo apt update
sudo ubuntu-drivers devices          # 查看推荐驱动
sudo apt install nvidia-driver-535   # 示例：535 为常用 LTS 驱动号
sudo reboot
  • 重启后再次检查： 
    nvidia-smi

2. 安装基础工具

sudo apt update
sudo apt install -y git curl wget vim build-essential 
    python3-pip python3-venv cmake pkg-config libgl1-mesa-glx 
    libglib2.0-0 libsm6 libxrender1 libxext6

二、安装 Python 与虚拟环境

1. 使用系统 Python（推荐 3.8+）

Ubuntu 通常自带 python3，检查版本：

python3 --version
pip3 --version

若版本过旧，可添加 deadsnakes PPA 安装新版：

sudo add-apt-repository ppa:deadsnakes/ppa
sudo apt update
sudo apt install python3.10 python3.10-venv python3.10-dev

2. 创建项目虚拟环境

mkdir ~/dl_env && cd ~/dl_env
python3.10 -m venv venv
source venv/bin/activate

激活后：

python --version
pip --version

三、安装 CUDA Toolkit 与 cuDNN

注意：CUDA Toolkit 版本要与你的 PyTorch/TensorFlow 版本匹配。
一般做法：先确定你要用的深度学习框架支持的 CUDA 版本，再安装对应 CUDA。

1. 选择 CUDA 版本

例如：PyTorch 2.3 支持 CUDA 12.1 / 11.8 等，具体见官方文档。
假设我们要用 CUDA 12.1。

2. 下载并安装 CUDA Toolkit

访问 https://developer.nvidia.com/cuda-downloads
选择：

  • Operating System: Linux
  • Distribution: Ubuntu
  • Architecture: x86_64
  • Installer Type: runfile (deb) 或 deb

方式 A：使用 .deb 包（推荐）

wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb
sudo dpkg -i cuda-keyring_1.1-1_all.deb
sudo apt update
sudo apt install -y cuda-toolkit-12-1

方式 B：使用 .run 文件

wget https://developer.download.nvidia.com/compute/cuda/12.1.0/local_installers/cuda_12.1.0_530.30.02_linux.run
chmod +x cuda_*.run
sudo ./cuda_*.run --toolkit --silent

3. 配置环境变量

编辑 ~/.bashrc（或当前用户的 .profile）：

export PATH=/usr/local/cuda-12.1/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/cuda-12.1/lib64:$LD_LIBRARY_PATH

然后：

source ~/.bashrc

验证：

nvcc --version

4. 安装 cuDNN

cuDNN 需从 NVIDIA 官网注册下载，按说明解压到 CUDA 目录：

# 假设已下载 cudnn-12.x-linux-x64-v9.x.x.tar.gz
tar -xzvf cudnn-*.tar.gz
sudo cp cuda/include/* /usr/local/cuda-12.1/include/
sudo cp cuda/lib64/* /usr/local/cuda-12.1/lib64/
sudo chmod a+r /usr/local/cuda-12.1/include/*
sudo chmod a+r /usr/local/cuda-12.1/lib64/*

四、安装深度学习框架（以 PyTorch 为例）

1. 升级 pip 和 setuptools

pip install --upgrade pip setuptools wheel

2. 根据 CUDA 版本安装 PyTorch

进入 PyTorch 官网：https://pytorch.org/get-started/locally/
选择：

  • Package: Pip
  • Language: Python
  • Compute Platform: CUDA 12.1（与你安装的 CUDA 一致）

示例命令（CUDA 12.1）：

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

验证：

python -c "import torch; print(torch.__version__); print(torch.cuda.is_available()); print(torch.cuda.get_device_name(0))"

五、安装常用依赖库

根据你的项目需求安装，例如：

pip install numpy pandas matplotlib seaborn scikit-learn pillow tqdm
pip install opencv-python-headless
pip install jupyterlab
pip install ipykernel
ipython kernel install --user --name=dl-env

如需多卡训练，可安装分布式相关包：

pip install torch.distributed

六、测试训练环境

创建一个简单脚本 test_gpu.py：

import torch

print("PyTorch version:", torch.__version__)
print("CUDA available:", torch.cuda.is_available())
if torch.cuda.is_available():
    print("GPU count:", torch.cuda.device_count())
    for i in range(torch.cuda.device_count()):
        print(f"GPU {i}:", torch.cuda.get_device_name(i))

    x = torch.randn(10, 10).cuda()
    y = x @ x.T
    print("Matrix multiplication on GPU succeeded.")
else:
    print("No GPU found.")

运行：

python test_gpu.py

七、可选：配置 Jupyter / VSCode

Jupyter

jupyter lab

浏览器打开 http://localhost:8888，选择 dl-env 内核。

VS Code

安装插件：

  • Python
  • Jupyter
  • Remote – SSH（如需远程开发）

在 VS Code 中选择解释器为 ~/dl_env/venv/bin/python。

八、常见问题排查

  1. torch.cuda.is_available() 返回 False

    • 检查 nvidia-smi 是否正常
    • 确认 CUDA Toolkit 与驱动版本兼容
    • 确认 PyTorch 安装的是 cuXXX 版本，而非 CPU 版

  2. nvcc 找不到

    • 检查 PATH 是否包含 /usr/local/cuda-XX.X/bin
    • 重新 source ~/.bashrc

  3. 显存不足 / OOM

    • 减小 batch size
    • 使用梯度累积 / 混合精度（AMP）
    • 检查是否有其他进程占用 GPU：nvidia-smi

如果你告诉我：

  • Ubuntu 版本（如 20.04 / 22.04）
  • 期望的框架（PyTorch / TensorFlow / JAX）
  • 目标 CUDA 版本（或希望自动适配）

我可以给出一套更精确、可直接复制粘贴的命令序列。