📚 课程目录第2章环境搭建与基准模型训练

2. 环境搭建与基准模型训练：PyTorch/CUDA环境配置、YOLOv11官方代码库下载与安装、使用官方脚本训练一个基线YOLOv11模型

说实话，做模型压缩和蒸馏，最怕什么？

最怕环境没搭好，折腾半天发现是CUDA版本不对，或者PyTorch和显卡驱动打架。我见过太多同学卡在这一步，一卡就是一整天。所以这一章，咱们把地基打牢。

2.1 硬件与软件环境概览

先看看我们都需要什么。我个人习惯把环境分成三层：硬件层、驱动层、框架层。

层级	组件	推荐版本/型号	备注
硬件层	GPU	NVIDIA RTX 3090 / 4090 / A100	显存至少8GB，推荐24GB
驱动层	NVIDIA Driver	>= 525.60.13	用 `nvidia-smi` 查看
驱动层	CUDA Toolkit	11.8 或 12.1	与PyTorch版本匹配
框架层	PyTorch	2.0.1 / 2.1.0	CUDA 11.8 或 12.1 版
框架层	Python	3.9 - 3.11	推荐3.10

⚠️ 我曾经踩过的坑： 有一次我直接在已有环境上装YOLOv11，结果torchvision和CUDA版本不匹配，训练时直接报 RuntimeError: CUDA error: no kernel image is available。后来全部重装，老老实实按版本表来，一次过。

2.2 安装CUDA与cuDNN

这一步其实很简单，但很多人会搞混。你想想看，CUDA Toolkit和显卡驱动是两个东西。驱动是让系统认识显卡，CUDA Toolkit是给PyTorch用的计算库。

检查驱动：终端输入 nvidia-smi，看右上角的CUDA Version。比如显示 CUDA Version: 12.1，说明驱动支持CUDA 12.1。
下载CUDA Toolkit：去NVIDIA官网，选对应版本。我一般用runfile安装，干净。
安装cuDNN：下载后解压，把文件复制到CUDA目录下。

# 检查驱动
nvidia-smi

# 安装CUDA 12.1（以Ubuntu为例）
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

# 配置环境变量
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

# 验证
nvcc --version

💡 小技巧： 如果你有多张显卡，或者想在同一台机器上装多个CUDA版本，用 update-alternatives 管理路径，切换起来很方便。

2.3 创建Python虚拟环境

我个人强烈建议用虚拟环境。为什么？因为YOLOv11的依赖和你的其他项目可能打架。我用conda，干净利落。

# 创建环境
conda create -n yolov11 python=3.10 -y

# 激活环境
conda activate yolov11

# 安装PyTorch（以CUDA 12.1为例）
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

# 验证PyTorch是否识别GPU
python -c "import torch; print(torch.cuda.is_available()); print(torch.cuda.get_device_name(0))"

如果输出 True 和你的显卡型号，恭喜，环境搭好了。

2.4 下载YOLOv11官方代码库

YOLOv11的官方代码在GitHub上。我建议直接clone，不要下载zip，因为后续更新方便。

# 克隆仓库
git clone https://github.com/ultralytics/ultralytics.git

# 进入目录
cd ultralytics

# 安装依赖
pip install -r requirements.txt

# 安装ultralytics包（可编辑模式，方便改代码）
pip install -e .

    📌 注意： 安装完成后，运行 yolo 命令，如果能显示帮助信息，说明安装成功。我第一次装的时候忘了装 requirements.txt，结果训练时缺包，一个个补，烦死了。

2.5 准备数据集

为了训练基线模型，我们需要一个数据集。这里以COCO128为例——它是COCO2017的一个子集，只有128张图，用来快速验证流程。

# 下载COCO128数据集
yolo train data=coco128.yaml model=yolov11n.pt epochs=1

# 或者手动下载
wget https://ultralytics.com/assets/coco128.zip
unzip coco128.zip -d datasets/

数据集结构是这样的：

datasets/
└── coco128/
    ├── images/
    │   ├── train2017/
    │   └── val2017/
    ├── labels/
    │   ├── train2017/
    │   └── val2017/
    └── coco128.yaml

2.6 训练基线YOLOv11模型

好，重头戏来了。我们用官方脚本训练一个YOLOv11n模型作为基线。为什么选nano？因为它小，训练快，适合验证流程。

# 训练YOLOv11n
yolo train \
    model=yolov11n.pt \
    data=coco128.yaml \
    epochs=100 \
    imgsz=640 \
    batch=16 \
    device=0 \
    project=runs/train \
    name=baseline_yolov11n \
    exist_ok=True \
    pretrained=True \
    amp=True

参数说明：

参数	值	说明
`model`	yolov11n.pt	预训练权重，自动下载
`data`	coco128.yaml	数据集配置文件
`epochs`	100	训练轮数
`imgsz`	640	输入图片尺寸
`batch`	16	批次大小，根据显存调整
`device`	0	GPU编号
`amp`	True	混合精度训练，省显存

💡 我的经验： 第一次训练时，先用 epochs=5 跑一轮，确认数据加载、模型前向、loss计算都没问题。没问题了再开100轮。不然等半天发现报错，心态会崩。

2.7 训练过程监控与结果评估

训练过程中，终端会实时打印loss、mAP等指标。我个人习惯同时开两个窗口：一个跑训练，一个用 tail -f 看日志。

# 实时查看训练日志
tail -f runs/train/baseline_yolov11n/train.log

训练完成后，结果保存在 runs/train/baseline_yolov11n/ 目录下：

weights/best.pt：验证集上mAP最高的权重
weights/last.pt：最后一轮的权重
results.png：训练曲线图
confusion_matrix.png：混淆矩阵

用以下命令验证模型效果：

# 验证
yolo val model=runs/train/baseline_yolov11n/weights/best.pt data=coco128.yaml

# 推理测试
yolo predict model=runs/train/baseline_yolov11n/weights/best.pt source=datasets/coco128/images/val2017/

2.8 本章知识体系

下面这张图把整个流程串起来了，你可以保存下来当参考。

嗯，到这里，你的环境就搭好了，基线模型也跑通了。接下来，我们就可以拿着这个模型，开始做知识蒸馏和模型压缩了。