3. 开发环境搭建：Ubuntu 22.04 + CUDA 12.1 + PyTorch 2.x + TensorRT 8.6 保姆级安装

环境搭建，说白了就是给模型训练和部署铺路。我见过太多同学在环境上卡一两天，最后发现是驱动版本不对。这一章，咱们一步不跳，把整个环境捋顺。

本章目标：从零搭建一套能跑 YOLO 训练、又能用 TensorRT 加速推理的完整环境。硬件假设你手头有一张 NVIDIA 显卡（RTX 30/40 系列最佳）。

3.1 整体架构一览

先看一张图，搞清楚每个组件的位置。这样你装的时候心里有数，不会装完发现少了东西。

YOLO 训练与边缘部署环境栈 YOLOv8 / YOLOv5 训练脚本 PyTorch 2.x + TorchVision TensorRT 8.6 + cuDNN 8.9 CUDA 12.1 + NVIDIA Driver 535+ Ubuntu 22.04 LTS 训练 & 导出 ONNX 训练框架 模型优化 & 推理 推理引擎 GPU 计算加速 底层驱动

安装顺序必须从下往上：系统 → 驱动 → CUDA → cuDNN → PyTorch → TensorRT。跳步容易出问题，我吃过这个亏。

3.2 第一步：Ubuntu 22.04 系统安装

如果你已经有 Ubuntu 22.04 系统，直接跳到 3.3 节。没有的话，我建议用 U 盘做启动盘安装。

3.2.1 系统要求

项目 推荐配置 最低配置
CPU Intel i7 / AMD Ryzen 7 以上 Intel i5 / AMD Ryzen 5
内存 32 GB 16 GB
显卡 NVIDIA RTX 4070 及以上 NVIDIA GTX 1660
硬盘 NVMe SSD 1 TB SSD 256 GB
网络 稳定外网（下载依赖） 可访问 PyPI / GitHub

3.2.2 安装要点

  • 分区建议： / 分区 100 GB，/home 分区剩余全部。数据放 /home 下，重装系统不丢数据。
  • 软件更新： 安装时勾选「为图形和 Wi-Fi 硬件安装第三方软件」。
  • 语言： 选 English，避免终端路径出现中文乱码。
我的习惯： 装完系统第一件事，跑 sudo apt update && sudo apt upgrade -y。把内核和基础库更新到最新，省得后面装驱动时报依赖错误。

3.3 第二步：NVIDIA 驱动安装

驱动是 GPU 和系统之间的翻译官。版本不对，CUDA 装上也白搭。

3.3.1 检查显卡型号

lspci | grep -i nvidia

如果输出为空，检查显卡是否插好，或者 BIOS 里是否禁用了独立显卡。

3.3.2 安装驱动（推荐方法）

我个人最推荐用 NVIDIA 官方 runfile 安装，干净可控。

  1. 去 NVIDIA 官网 下载对应显卡的驱动（Linux 64-bit .run 文件）。
  2. 禁用开源驱动 nouveau：
sudo bash -c "echo blacklist nouveau > /etc/modprobe.d/blacklist-nvidia-nouveau.conf"
sudo update-initramfs -u
sudo reboot
  1. 重启后进入文本模式（Ctrl+Alt+F2），登录后：
sudo service gdm stop
chmod +x NVIDIA-Linux-*.run
sudo ./NVIDIA-Linux-*.run
  1. 安装完成后重启，用 nvidia-smi 验证。
我曾经踩过的坑： 用 Ubuntu 自带的「附加驱动」工具安装，有时会装到 470 或 525 这种旧版本。CUDA 12.1 要求驱动版本 ≥ 535，装完才发现不兼容，又得重装。所以，老老实实去官网下载。

3.4 第三步：CUDA 12.1 安装

CUDA 是 NVIDIA 的并行计算平台。PyTorch 和 TensorRT 都依赖它。

3.4.1 下载与安装

去 NVIDIA 开发者网站下载 CUDA 12.1 runfile（local 版本）。

wget https://developer.download.nvidia.com/compute/cuda/12.1.0/local_installers/cuda_12.1.0_530.30.02_linux.run
sudo sh cuda_12.1.0_530.30.02_linux.run

安装时注意：不要勾选 Driver（驱动我们已经装好了），只选 CUDA Toolkit 和 CUDA Samples。

3.4.2 配置环境变量

echo 'export PATH=/usr/local/cuda-12.1/bin:$PATH' >> ~/.bashrc
echo 'export LD_LIBRARY_PATH=/usr/local/cuda-12.1/lib64:$LD_LIBRARY_PATH' >> ~/.bashrc
source ~/.bashrc

3.4.3 验证

nvcc --version

输出应显示 release 12.1。如果报错，检查 PATH 是否配置正确。

3.5 第四步：cuDNN 8.9 安装

cuDNN 是深度神经网络的加速库。TensorRT 需要它做卷积优化。

  1. 去 NVIDIA 官网下载 cuDNN 8.9 for CUDA 12.x（需要注册账号）。
  2. 下载 Local Installer for Linux x86_64 (Tar)。
tar -xvf cudnn-linux-x86_64-8.9.7.29_cuda12-archive.tar.xz
sudo cp cudnn-*-archive/include/cudnn*.h /usr/local/cuda-12.1/include/
sudo cp -P cudnn-*-archive/lib/libcudnn* /usr/local/cuda-12.1/lib64/
sudo chmod a+r /usr/local/cuda-12.1/include/cudnn*.h /usr/local/cuda-12.1/lib64/libcudnn*
验证 cuDNN 版本： cat /usr/local/cuda-12.1/include/cudnn_version.h | grep CUDNN_MAJOR -A 2。输出应该显示 8.9.x。

3.6 第五步：PyTorch 2.x 安装

PyTorch 2.x 引入了 torch.compile，训练速度能提升 30%-50%。

3.6.1 创建虚拟环境

我习惯用 conda 管理环境，干净且好切换。

wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
bash Miniconda3-latest-Linux-x86_64.sh
conda create -n yolov8 python=3.10 -y
conda activate yolov8

3.6.2 安装 PyTorch

去 PyTorch 官网选择对应配置。CUDA 12.1 对应命令：

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

3.6.3 验证

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

输出应为 2.x.x 和 True。如果 False，检查 CUDA 和驱动。

注意： 不要用 pip install torch 默认安装，它会装 CPU 版本。一定要指定 --index-url 为 cu121 版本。

3.7 第六步：TensorRT 8.6 安装

TensorRT 是 NVIDIA 的推理优化引擎。YOLO 模型转成 TensorRT 后，推理速度能快 2-3 倍。

3.7.1 下载 TensorRT

去 NVIDIA 官网下载 TensorRT 8.6 GA for Linux x86_64 and CUDA 12.0 TAR package。

tar -xvf TensorRT-8.6.1.6.Linux.x86_64-gnu.cuda-12.0.tar.gz
sudo mv TensorRT-8.6.1.6 /usr/local/
echo 'export LD_LIBRARY_PATH=/usr/local/TensorRT-8.6.1.6/lib:$LD_LIBRARY_PATH' >> ~/.bashrc
source ~/.bashrc

3.7.2 安装 Python 包

cd /usr/local/TensorRT-8.6.1.6/python
pip install tensorrt-8.6.1-cp310-none-linux_x86_64.whl

3.7.3 验证

python -c "import tensorrt as trt; print(trt.__version__)"

输出应为 8.6.1。

我的经验： TensorRT 的 Python 包版本必须和系统架构匹配。如果你用 Python 3.11，就下载 cp311 的 whl 文件。不匹配会报 ImportError: No module named tensorrt。

3.8 第七步：安装 YOLOv8 与验证全流程

环境搭好了，跑个 YOLOv8 验证一下。

pip install ultralytics
yolo predict model=yolov8n.pt source='https://ultralytics.com/images/bus.jpg'

如果一切正常，你会看到检测结果图片。这说明：

  • PyTorch 能调用 GPU
  • YOLO 模型能正常推理
  • 环境链路全部打通

3.9 常见问题与避坑指南

问题 原因 解决方案
nvidia-smi 报错 "No devices were found" 驱动未正确安装或显卡被禁用 检查 BIOS 设置，重新安装驱动
torch.cuda.is_available() 返回 False PyTorch 版本与 CUDA 不匹配 卸载 PyTorch，用 --index-url 重新安装
TensorRT 导入报错 libnvinfer.so 找不到 LD_LIBRARY_PATH 未配置 检查环境变量，确保指向 TensorRT 的 lib 目录
YOLO 训练时显存不足 batch size 太大或模型太大 减小 batch size，或使用 --device 0 指定单卡
总结一下： 环境搭建就像盖房子，地基（驱动和 CUDA）必须稳。每一步装完都验证一下，别攒到最后才发现问题。我当年第一次搭环境时，跳过了 cuDNN 的权限设置，结果 TensorRT 编译模型时一直报错，查了两天才找到原因。嗯，这些坑你就不用再踩了。
← 上一章 📖 返回目录 下一章 →