目标检测第三篇：基于SSD的目标检测算法

SSD简介

SSD，全称Single Shot MultiBox Detector，是Wei Liu在ECCV 2016上提出的一种目标检测算法，截至目前是主要的检测框架之一，相比Faster RCNN有明显的速度优势，相比YOLO又有明显的mAP优势，是在 RCNN 系列和 YOLO 系列之外属于单阶段的另外一系列的奠基之作。
论文链接：https://arxiv.org/pdf/1512.02325.pdf
官方代码：https://github.com/weiliu89/caffe/tree/ssd

SSD的主要设计理念：根据不同的特征层设置不同大小的先验框，在不同的特征层上建立检测头和分类头，以满足大、小目标检测的需求。关于SSD的具体的技术细节，网上大神给出的解释很多，这里不再赘述。
网络架构如下：

本文主要是使用 pytorch 对 SSD 进行简单的实现。包括各个模块的讲解、数据的前、后处理、训练参数解释等。

网络搭建

论文中以 VGG16 为主干网络，替换了 VGG16 5_3 层和后面的部分，换成了 3*3 的卷积，再加上多尺度特征层，来实现多个检测和分类头。

卷积块

论文中以 卷积 + 激活函数为一个卷积标准模块，这里实现的时候加上了 BatchNorm，具体实现如下：

def conv_blk(in_channels, out_channels, stride=1, padding=1):
    """卷积块"""
    return nn.Sequential(nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride,padding=padding),
                         nn.BatchNorm2d(out_channels),
                         nn.ReLU())

下采样块

论文中以 池化层作为下采样层，来实现特征层的大小减半，这里实现的时候，将 卷积 + 池化层作为一个标准下采样模块，具体实现如下：

def down_sample_blk(in_channels, out_channels, ceil_mode = False):
    """下采样块"""
    return nn.Sequential(conv_blk(in_channels, out_channels),
                         nn.MaxPool2d(2, ceil_mode=ceil_mode))

主干网

使用上面两个模块，按照 VGG16 网络的基本架构实现 主干网：具体实现代码如下：

def backbone(input_shape):
    """搭建vgg16主干网络"""
    return nn.Sequential(
        conv_blk(input_shape[0], 64),
        down_sample_blk(64, 64),
        conv_blk(64, 128),
        down_sample_blk(128, 128),
        conv_blk(128, 256),
        conv_blk(256, 256),
        down_sample_blk(256, 512, ceil_mode=True),
        conv_blk(512, 512),
        conv_blk(512, 512),
        conv_blk(512, 512) # vgg16 4-3层
    )

多层特征提起层

主干网仅仅到 vgg16 4-3层， 后面的我们需要加上特征提取层，来实现输出不同特征层的需求，具体实现代码如下：

class extra_feature(nn.Module):
    """根据backbone输出的特征图，生成额外的特征图"""
    def __init__(self, input_shape) -> None:
        super().__init__()
        # (3, 300, 300) -> (512, 38, 38)
        self.out_layer1 = backbone(input_shape)

        # (512, 38, 38) -> (512, 19, 19)
        self.out_layer2 = nn.Sequential(
            nn.MaxPool2d(2, stride=1, padding=1),
            conv_blk(512, 512),
            conv_blk(512, 512),
            down_sample_blk(512, 521),
            conv_blk(521, 1024),
            conv_blk(1024, 1024)
        )

        # (512, 19, 19) -> (256, 10, 10)
        self.out_layer3 = nn.Sequential(
            conv_blk(1024, 256),
            conv_blk(256, 512, stride=2)
        )

        # (256, 10, 10) -> (256, 5, 5)
        self.out_layer4 = nn.Sequential(
            conv_blk(512, 128),
            conv_blk(128, 256, stride=2)
        )

        # (256, 5, 5) -> (256, 3, 3)
        self.out_layer5 = nn.Sequential(
            conv_blk(256, 128),
            conv_blk(128, 256, stride=2)
        )

        # (256, 3, 3) -> (256, 1, 1)
        self.out_layer6 = nn.Sequential(
            conv_blk(256, 128),
            nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=0)
        )
    def forward(self, x):
        out1 = self.out_layer1(x)
        out2 = self.out_layer2(out1)
        out3 = self.out_layer3(out2)
        out4 = self.out_layer4(out3)
        out5 = self.out_layer5(out4)
        out6 = self.out_layer6(out5)
        return out1, out2, out3, out4, out5, out6

经过 extra_feature 层，我们将会得到六层特征层的输出，对应上面图片的六层输出层，特征层的大小依次为：（38, 38）、（19， 19）、（10， 10）、（5， 5）、（3， 3）、（1， 1）。

输出头

得到六个输出层之后，我们需要对六层输出的通道数目进行调整，使其为我们需要分类的数目和需要输出的检测框的数目的大小。具体代码如下：

def mutiboxhead(num_classes):
    """搭建SSD多尺度检测头"""
    num_anchors = [4, 6, 6, 6, 4, 4] # 每个尺度的锚框数量
    num_channels = [512, 1024, 512, 256, 256, 256] # 每个尺度的通道数量
    cls_predictors = [] # 类别预测器
    bbox_predictors = [] # 边界框预测器
    for i in range(6):
        cls_predictors.append(nn.Conv2d(num_channels[i], num_anchors[i] * (num_classes + 1), kernel_size=3, padding=1))
        bbox_predictors.append(nn.Conv2d(num_channels[i], num_anchors[i] * 4, kernel_size=3, padding=1))
    cls_predictors = nn.ModuleList(cls_predictors)
    bbox_predictors = nn.ModuleList(bbox_predictors)
    return cls_predictors, bbox_predictors

根据上面构建的所有东西来实现我们的SSD网络，实现代码如下：

class SSD(nn.Module):
    def __init__(self, input_shape=(3, 300, 300), num_classes=20, **kwargs):
        super(SSD, self).__init__(**kwargs)
        self.num_classes = num_classes
        self.input_shape = input_shape
        self.extra_feature = extra_feature(input_shape)
        self.cls_predictors, self.bbox_predictors = mutiboxhead(num_classes)
    
    def forward(self, x):
        x = self.extra_feature(x)
        cls_preds = []
        bbox_preds = []
        for i, feature in enumerate(x):
            cls_preds.append(self.cls_predictors[i](feature).permute(0, 2, 3, 1)) # (batch_size, num_classes, h, w) -> (batch_size, h, w, num_classes)
            bbox_preds.append(self.bbox_predictors[i](feature).permute(0, 2, 3, 1)) # (batch_size, 4, h, w) -> (batch_size, h, w, 4)
        
        bbox_preds = torch.cat([pred.reshape(pred.shape[0], -1, 4) for pred in bbox_preds], dim=1) # (batch_size, num_anchors, 4)
        cls_preds = torch.cat([pred.reshape(pred.shape[0], -1, self.num_classes + 1) for pred in cls_preds], dim=1) # (batch_size, num_anchors, num_classes + 1)
        return bbox_preds, cls_preds # (batch_size, num_anchors, 4), (batch_size, num_anchors, num_classes + 1)

如果对搭建的网络参数或者输出的形状不是很放心的话，可以输出和跟踪一下看看具体的网络的架构。

if __name__ == "__main__":
    x = torch.randn(1, 3, 300, 300)
    ssd = SSD((3, 300, 300), 20)
    cls_preds, bbox_preds = ssd(x)
    print(cls_preds[0].shape, bbox_preds[0].shape)

数据处理

一般情况下，我们使用 COCO 或者 VOC 数据集进行训练，这里以 VOC 数据进行讲解，如何加载数据。

形成训练TXT

VOC 的标签为 XML 文件，在网络加载数据的时候，不是很方便，这里将其处理为 TXT 文件，方便我们读取。处理的结果如下：

每一行分别为：图片的路径，x_min y_min x_max y_max label … x_min y_min x_max y_max label
在后面加载数据的时候，直接读取即可，稍微方便一点。
处理的代码如下;

# 从 voc 数据集中读取数据，并生成可用于训练、验证和测试txt文件
import os
try:
    import xml.etree.cElementTree as ET
except ImportError:
    import xml.etree.ElementTree as ET

def get_class_names(class_names_path):
    """
    获取数据集的类别名字
    Args:
        class_names_path (str): 数据集路径
    Returns:
        class_names (list): 类别名字列表
    """
    with open(class_names_path, "r") as f:
        class_names = f.readlines()
    class_names = {c.strip(): i  for i, c in enumerate(class_names)}

    return class_names, len(class_names)

def write_txt(txt_path, JPEGImages_path, annotation_path, txt_f, name_dict):
    """
    将txt文件写入到txt_path
    Args:
        txt_path (str): txt文件保存路径
        txt_f (list): txt文件内容
    """
    with open(txt_path, "r") as f:
        img_name_lists = f.readlines()
        train_name_lists = [img_name.strip() for img_name in img_name_lists]
        for train_name in train_name_lists:
            txt_f.write(os.path.join(JPEGImages_path, train_name + ".jpg").replace('\\', '/'))
            xml_name = train_name + ".xml"
            xml_path = os.path.join(annotation_path, xml_name)
            tree = ET.parse(xml_path)
            root = tree.getroot()
            for obj in root.iter("object"):
                cls = name_dict[obj.find("name").text]
                xmlbox = obj.find("bndbox")
                xmin = xmlbox.find("xmin").text
                ymin = xmlbox.find("ymin").text
                xmax = xmlbox.find("xmax").text
                ymax = xmlbox.find("ymax").text
                txt_f.write(f",{xmin} {ymin} {xmax} {ymax} {cls}")
            txt_f.write("\n")

def convert_annotation_2_txt(dataset_path, txt_path):
    """
    将voc数据集的xml标注转换为txt格式
    Args:
        dataset_path (str): 数据集路径
        txt_path (str): txt文件保存路径
    """
    annotation_path = os.path.join(dataset_path, "Annotations")
    ImageSets_path = os.path.join(dataset_path, "ImageSets/Main")
    JPEGImages_path = os.path.join(dataset_path, "JPEGImages")

    train_txt_f = open(os.path.join(txt_path, "train.txt"), "w")
    val_txt_f = open(os.path.join(txt_path, "val.txt"), "w")
    test_txt_f = open(os.path.join(txt_path, "test.txt"), "w")

    name_dict, _ = get_class_names(os.path.join(txt_path, "voc_classes.txt"))

    write_txt(os.path.join(ImageSets_path, "train.txt"), JPEGImages_path, annotation_path, train_txt_f, name_dict)
    train_txt_f.close()

    write_txt(os.path.join(ImageSets_path, "val.txt"), JPEGImages_path, annotation_path, val_txt_f, name_dict)
    val_txt_f.close()

    write_txt(os.path.join(ImageSets_path, "test.txt"), JPEGImages_path, annotation_path, test_txt_f, name_dict)
    test_txt_f.close()
    
if __name__ == "__main__":
    convert_annotation_2_txt("你的\VOC07+12", "输出txt想要保存的地点")

运行完以上程序，会形成 train.txt, test.txt val.txt，分别对应训练集、验证集和测试集。训练的时候读取对应的txt即可。

Dataset

获得完上面的 txt， 我们就可以编写自己的 Dataset 函数，形成自己在网络训练时候的数据加载方式。具体代码如下：

# 加载数据集的函数
class SSDDataset(Dataset):
    """
    定义自己的数据集加载方式
    """
    def __init__(self, data_lines, transform=None):
        """
        初始化文件路径和数据增强所需的信息
        Args:
            data (list): 文件路径列表
            transform: torchvision.transforms
        """
        self.data_lines = data_lines
        self.transform = transform

    def __len__(self):
        """
        返回数据集的总长度
        Returns:
            length (int): 数据集的总长度
        """
        return len(self.data_lines)
    
    def __getitem__(self, index):
        """
        通过索引返回一个数据
        Args:
            index (int): 索引
        Returns:
            img (Tensor): 图像数据
            label (Tensor): 标签数据
        """
        data_line = self.data_lines[index].split(",")
        image       = np.array((Image.open(data_line[0]).convert("RGB")))
        bboxes      = np.array([list(map(int, box.split(' '))) for box in data_line[1:]])

        if self.transform is not None:
            augmentations = self.transform(image=image, bboxes=bboxes)
            image = augmentations["image"]
            bboxes = augmentations["bboxes"]

        return image, bboxes

也就是通过输入的 index，来获取 图片 和对应的 标签信息， transform 是是否采用数据增强，数据增强的具体代码如下：

train_transforms = A.Compose(
    [
        # 保持图片的比例进行图片的放大
        A.LongestMaxSize(max_size=300),
        # 不保证图片比例进行放大，需要的时候会对图片添加 0 
        A.PadIfNeeded(
            min_height=300, min_width=300, border_mode=cv2.BORDER_CONSTANT
        ),
        # 进行正则化
        A.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010], max_pixel_value=255,),
        # 转变为 tensor
        ToTensorV2(),
    ],
    # 对检测框进行调整
    bbox_params=A.BboxParams(format="pascal_voc", label_fields=[],check_each_transform=False),
) 

val_transforms = A.Compose(
    [
        # 保持图片的比例进行图片的放大
        A.LongestMaxSize(max_size=300),
        # 不保证图片比例进行放大，需要的时候会对图片添加 0 
        A.PadIfNeeded(
            min_height=300, min_width=300, border_mode=cv2.BORDER_CONSTANT
        ),
        # 进行正则化
        A.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010], max_pixel_value=255,),
        # 转变为 tensor
        ToTensorV2(),
    ],
    # 对检测框进行调整
    bbox_params=A.BboxParams(format="pascal_voc", label_fields=[], check_each_transform=False),
)

因为我们使用的格式为 pascal_voc 的格式，所以这里的我们选择的增强方式为 pascal_voc，如果使用的为 coco 的话，换成 coco 即可。
一般情况下，在训练的时候和在验证的时候采用的数据增强手段是不一样的，训练的时候，我们强调数据的多样性，会使用较多的数据增强的手段，为反转、随机剪裁、色域变换等，但是在验证的时候，我们强调当前网络的性能，一般采取较少的数据增强手段。

注意：我们在对图片进行数据增强的时候，要保证检测的一致性。

DataLoader

刚刚只是定义了如何获取图片，但是如何设置为可以供 pytorch 使用，需要放到 DataLoader 中， 形成迭代器，来实现不间断的拿取 batch_size 的图片。
具体代码如下：

def load_data(root_path, batch_size=8, num_workers=0):
    # 第一步：加载数据集
    with open(os.path.join(root_path, "train.txt"), "r") as f:
        train_lines = f.readlines()
    with open(os.path.join(root_path, "train.txt"), "r") as f:
        val_lines = f.readlines()

    train_dataset = SSDDataset(train_lines, train_transforms)
    val_dataset = SSDDataset(val_lines, val_transforms)
    
    train_iter = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers, collate_fn=collate_fn)
    test_iter = DataLoader(val_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers, collate_fn=collate_fn)
    return train_iter, test_iter

def collate_fn(batch):
    """
    自定义一个函数，将一个batch的数据拼接起来
    Args:
        batch (list): batch数据
    Returns:
        images (Tensor): 图像数据
        bboxes (list): 边界框数据
    """
    images = []
    bboxes = []
    for image, bbox in batch:
        images.append(image)
        bboxes.append(torch.tensor(bbox))
    return torch.stack(images, dim=0), bboxes

Anchors

生成先验框

SSD 网络需要生成在不同的特征层生成先验框，以满足对不同大、小物体检测的需求；具体实现代码如下：

# SSD 生成 anchors 的方式, 中心宽高的方式
def generate_anchors(feature_maps_size=[38, 19, 10, 5, 3, 1], image_size=300, steps=[8, 16, 32, 64, 100, 300], 
                    min_sizes=[30, 60, 111, 162, 213, 264], max_sizes=[60, 111, 162, 213, 264, 315], 
                    aspect_ratios=[[2], [2, 3], [2, 3], [2, 3], [2], [2]], clip=True):
    mean = []
    for k, f in enumerate(feature_maps_size): #[38, 19, 10, 5, 3, 1],
        for i, j in product(range(f), repeat=2):
            f_k = image_size / steps[k]
            # unit center x,y
            cx = (j + 0.5) / f_k
            cy = (i + 0.5) / f_k

            # aspect_ratio: 1
            # rel size: min_size
            s_k = min_sizes[k] / image_size
            mean += [cx, cy, s_k, s_k]

            # aspect_ratio: 1
            # rel size: sqrt(s_k * s_(k+1))
            s_k_prime = sqrt(s_k * (max_sizes[k]/image_size))
            mean += [cx, cy, s_k_prime, s_k_prime]

            # rest of aspect ratios
            for ar in aspect_ratios[k]:
                mean += [cx, cy, s_k*sqrt(ar), s_k/sqrt(ar)]
                mean += [cx, cy, s_k/sqrt(ar), s_k*sqrt(ar)]
    # back to torch land
    prior_anchors = torch.Tensor(mean).view(-1, 4)  # * image_size ##[8732,4]

    # 对于超出边界的先验框进行裁剪到边界内
    if clip:
        prior_anchors.clamp_(max=1, min=0)
	# 返回的时候转化为 corner 的数据格式，并且乘上 图片 对应的尺度获取真实大小
    return box_center_to_corner(prior_anchors) * image_size

匹配先验框

在训练的时候，我们需要将真实的标签框和先验框进行匹配，来获取训练的正样本。也就是先确保每个标签框有一个先验框负责预测，完成 一对一 的匹配。如果其他的先验框没有被分配，且和其中某个标签框的 iou 也很大，也将对列为正样本，完成 先验框 -> 标签框 的 多对一 的匹配。具体技术细节，请参考：沐神的动手学深度学习的 13.4. 锚框 的章节。
具体实现代码如下：

def assign_anchor_to_bbox(ground_truth, anchors, device, iou_threshold=0.5):
    """ 为锚框分配边界框

    Args:
        ground_truth (Tensor): 边界框，形状为 (n, 4)
        anchors (Tensor): 锚框，形状为 (k, 4)
        device (torch.device): 用于分配张量的设备
        iou_threshold (float): 阈值
    Returns:
        Tensor: 所分配的边界框，形状为 (k, 4)
    """
    num_anchors, num_gt_boxes = anchors.shape[0], ground_truth.shape[0]

    # 计算锚框和边界框的交并比
    anchors_boxes_iou = box_iou(anchors, ground_truth)

    # 为每个锚框分配边界框的索引, 形状为 (k, ), 默认为 -1
    anchor_to_bbox_map = torch.full((num_anchors,), -1, dtype=torch.long, device=device)

    # 找到每个 anchors 对应的最大 iou 值和索引
    max_ious, indices = torch.max(anchors_boxes_iou, dim=1)
    
    # 将交并比大于阈值的锚框分配给边界框
    anchor_to_bbox_map[max_ious >= iou_threshold] = indices[max_ious >= iou_threshold]

    # 找到应该丢弃的行和列，也就是找到每一个 gt 真实值匹配的 anchor 所以对应的行和列
    col_discard = torch.full((num_anchors, ), -1)
    row_discard = torch.full((num_gt_boxes,), -1)
    # 每个 iou 较大的 anchor只保留一个gt对应
    for _ in range(num_gt_boxes):
        # 找到目前 iou 最大的位置，因为没有指定维度的话，默认将所有维度参与
        # anchors_boxes_iou 的维度为 ： [num_anchors, num_gt_boxes]
        max_idx = torch.argmax(anchors_boxes_iou)
        # 找到对应的列
        box_idx = (max_idx % num_gt_boxes).long()
        # 找到对应的行
        anc_idx = (max_idx / num_gt_boxes).long()
        # 将对应的 gt 作为 anchor 的真值
        anchor_to_bbox_map[anc_idx] = box_idx
        
        # 取出对应的行和列
        anchors_boxes_iou[:, box_idx] = col_discard
        anchors_boxes_iou[anc_idx, :] = row_discard
    
    return anchor_to_bbox_map

位置 offset

网络并不是直接预测标签的中心点和宽高，那样会增加网络预测的难度，我们预测的时候，先验框相对于标签框的偏移量，包括中心点的偏移量和宽高的偏移量。参考 沐神动手学深度的讲解如下;

具体代码如下：

def offset_boxes(c_anc, assigned_bbox, eps=1e-6):
    """根据所分配的边界框来调整锚框

    Args:
        c_anc (Tensor): 锚框，形状为 (n, 4) center 形状
        assigned_bbox (Tensor): 所分配的边界框，形状为 (n, 4) corner 形状
        eps (float): 一个极小值，防止被零整除
    Returns:
        Tensor: 调整后的锚框，形状为 (n, 4)
    """
    c_anc = box_corner_to_center(c_anc)
    c_assigned_bbox = box_corner_to_center(assigned_bbox)
    offset_xy = 10.0 * (c_assigned_bbox[:, :2] - c_anc[:, :2]) / (c_anc[:, 2:] + eps)
    offset_wh = 5.0 * torch.log((c_assigned_bbox[:, 2:] + eps) / (c_anc[:, 2:] + eps))
    offset = torch.cat([offset_xy, offset_wh], axis=1)
    return offset

这些偏移量将作为训练的时候，真正的标签。

损失函数

有了网络的输出和真正的偏移量的标签，我们就可以定义我们的损失函数了。
关于 SSD 损失讲解，可以参考大神的解释：链接
简单来说，分为两个部分：

前面是置信度损失函数，后面是位置损失函数，中间为权重因子。再详细划分如下：


根据上面公式，定义损失函数如下：

def loc_loss(self, loc_preds, bbox_labels, bbox_masks):
        """计算位置损失

        Args:
            loc_preds (Tensor): 预测的位置结果，形状为 (batch_size, num_anchors, 4)
            bbox_labels (Tensor): 真实的位置标签，形状为 (batch_size, num_anchors, 4)
            bbox_masks (Tensor): 真实的位置标签的掩码，形状为 (batch_size, num_anchors, 4)
        Returns:
            Tensor: 位置损失
        """
        # 第一步：计算损失
        return  F.l1_loss(loc_preds * bbox_masks, bbox_labels * bbox_masks, reduction='sum')
    
    def conf_loss(self, conf_preds, cls_labels):
        """计算分类损失

        Args:
            conf_preds (Tensor): 预测的分类结果，形状为 (batch_size, num_anchors, num_classes)
            cls_labels (Tensor): 真实的分类标签，形状为 (batch_size, num_anchors, 1)
        Returns:
            Tensor: 分类损失
        """
        # 第一步：计算正样本的损失
        # 获取正样本的索引
        pos_mask = cls_labels > 1
        # 获取正样本的数量
        num_pos = pos_mask.long().sum(dim=1, keepdim=True) # (batch_size, 1)
        # 计算正样本的损失, 交叉熵损失函数的输入为 (N, C) 和 (N, )，其中 N 为样本数量，C 为类别数量，会自动找到(N, C)中的对应的值，计算损失
        pos_loss = F.cross_entropy( conf_preds.reshape(-1, conf_preds.shape[-1])[pos_mask.reshape(-1)], cls_labels[pos_mask])

        # 第二步：计算负样本的损失
        # 获取负样本的索引
        neg_mask = cls_labels == 0
        # 计算负样本的损失
        neg_loss = F.cross_entropy(conf_preds.reshape(-1, conf_preds.shape[-1])[neg_mask.reshape(-1)], cls_labels[neg_mask])
        # 第三步：计算损失
        loss = pos_loss +  neg_loss
        # 返回损失
        return loss

根据上面的损失函数，可以定义我们训练的总的损失函数类：具体代码如下：

class MultiBoxLoss():
    """根据预测的结果计算损失函数"""
    def __init__(self, num_classes, device, alpha=1) -> None:
        super().__init__()
        self.num_classes   = num_classes
        self.image_size    = 300
        self.prior_anchors = generate_anchors()
        self.device        = device
        self.alpha         = alpha

    def __call__(self, net_output, labels):
        """计算损失函数

        Args:
            net_output (tuple): 网络输出，包含回归预测和分类预测，
            形状为 loc_preds:(batch_size, num_anchors, 4), conf_preds: (batch_size, num_anchors, 21)
            
            labels (Tensor): 标签，形状为 (n, 5)，其中 n 是所有边界框的数量, 5 表示 (类别, x, y, w, h)
        Returns:
            Tensor: 损失值
        """
        # 第一步：获取网络输出 (batch_size, num_anchors, 4), (batch_size, num_anchors, 21)
        loc_preds, conf_preds = net_output
        batch_size, device = len(labels), loc_preds.device
        
        num_anchors = self.prior_anchors.shape[0] # self.prior_anchors (num_anchors, 4)

        # 第二步：为不同大小的特征层生成对应的 anchors 找到对应 target 标号, 并且将其转换为相对于 anchor 的偏移量
        bbox_labels, bbox_masks, cls_labels = multibox_target(self.prior_anchors, labels)
        cls_labels = cls_labels.to(device)
        bbox_labels = bbox_labels.to(device)
        bbox_masks = bbox_masks.to(device)

        num_prior_anchors = bbox_masks[:,:, 0].sum() # (batch_size, 1)

        # 第三步：计算损失
        # 将预测的结果转换为 (batch_size, num_anchors, num_classes + 1)， + 1 为背景类
        conf_preds = conf_preds.view(batch_size, num_anchors, self.num_classes + 1)
        
        # 计算分类损失
        cls_loss = self.conf_loss(conf_preds, cls_labels)
        # 计算位置损失
        loc_loss = self.loc_loss(loc_preds, bbox_labels, bbox_masks)

        all_loss = (cls_loss + self.alpha * loc_loss) / num_prior_anchors

        return all_loss
    
    def loc_loss(self, loc_preds, bbox_labels, bbox_masks):
        """计算位置损失

        Args:
            loc_preds (Tensor): 预测的位置结果，形状为 (batch_size, num_anchors, 4)
            bbox_labels (Tensor): 真实的位置标签，形状为 (batch_size, num_anchors, 4)
            bbox_masks (Tensor): 真实的位置标签的掩码，形状为 (batch_size, num_anchors, 4)
        Returns:
            Tensor: 位置损失
        """
        # 第一步：计算损失
        return  F.l1_loss(loc_preds * bbox_masks, bbox_labels * bbox_masks, reduction='sum')
    
    def conf_loss(self, conf_preds, cls_labels):
        """计算分类损失

        Args:
            conf_preds (Tensor): 预测的分类结果，形状为 (batch_size, num_anchors, num_classes)
            cls_labels (Tensor): 真实的分类标签，形状为 (batch_size, num_anchors, 1)
        Returns:
            Tensor: 分类损失
        """
        # 第一步：计算正样本的损失
        # 获取正样本的索引
        pos_mask = cls_labels > 1
        # 获取正样本的数量
        num_pos = pos_mask.long().sum(dim=1, keepdim=True) # (batch_size, 1)
        # 计算正样本的损失, 交叉熵损失函数的输入为 (N, C) 和 (N, )，其中 N 为样本数量，C 为类别数量，会自动找到(N, C)中的对应的值，计算损失
        pos_loss = F.cross_entropy( conf_preds.reshape(-1, conf_preds.shape[-1])[pos_mask.reshape(-1)], cls_labels[pos_mask])

        # 第二步：计算负样本的损失
        # 获取负样本的索引
        neg_mask = cls_labels == 0
        # 计算负样本的损失
        neg_loss = F.cross_entropy(conf_preds.reshape(-1, conf_preds.shape[-1])[neg_mask.reshape(-1)], cls_labels[neg_mask])
        # 第三步：计算损失
        loss = pos_loss +  neg_loss
        # 返回损失
        return loss

    #-------------------------------#
    #   计算预测结果
    #-------------------------------#
    def eval(self, net_output, labels, show_img_flag=False):
        """计算准确率

        Args:
            net_output (tuple): 网络输出，包含回归预测和分类预测，形状为 (loc_preds (1, num_anchors, 4), conf_preds (1, num_anchors, 21))
            labels (Tensor): 标签，形状为 (n, 5)，其中 n 是所有边界框的数量, 5 表示 (类别, x, y, w, h)
        Returns:
            Tensor: 损失值
        """
        # 第一步：获取网络输出 (1, num_anchors, 4), (1, num_anchors, 21)
        loc_preds, conf_preds = net_output

        batch_size, device = len(labels), loc_preds.device
        # self.prior_anchors (num_anchors, 4)
        num_anchors = self.prior_anchors.shape[0]

        # 第二步：为不同大小的特征层生成对应的 anchors 找到对应 target 标号, 并且将其转换为相对于 anchor 的偏移量
        bbox_labels, bbox_masks, cls_labels = multibox_target(self.prior_anchors, labels)
        cls_labels = cls_labels.to(device)
        bbox_labels = bbox_labels.to(device)
        bbox_masks = bbox_masks.to(device)
		
		num_prior_anchors = bbox_masks[:,:, 0].sum() # (1,)

        # 第三步：计算准确度
        # 计算分类准确度
        cls_acc = self.cls_eval(conf_preds, cls_labels) / (batch_size * num_anchors)
        #计算分类损失
        cls_loss = self.conf_loss(conf_preds, cls_labels) # / (batch_size * num_anchors)
        # 计算位置损失
        bbox_loss = self.bbox_eval(loc_preds, bbox_labels, bbox_masks) #  / (batch_size * num_anchors)
        # 返回准确度
        return cls_acc, cls_loss / num_prior_anchors , bbox_loss / num_prior_anchors 
        
    def cls_eval(self, cls_preds, cls_labels):
        # 由于类别预测结果放在最后一维，argmax需要指定最后一维。
        return float((cls_preds.argmax(dim=-1).type(cls_labels.dtype) == cls_labels).sum())

    def bbox_eval(self, bbox_preds, bbox_labels, bbox_masks):
        return float((torch.abs((bbox_labels - bbox_preds) * bbox_masks)).sum())

其中 multibox_target 函数的作用为 ：为不同大小的特征层生成对应的 anchors 找到对应 target 标号, 并且将其转换为相对于 anchor 的偏移量

def multibox_target(anchors, labels):
    """ 使用真实边界框标记锚框

    Args:
        anchors (Tensor): 先验框，对应的大小为，(num_anchors, 4)
        labels (Tensor): 多个图片的真实标签，(batch_size, num_labels, 4)  都是 corner 的形式 
    Returns:
        bbox_offset  (batch_size, num_boxes, 4)
        bbox_mask    (batch_size, num_boxes, 4)
        class_labels (batch_size, num_boxes, 1)
    """
    # labels 为多个图片的box的标签合集，维度为 [batch_size, num_gt, 4 + 1]， 其中 1 为类别
    batch_size, anchors = len(labels), anchors.squeeze(0)
    batch_offset, batch_mask, batch_class_labels = [], [], []
    device, num_anchors = anchors.device, anchors.shape[0]
    
    for i in range(batch_size):
        # 取出第 i 张图片对应的 box 的 gt
        label = labels[i]
        # 找到和标签 box 最接近的真实标签框，作为训练的正样本
        anchor_bbox_map = assign_anchor_to_bbox(label[:, :4], anchors, device)
        # 构造 mask 矩阵，大于零的我们认为找打了相应的gt，扩展到4维的目的是为了后面bbox的四个点的mask
        bbox_mask = ((anchor_bbox_map >= 0).float().unsqueeze(-1)).repeat(1, 4)
        # 将类标签和分配的边界框坐标初始化为0
        class_labels = torch.zeros(num_anchors, dtype=torch.long, device=device)
        assigned_bbox = torch.zeros((num_anchors, 4), dtype=torch.float32, device=device)
        
        # 使用真实边界框来标记锚框的类别
        # 如果一个锚框没有被分配，标记为背景，值为零
        indices_true = torch.nonzero(anchor_bbox_map >= 0)
        # 拿出已经分配好的锚框对应的 gt bbox
        bbox_idx = anchor_bbox_map[indices_true]
        # 将类别和对应的标签进行赋值。label + 1 的原因是为了和 0 进行区别。
        class_labels[indices_true] = label[bbox_idx, 4].long() + 1
        assigned_bbox[indices_true] = label[bbox_idx, :4].float()

        # 偏移量的转换，得到真实框和对应anchor之间的位置偏移
        offset = offset_boxes(anchors, assigned_bbox) * bbox_mask
        # 记录一下对应的关系
        batch_offset.append(offset)
        batch_mask.append(bbox_mask)
        batch_class_labels.append(class_labels)
        
    bbox_offset = torch.stack(batch_offset)
    bbox_mask = torch.stack(batch_mask)
    class_labels = torch.stack(batch_class_labels)
    
    return bbox_offset, bbox_mask, class_labels

训练

定义相关的优化器、学习率等，我们就可以开始我们的训练了；具体训练代码如下;

if __name__ == "__main__":

    # -------------------超参数-------------------#
    batch_size = 32
    num_workers = 4
    input_size = (3, 300, 300)
    num_classes = 20 
    learning_rate = 1e-2
    epochs = 100
    weight_decay = 5e-4
    root_path = "./data"
    # -------------------超参数-------------------#

    # 设置使用的设备
    # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")
    
    # 第一步：加载数据集
    train_iter, test_iter = load_data(root_path, batch_size, num_workers)

    # 第二步：定义模型
    net = SSD(input_size, num_classes)
    net.to(device)
    # summary(net,input_size=(3, 300, 300))

    # 第三步：定义损失函数
    Loss = MultiBoxLoss(num_classes, device)

    # 第四步：定义优化器
    optimizer = torch.optim.SGD(net.parameters(), lr=learning_rate, momentum=0.9, weight_decay=weight_decay)

    # 第五步：训练模型
    train_day = datetime.now().strftime("%Y_%m_%d_%H_%M_%S")

    print("#", "-"*30, "训练参数", "-"*30, "#")
    print("训练日期:{0:>65}".format(train_day))
    print("训练设备:{0:>65}".format(str(device)))
    print("训练数据集大小:{0:>59}".format(len(train_iter) * batch_size))
    print("测试数据集大小:{0:>59}".format(len(test_iter) * batch_size))
    print("训练批次大小:{0:>61}".format(batch_size))
    print("训练轮次:{0:>65}".format(epochs))
    print("学习率:{0:>67}".format(learning_rate))
    print("权重衰减:{0:>65}".format(weight_decay))
    print("#", "-"*30, "训练开始", "-"*30, "#")

    train_file = "./log/train_log_" + train_day
    os.makedirs(train_file, exist_ok=True)

    min_loss = 1e+10
    for epoch in range(epochs):
        net.train()
        phgr = tqdm(train_iter, total=len(train_iter), desc="Train epoch " + str(epoch + 1) + "/" + str(epochs))
        epoch_train_losses = []
        for i, (images, labels) in enumerate(phgr):
            # 将数据转换为 cuda 的 tensor
            images = images.to(device)

            # 前向传播
            net_output = net(images)

            # 计算损失
            net_loss = Loss(net_output, labels)
            loss = net_loss.item()
            epoch_train_losses.append(loss)
            # 反向传播
            optimizer.zero_grad()
            net_loss.backward()
            optimizer.step()

            phgr.set_postfix({"loss": loss})
        
        # 保存最优模型
        epoch_train_loss = (sum(epoch_train_losses) / len(epoch_train_losses))
        print("epoch_train_loss:", epoch_train_loss, " min_loss:", min_loss)
        if epoch_train_loss < min_loss:
            min_loss = epoch_train_loss
            save_weight_path = train_file + f"/ssd_best_loss.pth"
            torch.save(net.state_dict(), save_weight_path) 

        if epoch % 5 == 0 or epoch == epochs - 1:
            net.eval()
            epoch_test_losses = []
            show_img_flag = True
            with torch.no_grad():
                phgr_test = tqdm(test_iter, total=len(test_iter), desc="Test epoch " + str(epoch + 1) + "/" + str(epochs))
                for images, labels in phgr_test:
                    # 将数据转换为 cuda 的 tensor
                    tensor_images = images.to(device)

                    # 前向传播
                    net_output = net(tensor_images)

                    if show_img_flag:
                        # 显示第一张图片的预测结果作为显示
                        # decoder_out = decoder(net_output, Loss.prior_anchors)
                        # out_frame = draw_infer_box(images[0], decoder_out[0], decoder_out[1], decoder_out[2], class_names, colors)
                        # cv2.imwrite("train_eval_result.png", out_frame)
                        show_img_flag = False

                    # 计算准确度
                    cls_acc, cls_loss, bbox_loss = Loss.eval(net_output, labels)
                    epoch_test_losses.append(cls_loss + bbox_loss)
                    phgr_test.set_postfix({"cls_acc": cls_acc," cls_loss": cls_loss.item(), " bbox_loss": bbox_loss})

            epoch_test_loss = (sum(epoch_test_losses) / len(epoch_test_losses))
            save_weight_path = train_file + f"/ssd_epoch_{epoch + 1}_testloss_{0:.3f}_trainloss_{0:.3f}.pth".format(epoch_test_loss, epoch_train_loss)

代码及参考

完整包含预测的代码地址为：
参考
1、沐神的动手学深度学习