秃姐学AI系列之：SSD代码实现

多尺度目标检测锚框

为了演示如何在多个尺度下生成锚框，让我们先读取一张图像。它的高度和宽度分别为561和728像素

%matplotlib inline
import torch
from d2l import torch as d2l

img = d2l.plt.imread('../img/catdog.jpg')
h, w = img.shape[:2]
h, w

# 输出
(561, 728)

我们将卷积图层的二维数组输出称为特征图。通过定义特征图的形状，我们可以确定任何图像上均匀采样锚框的中心。

多尺度生成锚框

display_anchors函数定义如下：

我们在特征图（fmap）上生成锚框（anchors），每个单位（像素）作为锚框的中心。由于锚框中的(x,y)轴坐标值（anchors）已经被除以特征图（fmap）的宽度和高度，因此 anchors 这些值介于0和1之间，表示特征图中锚框的相对位置。

由于锚框（anchors）的中心分布于特征图（fmap）上的所有单位，因此这些中心必须根据其相对空间位置在任何输入图像上均匀分布。更具体地说，给定特征图的宽度和高度fmap_w和fmap_h，以下函数将均匀地对任何输入图像中 fmap_h行和 fmap_w列中的像素进行采样。以这些均匀采样的像素为中心，将会生成大小为s（假设列表s的长度为1）且宽高比（ratios）不同的锚框。

def display_anchors(fmap_w, fmap_h, s):
    d2l.set_figsize()
    # 前两个维度上的值不影响输出，只关心高宽，并不关心通道数
    fmap = torch.zeros((1, 10, fmap_h, fmap_w))
    anchors = d2l.multibox_prior(fmap, sizes=s, ratios=[1, 2, 0.5])
    bbox_scale = torch.tensor((w, h, w, h))
    d2l.show_bboxes(d2l.plt.imshow(img).axes, anchors[0] * bbox_scale)

上面的函数相当于定死了radios（宽=高，宽=2高，2宽=高），可选的只有锚框的大小s

上面的multibox_prior()函数详见锚框博文的代码实现

探测目标

为了在显示时更容易分辨，在这里具有不同中心的锚框不会重叠：锚框的尺度设置为0.15，特征图的高度和宽度设置为4。我们可以看到，图像上4行和4列的锚框的中心是均匀分布的。

display_anchors(fmap_w=4, fmap_h=4, s=[0.15])

然后，我们将特征图的高度和宽度减小一半，然后使用较大的锚框来检测较大的目标。当尺度设置为0.4时，一些锚框将彼此重叠。

display_anchors(fmap_w=2, fmap_h=2, s=[0.4])

最后，我们进一步将特征图的高度和宽度减小一半，然后将锚框的尺度增加到0.8。此时，锚框的中心即是图像的中心。

display_anchors(fmap_w=1, fmap_h=1, s=[0.8])

SSD实现

类别预测层

预测锚框的类别

设目标类别的数量为 q。这样一来，锚框有 q+1个类别，其中 0 类是背景。在某个尺度下，设特征图的高和宽分别为 h 和 w。如果以其中每个单元为中心生成 a 个锚框，那么我们需要对 hwa 个锚框进行分类。如果使用全连接层作为输出，很容易导致模型参数过多。回忆我们之前介绍的使用卷积层的通道来输出类别预测的方法， 单发多框检测采用同样的方法来降低模型复杂度。

具体来说，类别预测层使用一个保持输入高和宽的卷积层（通过 3*3 卷积核和 1 的填充来保持）。这样一来，输出和输入在特征图宽和高上的空间坐标一一对应。考虑输出和输入同一空间坐标（x、y）：输出特征图上（x、y）坐标的通道里包含了以输入特征图（x、y）坐标为中心生成的所有锚框的类别预测。因此输出通道数为 a*(q+1)，其中索引为 i*(q+1)+j（0≤j≤q）的通道代表了索引为 i 的锚框有关类别索引为 j 的预测（因为存在第0类的背景类）。

在下面，我们定义了这样一个类别预测层，通过参数 num_anchors和 num_classes分别指定了 a 和 q。该图层使用填充为 1 的 3×3 的卷积层。此卷积层的输入和输出的宽度和高度保持不变。

%matplotlib inline
import torch
import torchvision
from torch import nn
from torch.nn import functional as F
from d2l import torch as d2l


def cls_predictor(num_inputs, num_anchors, num_classes):
"""num_inputs:输入通道数    num_anchors:多少个锚框    num_classes:有多少个类
输出的通道数 = 锚框个数 * （类别数 + 1）    +1是背景类的作用
其实就是预测总像素数*每个像素的锚框数*（每个锚框的预测类别 + 背景类）
即 对每个锚框都要预测其所有类别
这个kernel与padding大小的组合不会改变输入的大小"""
    return nn.Conv2d(num_inputs, num_anchors * (num_classes + 1),
                     kernel_size=3, padding=1)

边界框预测层

边界框预测层的设计与类别预测层的设计类似。唯一不同的是，这里需要为每个锚框预测4个偏移量（通常是中心点的X、Y坐标偏移和边界框 H、W偏移；是相对偏移），而不是 q+1 个类别。

def bbox_predictor(num_inputs, num_anchors):
    return nn.Conv2d(num_inputs, num_anchors * 4, kernel_size=3, padding=1)

连接多尺度的预测

正如我们所提到的，单发多框检测（SSD）使用多尺度特征图来生成锚框并预测其类别和偏移量。在不同的尺度下，特征图的形状或以同一单元为中心的锚框的数量可能会有所不同。因此，不同尺度下预测输出的形状可能会有所不同。

在以下示例中，我们为同一个小批量构建两个不同比例（Y1和Y2）的特征图，其中Y2的高度和宽度是Y1的一半。以类别预测为例，假设Y1和Y2的每个单元分别生成了5个和3个锚框。进一步假设目标类别的数量为10，对于特征图Y1和Y2，类别预测输出中的通道数分别为5×(10+1)=55和3×(10+1)=33，其中任一输出的形状是（批量大小，通道数，高度，宽度）。

def forward(x, block):
    return block(x)

# 这里的batch_size=2不需要纠结，举例而已，只要知道它是不变的就可以了
Y1 = forward(torch.zeros((2, 8, 20, 20)), cls_predictor(8, 5, 10))
Y2 = forward(torch.zeros((2, 16, 10, 10)), cls_predictor(16, 3, 10))
Y1.shape, Y2.shape

# 输出
(torch.Size([2, 55, 20, 20]), torch.Size([2, 33, 10, 10]))

正如我们所看到的，除了批量大小这一维度外，其他三个维度都具有不同的尺寸。为了将这两个预测输出链接起来以提高计算效率，我们将把这些张量转换为更一致的格式。

通道维包含中心相同的锚框的预测结果。我们首先将通道维移到最后一维（这样下来，每个像素的预测值在展平 flatten之后是一个连续值，即对像素 flatten；也就是把同类别的锚框放到一起）（如果不 permute的话就是一个类别后面跟着所有像素）。因为不同尺度下批量大小仍保持不变，我们可以将预测结果转成二维的（批量大小，高×宽×通道数）的格式，以方便之后在维度1上的连结。

def flatten_pred(pred):
    return torch.flatten(pred.permute(0, 2, 3, 1), start_dim=1)

def concat_preds(preds):
    return torch.cat([flatten_pred(p) for p in preds], dim=1)

这样一来，尽管Y1和Y2在通道数、高度和宽度方面具有不同的大小，我们仍然可以在同一个小批量的两个不同尺度上连接这两个预测输出。

concat_preds([Y1, Y2]).shape

# 输出
torch.Size([2, 25300])

高宽减半块

为了在多个尺度下检测目标，我们在下面定义了高和宽减半块down_sample_blk，该模块将输入特征图的高度和宽度减半。事实上，该块应用了在 subsec_vgg-blocks中的VGG模块设计。更具体地说，每个高和宽减半块由两个填充为 1 的 3×3 的卷积层、以及步幅为 2 的 2×2 最大汇聚层组成。我们知道填充为 1 的 3×3 卷积层不改变特征图的形状。但是，其后的 2×2 的最大汇聚层将输入特征图的高度和宽度减少了一半。对于此高和宽减半块的输入和输出特征图，因为 1×2+(3−1)+(3−1)=6，所以输出中的每个单元在输入上都有一个6×6的感受野。因此，高和宽减半块会扩大每个单元在其输出特征图中的感受野。

def down_sample_blk(in_channels, out_channels):
    blk = []
    for _ in range(2):
        blk.append(nn.Conv2d(in_channels, out_channels,
                             kernel_size=3, padding=1))
        blk.append(nn.BatchNorm2d(out_channels))
        blk.append(nn.ReLU())
        in_channels = out_channels
    blk.append(nn.MaxPool2d(2))  # 窗口为2的时候默认stride也为2
    return nn.Sequential(*blk)

基本网络块

基本网络块用于从输入图像中抽取特征。

为了计算简洁，我们构造了一个小的基础网络，该网络串联3个高和宽减半块，并逐步将通道数翻倍。给定输入图像的形状为 256×256，此基本网络块输出的特征图形状为 32×32（256/23=32）。

def base_net():
    blk = []
    num_filters = [3, 16, 32, 64]
    for i in range(len(num_filters) - 1):
        blk.append(down_sample_blk(num_filters[i], num_filters[i+1]))
    return nn.Sequential(*blk)

forward(torch.zeros((2, 3, 256, 256)), base_net()).shape

# 输出
torch.Size([2, 64, 32, 32])

完整模型

完整的SSD模型由五个模块组成。每个块生成的特征图既用于生成锚框，又用于预测这些锚框的类别和偏移量。

在这五个模块中，第一个是基本网络块，第二个到第四个是高和宽减半块，最后一个模块使用全局最大池将高度和宽度都降到1。

def get_blk(i):
    if i == 0:
        blk = base_net()
    elif i == 1:
        blk = down_sample_blk(64, 128)
    elif i == 4:
        blk = nn.AdaptiveMaxPool2d((1,1))
    else:
        blk = down_sample_blk(128, 128)
    return blk

现在我们为每个块定义前向传播

与图像分类任务不同，此处的输出包括：CNN特征图 Y；在当前尺度下根据 Y生成的锚框；预测的这些锚框的类别和偏移量（基于Y）。

def blk_forward(X, blk, size, ratio, cls_predictor, bbox_predictor):
    Y = blk(X)
    anchors = d2l.multibox_prior(Y, sizes=size, ratios=ratio)
    cls_preds = cls_predictor(Y)
    bbox_preds = bbox_predictor(Y)
    return (Y, anchors, cls_preds, bbox_preds)

超参数

回想一下，一个较接近顶部的多尺度特征块是用于检测较大目标的，因此需要生成更大的锚框。在上面的前向传播中，在每个多尺度特征块上，我们通过调用的multibox_prior函数的sizes参数传递两个比例值的列表。

在下面，0.2和1.05之间的区间被均匀分成五个部分，以确定五个模块的在不同尺度下的较小值：0.2、0.37、0.54、0.71和0.88。之后，他们较大的值由 $\sqrt{0.2*0.37}$ =0.272、 $\sqrt{0.37*0.54}$ =0.447等给出。

sizes = [[0.2, 0.272], [0.37, 0.447], [0.54, 0.619], [0.71, 0.79],
         [0.88, 0.961]]
ratios = [[1, 2, 0.5]] * 5    # 不用太关心 一个常用组合
num_anchors = len(sizes[0]) + len(ratios[0]) - 1    # 3 + 1 - 1 = 4，以一个像素为中心生成4个锚框

定义完整模型

现在，我们就可以按如下方式定义完整的模型TinySSD了。

class TinySSD(nn.Module):
    def __init__(self, num_classes, **kwargs):
        super(TinySSD, self).__init__(**kwargs)
        self.num_classes = num_classes
        idx_to_in_channels = [64, 128, 128, 128, 128]
        for i in range(5):
            # setattr() 函数对应函数 getattr()，用于设置属性值，该属性不一定是存在的
            # 若不存在会创建并赋值
            # 即赋值语句self.blk_i=get_blk(i)
            setattr(self, f'blk_{i}', get_blk(i))
            setattr(self, f'cls_{i}', cls_predictor(idx_to_in_channels[i],
                                                    num_anchors, num_classes))
            setattr(self, f'bbox_{i}', bbox_predictor(idx_to_in_channels[i],
                                                      num_anchors))

    def forward(self, X):
        anchors, cls_preds, bbox_preds = [None] * 5, [None] * 5, [None] * 5
        for i in range(5):
            # getattr(self,'blk_%d'%i)即访问self.blk_i
            X, anchors[i], cls_preds[i], bbox_preds[i] = blk_forward(
                X, getattr(self, f'blk_{i}'), sizes[i], ratios[i],
                getattr(self, f'cls_{i}'), getattr(self, f'bbox_{i}'))
        anchors = torch.cat(anchors, dim=1)
        cls_preds = concat_preds(cls_preds)
        cls_preds = cls_preds.reshape(
            cls_preds.shape[0], -1, self.num_classes + 1)
        bbox_preds = concat_preds(bbox_preds)
        return anchors, cls_preds, bbox_preds

我们创建一个模型实例，然后使用它对一个256×256像素的小批量图像X执行前向传播。

如本节前面部分所示，第一个模块输出特征图的形状为32×32。回想一下，第二到第四个模块为高和宽减半块，第五个模块为全局汇聚层。由于以特征图的每个单元为中心有4个锚框生成，因此在所有五个尺度下，每个图像总共生成(322+162+82+42+1)×4=5444个锚框。

net = TinySSD(num_classes=1)
X = torch.zeros((32, 3, 256, 256))
anchors, cls_preds, bbox_preds = net(X)

print('output anchors:', anchors.shape)
print('output class preds:', cls_preds.shape)
print('output bbox preds:', bbox_preds.shape)

# 输出
output anchors: torch.Size([1, 5444, 4])
output class preds: torch.Size([32, 5444, 2])
output bbox preds: torch.Size([32, 21776])

训练模型

现在，我们将描述如何训练用于目标检测的单发多框检测模型

读取数据集和初始化

这里读取的是之前博文里讲过的香蕉检测数据集

batch_size = 32
train_iter, _ = d2l.load_data_bananas(batch_size)

# 输出
read 1000 training examples
read 100 validation examples

香蕉检测数据集中，目标的类别数为1。定义好模型后，我们需要初始化其参数并定义优化算法

device, net = d2l.try_gpu(), TinySSD(num_classes=1)
trainer = torch.optim.SGD(net.parameters(), lr=0.2, weight_decay=5e-4)

定义损失函数和评价函数

目标检测有两种类型的损失。

第一种有关锚框类别的损失：我们可以简单地复用之前图像分类问题里一直使用的交叉熵损失函数来计算；

第二种有关正类锚框偏移量的损失：预测偏移量是一个回归问题。但是，对于这个回归问题，我们在这里不使用平方损失（L2范数损失）（因为很有可能我的预测会和正确的数特别远，那这个时候做平方就会变得特别远），而是使用L1范数损失，即预测值和真实值之差的绝对值（这样我们预测特别不靠谱的时候数值也不会特别大）。

掩码变量 bbox_masks令负类锚框和填充锚框不参与损失的计算。最后，我们将锚框类别和偏移量的损失相加，以获得模型的最终损失函数。

cls_loss = nn.CrossEntropyLoss(reduction='none')
bbox_loss = nn.L1Loss(reduction='none')

def calc_loss(cls_preds, cls_labels, bbox_preds, bbox_labels, bbox_masks):
    batch_size, num_classes = cls_preds.shape[0], cls_preds.shape[2]
    cls = cls_loss(cls_preds.reshape(-1, num_classes),
                   cls_labels.reshape(-1)).reshape(batch_size, -1).mean(dim=1)
    bbox = bbox_loss(bbox_preds * bbox_masks,
                     bbox_labels * bbox_masks).mean(dim=1)
    return cls + bbox