从代码学习深度学习 - 单发多框检测（SSD）PyTorch版

前言

大家好！欢迎来到“从代码学习深度学习”系列博客。目标检测是计算机视觉领域的核心任务之一，旨在识别图像或视频中特定类别的对象实例，并确定它们的位置和范围。近年来，深度学习技术极大地推动了目标检测的发展，涌现出许多优秀的算法，如 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 torchvisiondef 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) / 256return images, targets_tensorclass 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.0def 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.0def __getitem__(self, idx):"""获取指定索引处的值输入：idx: 索引值  # 输入参数：要访问的数据索引输出：float: 指定索引处的值  # 返回指定位置的累加值"""return self.data[idx]  # 返回指定索引处的数据值def try_gpu(i=0):"""如果存在，则返回gpu(i)，否则返回cpu()Args:i (int, optional): GPU设备的编号，默认为0，表示尝试使用第0号GPUReturns:torch.device: 返回可用的设备对象，如果指定编号的GPU可用则返回GPU，否则返回CPU"""# 检查系统中可用的GPU数量是否大于等于i+1if 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) / 2cy = (y1 + y2) / 2  # 中心点y坐标 = (上边界y + 下边界y) / 2# 计算宽度和高度w = x2 - x1  # 宽度 = 右边界x - 左边界xh = y2 - y1  # 高度 = 下边界y - 上边界y# 将计算得到的中心点坐标、宽度和高度堆叠成新的张量boxes = torch.stack((cx, cy, w, h), axis=-1)return boxesdef 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 - 宽度/2y1 = cy - 0.5 * h  # 上边界y = 中心点y - 高度/2# 计算右下角坐标x2 = cx + 0.5 * w  # 右边界x = 中心点x + 宽度/2y2 = cy + 0.5 * h  # 下边界y = 中心点y + 高度/2# 将计算得到的左上角和右下角坐标堆叠成新的张量boxes = torch.stack((x1, y1, x2, y2), axis=-1)return boxesdef 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_hcenter_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:]</