2. 环境搭建

说实话,工业质检的部署环境,比训练环境要「脏」得多。训练可以在云上随便折腾,但到了产线工控机,系统版本、驱动、库依赖,任何一个环节出问题,模型就跑不起来。这一章,我把这几年踩过的坑和沉淀下来的标准流程,一次性讲清楚。

2.1 Ubuntu系统基础配置

我个人习惯用Ubuntu 20.04 LTS或22.04 LTS。为什么?因为NVIDIA的驱动和CUDA对这两个版本支持最稳。我曾经在18.04上折腾过RTX 4090的驱动,那叫一个痛苦。

2.1.1 系统安装后的「三板斧」

装完系统第一件事,不是装CUDA,而是先把系统源换掉,再装基础工具链。你想想看,如果源在国外,下载一个驱动包等半小时,产线那边等着验收,心态直接崩。

# 1. 备份并更换apt源为清华源
sudo cp /etc/apt/sources.list /etc/apt/sources.list.bak
sudo sed -i 's/archive.ubuntu.com/mirrors.tuna.tsinghua.edu.cn/g' /etc/apt/sources.list
sudo sed -i 's/security.ubuntu.com/mirrors.tuna.tsinghua.edu.cn/g' /etc/apt/sources.list

# 2. 更新系统并安装基础工具
sudo apt update && sudo apt upgrade -y
sudo apt install -y build-essential dkms git curl wget vim net-tools

# 3. 禁用系统自带的nouveau驱动(NVIDIA驱动安装的前置条件)
sudo bash -c "echo 'blacklist nouveau' >> /etc/modprobe.d/blacklist-nvidia-nouveau.conf"
sudo bash -c "echo 'options nouveau modeset=0' >> /etc/modprobe.d/blacklist-nvidia-nouveau.conf"
sudo update-initramfs -u
注意:禁用nouveau后需要重启系统。重启后运行 lsmod | grep nouveau,如果没有输出,说明禁用成功。这一步漏掉,NVIDIA驱动安装时会直接报错。

2.2 CUDA与cuDNN安装

工业场景下,我建议不要追最新版CUDA。最新版往往意味着兼容性问题。YOLOv8官方推荐CUDA 11.8或12.1,我实测下来,11.8最稳,几乎所有工控机都能跑。

2.2.1 安装NVIDIA驱动

驱动版本和CUDA版本有对应关系。我习惯用runfile安装,因为apt安装的驱动版本往往太老。

# 查看推荐驱动版本
ubuntu-drivers devices

# 安装推荐驱动(比如nvidia-driver-535)
sudo apt install -y nvidia-driver-535

# 或者从NVIDIA官网下载runfile安装(推荐)
# 先停止图形界面服务
sudo service gdm3 stop
# 赋予执行权限并安装
chmod +x NVIDIA-Linux-x86_64-535.154.05.run
sudo ./NVIDIA-Linux-x86_64-535.154.05.run --no-opengl-files
我的经验:安装完成后运行 nvidia-smi,如果能看到GPU信息,驱动就装好了。注意看右上角的CUDA Version,那个是驱动支持的最高CUDA版本,不是实际安装的CUDA版本。

2.2.2 安装CUDA Toolkit

去NVIDIA官网下载CUDA 11.8的runfile。我不用deb包,因为runfile可以自定义安装路径,方便多版本共存。

# 下载CUDA 11.8
wget https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda_11.8.0_520.61.05_linux.run

# 安装(注意不要安装驱动,只安装Toolkit)
sudo sh cuda_11.8.0_520.61.05_linux.run --toolkit --silent --override

# 配置环境变量
echo 'export PATH=/usr/local/cuda-11.8/bin:$PATH' >> ~/.bashrc
echo 'export LD_LIBRARY_PATH=/usr/local/cuda-11.8/lib64:$LD_LIBRARY_PATH' >> ~/.bashrc
source ~/.bashrc

# 验证安装
nvcc --version

2.2.3 安装cuDNN

cuDNN是深度学习的加速库,YOLO的卷积操作依赖它。需要去NVIDIA官网注册下载,这里假设你已经下载了cudnn-linux-x86_64-8.9.7.29_cuda11-archive.tar.xz。

# 解压并复制到CUDA目录
tar -xvf cudnn-linux-x86_64-8.9.7.29_cuda11-archive.tar.xz
sudo cp cudnn-linux-x86_64-8.9.7.29_cuda11-archive/include/cudnn*.h /usr/local/cuda-11.8/include/
sudo cp cudnn-linux-x86_64-8.9.7.29_cuda11-archive/lib/libcudnn* /usr/local/cuda-11.8/lib64/
sudo chmod a+r /usr/local/cuda-11.8/include/cudnn*.h /usr/local/cuda-11.8/lib64/libcudnn*

# 验证cuDNN版本
cat /usr/local/cuda-11.8/include/cudnn_version.h | grep CUDNN_MAJOR -A 2
避坑指南:我曾经在cuDNN版本上吃过亏。YOLOv8要求cuDNN >= 8.5,但8.9版本在某些RTX 30系列卡上有兼容问题。如果训练时出现 cudnn_status_not_initialized 错误,试试降级到8.7。

2.3 Docker部署环境

说实话,Docker是工业部署的「救命稻草」。产线上的工控机环境千奇百怪,有了Docker,你只需要保证宿主机有NVIDIA驱动和Docker,剩下的全部交给镜像。

2.3.1 安装Docker与NVIDIA Container Toolkit

# 安装Docker
curl -fsSL https://mirrors.tuna.tsinghua.edu.cn/docker-ce/linux/ubuntu/gpg | sudo apt-key add -
sudo add-apt-repository "deb [arch=amd64] https://mirrors.tuna.tsinghua.edu.cn/docker-ce/linux/ubuntu $(lsb_release -cs) stable"
sudo apt update
sudo apt install -y docker-ce docker-ce-cli containerd.io

# 安装NVIDIA Container Toolkit
distribution=$(. /etc/os-release;echo $ID$VERSION_ID)
curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add -
curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list
sudo apt update
sudo apt install -y nvidia-container-toolkit
sudo systemctl restart docker

# 验证Docker能否调用GPU
sudo docker run --rm --gpus all nvidia/cuda:11.8.0-base-ubuntu20.04 nvidia-smi
小技巧:把当前用户加入docker组,避免每次运行docker都要加sudo:sudo usermod -aG docker $USER,然后重新登录。

2.3.2 构建YOLO部署镜像

我习惯基于nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu20.04构建镜像,这个镜像已经包含了CUDA和cuDNN,省去手动安装的麻烦。

# Dockerfile
FROM nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu20.04

# 换源并安装Python和基础库
RUN sed -i 's/archive.ubuntu.com/mirrors.tuna.tsinghua.edu.cn/g' /etc/apt/sources.list && \
    apt update && apt install -y python3 python3-pip libgl1-mesa-glx libglib2.0-0 && \
    pip3 install --upgrade pip -i https://pypi.tuna.tsinghua.edu.cn/simple

# 安装PyTorch和YOLO依赖
RUN pip3 install torch==1.13.1+cu117 torchvision==0.14.1+cu117 --extra-index-url https://download.pytorch.org/whl/cu117 -i https://pypi.tuna.tsinghua.edu.cn/simple && \
    pip3 install ultralytics opencv-python-headless -i https://pypi.tuna.tsinghua.edu.cn/simple

WORKDIR /workspace
COPY . /workspace

CMD ["python3", "deploy.py"]

2.4 PyTorch环境配置

PyTorch的安装其实是最简单的,但版本选择有讲究。工业部署我推荐PyTorch 1.13.1或2.0.1,这两个版本对ONNX导出支持最好。

2.4.1 安装PyTorch

# 创建虚拟环境(推荐)
python3 -m venv yolov8_env
source yolov8_env/bin/activate

# 安装PyTorch 1.13.1(对应CUDA 11.7/11.8)
pip install torch==1.13.1+cu117 torchvision==0.14.1+cu117 --extra-index-url https://download.pytorch.org/whl/cu117

# 验证PyTorch能否调用GPU
python -c "import torch; print(torch.cuda.is_available()); print(torch.cuda.get_device_name(0))"
注意:如果输出 False,说明PyTorch没有检测到CUDA。检查一下是不是装了CPU版本的PyTorch。运行 pip list | grep torch,看看版本号后面有没有 +cu117 或 +cu118。

2.4.2 安装YOLOv8与验证

# 安装ultralytics库
pip install ultralytics

# 下载预训练模型并测试推理
yolo predict model=yolov8n.pt source='https://ultralytics.com/images/bus.jpg'

# 检查推理结果
ls -la runs/detect/predict/

2.5 环境搭建知识体系

下面这张图,是我总结的环境搭建全流程。每次在新工控机上部署,我就照着这个流程走一遍,基本不会出问题。

环境搭建全流程 Ubuntu 20.04/22.04 NVIDIA驱动 535+ 基础工具链 CUDA 11.8 Toolkit cuDNN 8.7/8.9 Docker + NVIDIA Container Toolkit Python 3.8+ 虚拟环境 PyTorch 1.13.1/2.0.1 Ultralytics YOLOv8 建议:每层安装完成后都做一次验证,不要等到最后才发现问题

2.6 环境验证清单

每次搭建完环境,我都会跑一遍下面的验证清单。全部通过,才敢把工控机交给产线。

验证项 命令 预期结果
NVIDIA驱动 nvidia-smi 显示GPU型号、驱动版本、CUDA版本
CUDA Toolkit nvcc --version 显示CUDA 11.8
cuDNN cat /usr/local/cuda/include/cudnn_version.h | grep CUDNN_MAJOR 显示8.x
Docker GPU docker run --rm --gpus all nvidia/cuda:11.8.0-base-ubuntu20.04 nvidia-smi 容器内显示GPU信息
PyTorch GPU python -c "import torch; print(torch.cuda.is_available())" True
YOLO推理 yolo predict model=yolov8n.pt source='https://ultralytics.com/images/bus.jpg' 生成推理结果图片
最后说一句:环境搭建是工业部署里最枯燥但最关键的一步。我见过太多项目,模型精度99%,结果因为环境问题在产线跑不起来。按照这个流程走,至少能帮你省掉80%的排查时间。