4. 训练环境搭建：Docker容器化、PyTorch/TensorRT环境配置、数据集准备与标注格式转换

环境搭建这事儿，看着简单，其实坑不少。我见过太多团队，算法调得挺好，一到部署就翻车。说白了，就是环境不一致闹的。

这一章，咱们就把工业缺陷检测的环境搭建捋清楚。从Docker容器化开始，到PyTorch训练环境，再到TensorRT推理优化，最后聊聊数据集那点事儿。

4.1 Docker容器化：让环境“随身携带”

为什么要用Docker？你想想看，算法工程师的电脑是Ubuntu 20.04，现场工控机是CentOS 7，库版本对不上，驱动装不上，这种事我遇到过不下十次。

Docker的好处就是：一次构建，到处运行。把整个环境打包成镜像，到哪都能跑。

4.1.1 Dockerfile编写要点

我习惯用NVIDIA的官方镜像做基础层。别自己从零装CUDA，太容易出问题。

# 基础镜像：Ubuntu 20.04 + CUDA 11.8
FROM nvidia/cuda:11.8.0-cudnn8-devel-ubuntu20.04

# 设置非交互模式
ENV DEBIAN_FRONTEND=noninteractive

# 安装系统依赖
RUN apt-get update && apt-get install -y \
    python3.8 \
    python3-pip \
    libgl1-mesa-glx \
    libglib2.0-0 \
    libsm6 \
    libxext6 \
    libxrender-dev \
    libgomp1 \
    && rm -rf /var/lib/apt/lists/*

# 设置工作目录
WORKDIR /workspace

# 安装Python依赖
COPY requirements.txt .
RUN pip3 install --no-cache-dir -r requirements.txt

# 设置时区
ENV TZ=Asia/Shanghai
RUN ln -snf /usr/share/zoneinfo/$TZ /etc/localtime
我的习惯：requirements.txt里只写核心依赖，版本号锁死。比如torch==1.13.1，别用>=，否则不同时间构建的镜像可能不一样。

4.1.2 构建与运行命令

# 构建镜像
docker build -t defect-detection:v1.0 .

# 运行容器（挂载数据卷和GPU）
docker run -it --gpus all \
    -v /home/user/data:/workspace/data \
    -v /home/user/models:/workspace/models \
    --shm-size=16g \
    defect-detection:v1.0
注意：--shm-size一定要设。PyTorch DataLoader多进程读取数据时，共享内存不够会报错。我一开始没设，排查了半天。

4.2 PyTorch环境配置：选对版本很重要

PyTorch版本和CUDA版本要匹配。这个表我经常翻，直接贴出来：

PyTorch版本 CUDA版本 cuDNN版本 适用场景
1.10.x 11.3 8.2 老项目兼容
1.13.x 11.7 / 11.8 8.5 工业部署主流
2.0.x 11.8 / 12.1 8.7 新项目推荐

我个人建议：工业项目用PyTorch 1.13 + CUDA 11.8。这个组合最稳，TensorRT支持也最好。

4.2.1 安装命令

# PyTorch 1.13.1 + CUDA 11.8
pip3 install torch==1.13.1+cu118 torchvision==0.14.1+cu118 \
    --extra-index-url https://download.pytorch.org/whl/cu118
验证安装： 
python3 -c "import torch; print(torch.cuda.is_available())"
# 输出 True 才算成功

4.3 TensorRT环境配置：推理加速的关键

TensorRT能把模型推理速度提升2-5倍。但配置起来有点麻烦。

4.3.1 TensorRT版本选择

记住一个原则：TensorRT版本必须和CUDA版本对应。我用的是TensorRT 8.5.3 + CUDA 11.8。

# 下载TensorRT（需要NVIDIA账号）
# 选择：TensorRT 8.5.3 for Ubuntu 20.04 and CUDA 11.8

# 安装
dpkg -i nv-tensorrt-repo-ubuntu2004-cuda11.8-trt8.5.3.0-ga-20230314_1-1_amd64.deb
apt-get update
apt-get install tensorrt

# 验证
python3 -c "import tensorrt; print(tensorrt.__version__)"
避坑指南：我曾经在Docker里装TensorRT，忘了装pycuda，结果跑的时候报错。记得一起装上： 
pip3 install pycuda

4.3.2 ONNX导出与TensorRT转换

这是模型移植的核心流程：

  1. PyTorch模型 → ONNX：固定输入尺寸，设置动态轴
  2. ONNX → TensorRT：选择精度（FP32/FP16/INT8）
  3. TensorRT推理：序列化引擎，保存为.engine文件
# PyTorch转ONNX
import torch
import torch.onnx

model = torch.load('model.pth')
dummy_input = torch.randn(1, 3, 640, 640)

torch.onnx.export(
    model,
    dummy_input,
    'model.onnx',
    opset_version=11,
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={'input': {0: 'batch_size'}}
)

# ONNX转TensorRT（命令行）
trtexec --onnx=model.onnx \
        --saveEngine=model.engine \
        --fp16 \
        --workspace=4096
注意：转INT8需要校准数据集，不然精度掉得厉害。我有个项目INT8后mAP掉了5个点，后来改FP16才稳住。

4.4 数据集准备与标注格式转换

工业缺陷检测的数据集，格式五花八门。VOC、COCO、YOLO、自定义格式... 转换起来很烦。

4.4.1 常见标注格式对比

格式 存储方式 标注内容 适用场景
VOC XML文件 边界框 + 类别 分类、检测
COCO JSON文件 边界框 + 分割掩码 检测、分割
YOLO TXT文件 归一化坐标 实时检测
自定义 CSV/TXT 按项目定义 特殊需求

4.4.2 格式转换脚本（VOC转YOLO）

这是我项目里常用的转换脚本：

import xml.etree.ElementTree as ET
import os

def voc_to_yolo(xml_file, output_dir, classes):
    tree = ET.parse(xml_file)
    root = tree.getroot()
    
    # 获取图像尺寸
    size = root.find('size')
    width = int(size.find('width').text)
    height = int(size.find('height').text)
    
    # 解析每个目标
    yolo_lines = []
    for obj in root.iter('object'):
        class_name = obj.find('name').text
        if class_name not in classes:
            continue
        
        class_id = classes.index(class_name)
        bbox = obj.find('bndbox')
        
        # VOC格式：xmin, ymin, xmax, ymax
        xmin = int(bbox.find('xmin').text)
        ymin = int(bbox.find('ymin').text)
        xmax = int(bbox.find('xmax').text)
        ymax = int(bbox.find('ymax').text)
        
        # 转YOLO格式：中心点 + 宽高（归一化）
        x_center = (xmin + xmax) / 2.0 / width
        y_center = (ymin + ymax) / 2.0 / height
        w = (xmax - xmin) / width
        h = (ymax - ymin) / height
        
        yolo_lines.append(f"{class_id} {x_center:.6f} {y_center:.6f} {w:.6f} {h:.6f}")
    
    # 写入文件
    txt_name = os.path.splitext(os.path.basename(xml_file))[0] + '.txt'
    with open(os.path.join(output_dir, txt_name), 'w') as f:
        f.write('\n'.join(yolo_lines))

# 使用示例
classes = ['scratch', 'dent', 'crack']  # 缺陷类别
voc_to_yolo('annotations/001.xml', 'labels/', classes)
我的经验：转换后一定要可视化检查。写个脚本把YOLO标签画回图片上，看一眼就知道对不对。我吃过亏，坐标转换公式写反了，训练了三天才发现。

4.4.3 数据集划分

工业场景下，我建议按产线批次划分，而不是随机划分。这样可以避免同一批次的数据既在训练集又在验证集。

# 按批次划分
train_batches = ['batch_01', 'batch_02', 'batch_03']
val_batches = ['batch_04']
test_batches = ['batch_05']

# 生成文件列表
train_images = []
for batch in train_batches:
    train_images.extend(glob.glob(f'data/{batch}/*.jpg'))
关键指标：
  • 训练集：验证集：测试集 = 7:2:1
  • 每个缺陷类别至少100张
  • 正负样本比例不超过1:10

4.5 本章知识结构

下面这张图，把环境搭建的整个流程串起来了：

工业缺陷检测环境搭建流程 Docker容器化 基础镜像 + 依赖安装 PyTorch环境 CUDA + cuDNN + Torch TensorRT环境 ONNX导出 + 引擎转换 数据集准备 标注格式转换 → 数据集划分 → 数据增强 模型移植流程 PyTorch训练 → ONNX导出 → TensorRT优化 → 部署 关键要点： • 版本匹配：PyTorch ↔ CUDA ↔ TensorRT 版本必须对应 • 数据验证：转换后一定要可视化检查，避免坐标错误

4.6 常见问题与解决方案

问题 原因 解决方案
Docker内无法调用GPU 未安装nvidia-docker 安装nvidia-docker2并重启
PyTorch报CUDA版本错误 PyTorch与CUDA不匹配 卸载重装对应版本
TensorRT转换失败 ONNX算子不支持 简化模型结构，替换自定义算子
标注坐标偏移 格式转换公式错误 可视化检查，修正公式
最后提醒一句：环境搭建别图快。每装一个组件，都验证一下。我见过太多人一口气装完，跑不起来，又从头排查，反而更慢。
← 上一章 📖 返回目录 下一章 →