📚 课程目录第2章 PyTorch模型基础

2. PyTorch模型基础：环境搭建与YOLOv8初探

说实话，很多同学一上来就急着跑模型，结果环境配了两天还没跑通。我当年也踩过这个坑——记得第一次配深度学习环境，CUDA版本不对，折腾了一整晚。所以这一章，咱们先把地基打牢。

2.1 PyTorch环境搭建

我个人习惯用conda来管理环境，干净、隔离、好清理。你想想看，项目多了，每个依赖版本都不一样，用conda能省不少心。

推荐环境配置：

Python 3.10+
PyTorch 2.0+（CUDA 11.8或12.1）
CUDA Toolkit + cuDNN（版本匹配）
Ultralytics（YOLOv8官方库）

搭建步骤其实就三步：

创建虚拟环境：conda create -n yolov8 python=3.10 -y
安装PyTorch：去官网选对应CUDA版本，比如 pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
安装Ultralytics：pip install ultralytics

避坑指南：我曾经因为CUDA版本和显卡驱动不匹配，训练时直接报"CUDA error: out of memory"。后来发现是驱动太老，不支持新CUDA。建议先跑 nvidia-smi 看驱动版本，再选对应的CUDA。

2.2 YOLOv8源码下载与结构

YOLOv8的源码托管在GitHub上，Ultralytics团队维护得很勤快。下载方式有两种：

pip安装：pip install ultralytics，直接当库用
git clone：git clone https://github.com/ultralytics/ultralytics.git，方便改源码

我个人更推荐第二种，因为后面做ONNX导出时，可能需要调整一些细节。

YOLOv8的核心结构，说白了就是三部分：

模块	作用	关键组件
Backbone	特征提取	CSPDarknet + Conv + C2f
Neck	特征融合	FPN + PAN结构
Head	检测输出	Decoupled Head + Loss

嗯，这里有个关键点——YOLOv8把之前的耦合头换成了解耦头（Decoupled Head）。分类和回归分支分开，收敛更快。我在项目里对比过，确实比v5好调。

2.3 模型加载与推理测试

环境搭好了，源码也下了，咱们先跑个推理试试。说白了，就是验证一下整个链路通不通。

小技巧：第一次跑推理，建议用官方预训练权重，别自己从头训。我刚开始总想自己训，结果显卡烧了三天，效果还不如人家预训练的。

加载模型就两行代码：

from ultralytics import YOLO

# 加载预训练模型
model = YOLO('yolov8n.pt')  # n: nano, s: small, m: medium, l: large, x: xlarge

推理更简单：

# 对图片进行推理
results = model('bus.jpg')

# 查看结果
for r in results:
    print(r.boxes)  # 边界框
    print(r.names)  # 类别名称

嗯，这里要注意——results 返回的是一个列表，每个元素对应一张图的结果。如果你传了一张图，那就是一个元素。

2.4 理解模型输入输出

做部署，最怕的就是输入输出搞不清楚。我见过有人把输入尺寸搞错，结果ONNX导出来推理全是错的。

输入格式：

形状：[1, 3, 640, 640]（batch, channel, height, width）
数据类型：float32
归一化：像素值除以255，范围[0,1]
颜色通道：RGB（注意OpenCV默认是BGR！）

输出格式：

形状：[1, 84, 8400]（以YOLOv8n为例）
84 = 4个坐标 + 80个类别概率
8400 = 所有检测框（不同尺度特征图加起来）

为什么是8400？

YOLOv8在三个尺度上做检测：80×80、40×40、20×20。加起来就是 6400 + 1600 + 400 = 8400。每个格子预测一个框，所以总共8400个候选框。

输出后处理一般要经过：

阈值过滤：去掉置信度低的框（比如 < 0.5）
NMS：非极大值抑制，去掉重复框
坐标转换：从归一化坐标转回原图坐标

我建议你写个简单的后处理函数，把结果可视化出来，这样心里有数。代码大概长这样：

import cv2
import numpy as np

def postprocess(output, img_shape, conf_thres=0.5):
    # output: [1, 84, 8400]
    # 转置为 [8400, 84]
    pred = output[0].transpose(1, 0)
    
    # 过滤低置信度
    mask = pred[:, 4] > conf_thres
    pred = pred[mask]
    
    # 取类别
    class_ids = np.argmax(pred[:, 5:], axis=1)
    confidences = pred[:, 4]
    boxes = pred[:, :4]
    
    return boxes, class_ids, confidences

好了，这一章的内容就到这儿。环境搭好了，模型能跑了，输入输出也搞清楚了——后面做ONNX导出和TensorRT部署，心里就有底了。