文章目录
前言
大家好！欢迎来到“从代码学习深度学习”系列博客。目标检测是计算机视觉领域的核心任务之一，旨在识别图像或视频中特定类别的对象实例，并确定它们的位置和范围。近年来，深度学习技术极大地推动了目标检测的发展，涌现出许多优秀的算法，如 R-CNN 系列、YOLO 系列以及我们今天要重点介绍的单发多框检测（Single Shot MultiBox Detector, SSD）。
SSD 是一种流行的单阶段目标检测器，以其在速度和精度之间的良好平衡而闻名。与两阶段检测器（如 Faster R-CNN）先生成区域提议再进行分类和回归不同，SSD 直接在不同尺度的特征图上预测边界框和类别，从而实现了更快的检测速度。
本篇博客旨在通过一个具体的 PyTorch 实现（基于香蕉检测数据集），带领大家深入理解 SSD 的核心原理和代码实现细节。我们将逐步剖析模型结构、损失函数、训练过程以及预测可视化等关键环节，真正做到“从代码中学习”。
完整代码：下载链接
在深入 SSD 模型之前，我们先引入一些在整个项目中会用到的工具函数，它们主要负责数据处理、模型训练辅助以及结果可视化。
工具函数
在实现和训练 SSD 模型以及可视化结果的过程中，我们会用到一些辅助函数。这些函数分散在不同的工具文件中。
数据处理工具 (utils_for_data.py)
这部分代码负责读取和加载香蕉检测数据集。read_data_bananas 函数读取图像和对应的 CSV 标签文件，并将它们转换成 PyTorch 张量。BananasDataset 类继承了 torch.utils.data.Dataset，方便我们构建数据加载器。load_data_bananas 函数则利用 BananasDataset 创建了训练和验证数据的数据加载器（DataLoader）。
# --- START OF FILE utils_for_data.py ---
import os
import pandas as pd
import torch
import torchvision
def read_data_bananas(is_train=True):
"""
读取香蕉检测数据集中的图像和标签
参数:
is_train (bool): 是否读取训练集数据，True表示读取训练集，False表示读取验证集
返回:
tuple: (images, targets)
- images: 图像列表，每个元素是一个形状为[C, H, W]的张量
- targets: 标注信息张量，形状为[N, 1, 5]，每行包含[类别, 左上角x, 左上角y, 右下角x, 右下角y]
"""
# 设置数据目录路径
data_dir = 'banana-detection'
# 根据is_train确定使用训练集还是验证集路径
subset_name = 'bananas_train' if is_train else 'bananas_val'
# 构建标签CSV文件的完整路径
# csv_fname: 字符串，表示CSV文件的完整路径
csv_fname = os.path.join(data_dir, subset_name, 'label.csv')
# 读取CSV文件到pandas DataFrame
# csv_data: DataFrame，包含图像名称和对应的标注信息
csv_data = pd.read_csv(csv_fname)
# 将img_name列设置为索引，便于后续访问
# csv_data: DataFrame，索引为图像名称，列为标注信息
csv_data = csv_data.set_index('img_name')
# 初始化存储图像和标注的列表
# images: 列表，用于存储读取的图像张量
# targets: 列表，用于存储对应的标注信息
images, targets = [], []
# 遍历DataFrame中的每一行，读取图像和对应的标注信息
for img_name, target in csv_data.iterrows():
# 读取图像并添加到images列表中
# img_name: 字符串，图像文件名
# 读取的图像: 张量，形状为[C, H, W]，C是通道数，H是高度，W是宽度
images.append(torchvision.io.read_image(
os.path.join(data_dir, subset_name, 'images', f'{
img_name}')))
# 添加标注信息到targets列表中
# target: Series，包含类别和边界框坐标信息
# list(target): 列表，形状为[5]，包含[类别, 左上角x, 左上角y, 右下角x, 右下角y]
targets.append(list(target))
# 将targets列表转换为张量，并添加一个维度，然后将值归一化到0-1范围
# torch.tensor(targets): 张量，形状为[N, 5]，N是样本数量
# torch.tensor(targets).unsqueeze(1): 张量，形状为[N, 1, 5]
# 最终返回的targets: 张量，形状为[N, 1, 5]，值范围在0-1之间
targets_tensor = torch.tensor(targets).unsqueeze(1) / 256
return images, targets_tensor
class BananasDataset(torch.utils.data.Dataset):
"""
一个用于加载香蕉检测数据集的自定义数据集类
继承自torch.utils.data.Dataset基类，实现了必要的__init__、__getitem__和__len__方法
用于提供数据加载器(DataLoader)访问数据集的接口
"""
def __init__(self, is_train):
"""
初始化香蕉检测数据集
参数:
is_train (bool): 是否加载训练集数据，True表示加载训练集，False表示加载验证集
属性:
self.features: 列表，包含所有图像张量，每个张量形状为[C, H, W]
self.labels: 张量，形状为[N, 1, 5]，其中N是样本数量，1是类别数量
每个样本包含[类别, 左上角x, 左上角y, 右下角x, 右下角y]
"""
# 调用read_data_bananas函数读取数据集
# self.features: 列表，包含N个形状为[C, H, W]的图像张量
# self.labels: 张量，形状为[N, 1, 5]
self.features, self.labels = read_data_bananas(is_train)
# 打印读取的数据集信息
dataset_type = '训练样本' if is_train else '验证样本'
print(f'读取了 {
len(self.features)} 个{
dataset_type}')
def __getitem__(self, idx):
"""
获取指定索引的样本
参数:
idx (int): 样本索引
返回:
tuple: (feature, label)
- feature: 张量，形状为[C, H, W]，图像数据，已转换为float类型
- label: 张量，形状为[1, 5]，对应的标注信息
"""
# 返回索引为idx的特征和标签对
# self.features[idx]: 张量，形状为[C, H, W]
# self.features[idx].float(): 将图像张量转换为float类型，形状不变，仍为[C, H, W]
# self.labels[idx]: 张量，形状为[1, 5]，包含一个目标的类别和边界框信息
return (self.features[idx].float(), self.labels[idx])
def __len__(self):
"""
获取数据集中样本的数量
返回:
int: 数据集中的样本数量
"""
# 返回数据集中的样本数量
# len(self.features): int，表示数据集中图像的总数
return len(self.features)
def load_data_bananas(batch_size):
"""
加载香蕉检测数据集，并创建数据加载器
参数:
batch_size (int): 批量大小，指定每次加载的样本数量
返回:
tuple: (train_iter, val_iter)
- train_iter: 训练数据加载器，每次返回batch_size个训练样本
每个批次包含:
- 特征张量，形状为[batch_size, C, H, W]
- 标签张量，形状为[batch_size, 1, 5]
- val_iter: 验证数据加载器，每次返回batch_size个验证样本
批次格式与train_iter相同
"""
# 创建训练集数据加载器
# BananasDataset(is_train=True): 实例化训练集数据集对象
# batch_size: 每个批次的样本数量
# shuffle=True: 打乱数据顺序，增强模型的泛化能力
# train_iter的每个批次包含:
# - 特征张量，形状为[batch_size, C, H, W]，C是通道数，H是高度，W是宽度
# - 标签张量，形状为[batch_size, 1, 5]，每行包含[类别, 左上角x, 左上角y, 右下角x, 右下角y]
train_iter = torch.utils.data.DataLoader(
BananasDataset(is_train=True),
batch_size=batch_size,
shuffle=True
)
# 创建验证集数据加载器
# BananasDataset(is_train=False): 实例化验证集数据集对象
# batch_size: 每个批次的样本数量
# shuffle默认为False: 不打乱验证数据的顺序，保持一致性
# val_iter的每个批次包含:
# - 特征张量，形状为[batch_size, C, H, W]
# - 标签张量，形状为[batch_size, 1, 5]
val_iter = torch.utils.data.DataLoader(
BananasDataset(is_train=False),
batch_size=batch_size
)
return train_iter, val_iter
# --- END OF FILE utils_for_data.py ---
训练工具 (utils_for_train.py)
这部分包含通用的训练辅助类。Timer 类用于记录和计算代码块的执行时间。Accumulator 类则方便我们在训练过程中累加损失、准确率等多个指标。try_gpu 函数尝试获取可用的 GPU 设备，否则回退到 CPU。
# --- START OF FILE utils_for_train.py ---
import torch
import math # 导入math包，用于计算指数
from torch import nn
import time
import numpy as np # 导入numpy 用于cumsum计算
class Timer:
"""记录多次运行时间"""
def __init__(self):
"""Defined in :numref:`subsec_linear_model`"""
self.times = []
self.start()
def start(self):
"""启动计时器"""
self.tik = time.time()
def stop(self):
"""停止计时器并将时间记录在列表中"""
self.times.append(time.time() - self.tik)
return self.times[-1]
def avg(self):
"""返回平均时间"""
return sum(self.times) / len(self.times)
def sum(self):
"""返回时间总和"""
return sum(self.times)
def cumsum(self):
"""返回累计时间"""
return np.array(self.times).cumsum().tolist()
class Accumulator:
"""在 n 个变量上累加
"""
def __init__(self, n):
"""初始化 Accumulator 类
输入：
n: 需要累加的变量数量 # 输入参数：变量数量
输出：
无返回值 # 方法无显式返回值
"""
self.data = [0.0] * n # 初始化一个长度为 n 的浮点数列表，初始值为 0.0
def add(self, *args):
"""向累加器中添加多个值
输入：
*args: 可变数量的数值，用于累加 # 输入参数：可变参数，表示要累加的值
输出：
无返回值 # 方法无显式返回值
"""
self.data = [a + float(b) for a, b in zip(self.data, args)] # 将输入值累加到对应位置的数据上
def reset(self):
"""重置累加器中的所有值为 0
输入：
无 # 方法无输入参数
输出：
无返回值 # 方法无显式返回值
"""
self.data = [0.0] * len(self.data) # 重置数据列表，所有值设为 0.0
def __getitem__(self, idx):
"""获取指定索引处的值
输入：
idx: 索引值 # 输入参数：要访问的数据索引
输出：
float: 指定索引处的值 # 返回指定位置的累加值
"""
return self.data[idx] # 返回指定索引处的数据值
def try_gpu(i=0):
"""如果存在，则返回gpu(i)，否则返回cpu()
Args:
i (int, optional): GPU设备的编号，默认为0，表示尝试使用第0号GPU
Returns:
torch.device: 返回可用的设备对象，如果指定编号的GPU可用则返回GPU，否则返回CPU
"""
# 检查系统中可用的GPU数量是否大于等于i+1
if torch.cuda.device_count() >= i + 1:
# 如果条件满足，返回指定编号i的GPU设备
return torch.device(f'cuda:{
i}')
# 如果没有足够的GPU设备，返回CPU设备
return torch.device('cpu')
# --- END OF FILE utils_for_train.py ---
检测相关工具 (utils_for_detection.py)
这是 SSD 实现的核心工具集。包含了以下关键功能：
- 边界框表示转换:
box_corner_to_center和box_center_to_corner用于在 (左上角, 右下角) 和 (中心点, 宽高) 两种坐标表示法之间转换。 - 锚框生成:
multibox_prior根据输入的特征图、尺寸比例 (sizes) 和宽高比 (ratios) 生成大量的锚框。 - IoU 计算:
box_iou计算两组边界框之间的交并比 (Intersection over Union)，这是目标检测中的基本度量。 - 锚框分配:
assign_anchor_to_bbox将真实边界框 (ground truth) 分配给最匹配的锚框。 - 偏移量计算:
offset_boxes计算预测边界框相对于锚框的偏移量（中心点坐标和宽高），这是回归任务的目标。offset_inverse则根据锚框和预测的偏移量反算出预测的边界框坐标。 - 目标生成:
multibox_target是关键函数，它整合了锚框分配和偏移量计算，为每个锚框生成对应的类别标签和边界框回归目标。 - 非极大值抑制 (NMS):
nms用于在预测阶段去除高度重叠的冗余检测框，保留置信度最高的框。 - 多框检测:
multibox_detection结合类别概率预测、边界框偏移量预测、锚框以及 NMS，生成最终的检测结果。 - 可视化辅助:
bbox_to_rect将边界框转换为 Matplotlib 绘图格式，show_bboxes则用于在图像上绘制边界框和标签。
# --- START OF FILE utils_for_detection.py ---
import torch
import matplotlib.pyplot as plt
torch.set_printoptions(2) # 精简输出精度
def box_corner_to_center(boxes):
"""将边界框从（左上角，右下角）表示法转换为（中心点，宽度，高度）表示法
该函数接收以(x1, y1, x2, y2)格式表示的边界框张量，其中：
- (x1, y1)：表示边界框左上角的坐标
- (x2, y2)：表示边界框右下角的坐标
然后将其转换为(cx, cy, w, h)格式，其中：
- (cx, cy)：表示边界框中心点的坐标
- w：表示边界框的宽度
- h：表示边界框的高度
参数:
boxes (torch.Tensor): 形状为(N, 4)的张量，包含N个边界框的左上角和右下角坐标
返回:
torch.Tensor: 形状为(N, 4)的张量，包含N个边界框的中心点坐标、宽度和高度
"""
# 分别提取所有边界框的左上角和右下角坐标
x1, y1, x2, y2 = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3]
# 计算中心点坐标
cx = (x1 + x2) / 2 # 中心点x坐标 = (左边界x + 右边界x) / 2
cy = (y1 + y2) / 2 # 中心点y坐标 = (上边界y + 下边界y) / 2
# 计算宽度和高度
w = x2 - x1 # 宽度 = 右边界x - 左边界x
h = y2 - y1 # 高度 = 下边界y - 上边界y
# 将计算得到的中心点坐标、宽度和高度堆叠成新的张量
boxes = torch.stack((cx, cy, w, h), axis=-1)
return boxes
def box_center_to_corner(boxes):
"""将边界框从（中心点，宽度，高度）表示法转换为（左上角，右下角）表示法
该函数接收以(cx, cy, w, h)格式表示的边界框张量，其中：
- (cx, cy)：表示边界框中心点的坐标
- w：表示边界框的宽度
- h：表示边界框的高度
然后将其转换为(x1, y1, x2, y2)格式，其中：
- (x1, y1)：表示边界框左上角的坐标
- (x2, y2)：表示边界框右下角的坐标
参数:
boxes (torch.Tensor): 形状为(N, 4)的张量，包含N个边界框的中心点坐标、宽度和高度
返回:
torch.Tensor: 形状为(N, 4)的张量，包含N个边界框的左上角和右下角坐标
"""
# 分别提取所有边界框的中心点坐标、宽度和高度
cx, cy, w, h = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3]
# 计算左上角坐标
x1 = cx - 0.5 * w # 左边界x = 中心点x - 宽度/2
y1 = cy - 0.5 * h # 上边界y = 中心点y - 高度/2
# 计算右下角坐标
x2 = cx + 0.5 * w # 右边界x = 中心点x + 宽度/2
y2 = cy + 0.5 * h # 下边界y = 中心点y + 高度/2
# 将计算得到的左上角和右下角坐标堆叠成新的张量
boxes = torch.stack((x1, y1, x2, y2), axis=-1)
return boxes
def multibox_prior(data, sizes, ratios):
"""生成以每个像素为中心具有不同形状的锚框
参数：
data：输入图像张量，维度为(批量大小, 通道数, 高度, 宽度)
sizes：锚框缩放比列表，元素个数为num_sizes，每个元素∈(0,1]
ratios：锚框宽高比列表，元素个数为num_ratios，每个元素>0
返回：
输出张量，维度为(1, 像素总数*每像素锚框数, 4)，表示所有锚框的坐标
"""
# 获取输入数据的高度和宽度
# in_height, in_width: 标量
in_height, in_width = data.shape[-2:]
# 获取设备信息以及尺寸和比例的数量
# device: 字符串; num_sizes, num_ratios: 标量
device, num_sizes, num_ratios = data.device, len(sizes), len(ratios)
# 计算每个像素点产生的锚框数量 = 尺寸数 + 宽高比数 - 1
# boxes_per_pixel: 标量
boxes_per_pixel = (num_sizes + num_ratios - 1)
# 将尺寸和比例转换为张量
# size_tensor: 维度为(num_sizes,)
# ratio_tensor: 维度为(num_ratios,)
size_tensor = torch.tensor(sizes, device=device)
ratio_tensor = torch.tensor(ratios, device=device)
# 为了将锚点移动到像素的中心，需要设置偏移量
# 因为一个像素的高为1且宽为1，我们选择偏移中心0.5
# offset_h, offset_w: 标量
offset_h, offset_w = 0.5, 0.5
# 计算高度和宽度方向上的步长（归一化）
# steps_h, steps_w: 标量
steps_h = 1.0 / in_height # 在y轴上缩放步长
steps_w = 1.0 / in_width # 在x轴上缩放步长
# 生成锚框的所有中心点
# center_h: 维度为(in_height,)
# center_w: 维度为(in_width,)
center_h = (torch.arange(in_height, device=device) + offset_h) * steps_h
center_w = (torch.arange(in_width, device=device) + offset_w) * steps_w
# 使用meshgrid生成网格坐标
# shift_y, shift_x: 维度均为(in_height, in_width)
shift_y, shift_x = torch.meshgrid(center_h, center_w, indexing='ij')
# 将坐标展平为一维
# shift_y, shift_x: 展平后维度均为(in_height*in_width,)
shift_y, shift_x = shift_y.reshape(-1), shift_x.reshape(-1)
# 生成"boxes_per_pixel"个高和宽，
# 之后用于创建锚框的四角坐标(xmin,ymin,xmax,ymax)
# 计算锚框宽度：先计算尺寸与第一个比例的组合，再计算第一个尺寸与其余比例的组合
# w: 维度为(num_sizes + num_ratios - 1,)
w = torch.cat((size_tensor * torch.sqrt(ratio_tensor[0]),
sizes[0] * torch.sqrt(ratio_tensor[1:])))\
* in_height / in_width # 处理矩形输入，调整宽度
# 计算锚框高度：对应于宽度的计算方式
# h: 维度为(num_sizes + num_ratios - 1,)
h = torch.cat((size_tensor / torch.sqrt(ratio_tensor[0]),
sizes[0] / torch.sqrt(ratio_tensor[1:]</