Skip to content

单发多框检测(SSD)

一、什么是单发多框检测(SSD)

SSD 是一种常用的目标检测模型,它的特点是 “单发”,也就是只用一个神经网络就可以完成目标检测,不需要像 R-CNN 系列那样分多个步骤。它可以在不同的尺度下生成锚框,然后预测每个锚框的类别和偏移量,从而检测不同大小的目标。

二、SSD 模型的结构

SSD 模型主要由两部分组成:

  1. 基础网络:用来从输入图像中提取特征,通常用深度卷积神经网络,比如 VGG、ResNet 等。基础网络输出的特征图尺寸比较大,用来生成小锚框,检测小目标。

  2. 多尺度特征块:在基础网络之后,有多个多尺度特征块,每个特征块会把上一层的特征图尺寸缩小(比如减半),同时每个单元的感受野更大,用来生成大锚框,检测大目标。

alt text

单发多框检测模型主要由一个基础网络块和若干多尺度特征块串联而成

1. 类别预测层

对于每个尺度的特征图,我们需要预测每个锚框的类别。假设目标类别有 q 个,那么每个锚框有 q+1 个类别(0 类是背景)。

类别预测层用一个 3×3 的卷积层,保持输入和输出的尺寸一样,这样输出特征图的每个位置对应输入特征图的每个位置,每个位置的通道数是 a*(q+1),其中 a 是每个位置生成的锚框数量,q 是目标类别数。

2. 边界框预测层

边界框预测层和类别预测层类似,但是每个锚框需要预测 4 个偏移量(x1, y1, x2, y2),所以输出通道数是 a*4。

3. 多尺度预测的连接

不同尺度的特征图输出的预测结果形状不一样,我们需要把这些结果连接起来,形成最终的预测结果。

三、SSD 模型的原理

SSD 是一个多尺度目标检测模型,它的原理是:

  1. 在不同尺度的特征图上生成不同大小的锚框,小特征图(感受野大)生成大锚框,检测大目标;大特征图(感受野小)生成小锚框,检测小目标。

  2. 对于每个锚框,预测它的类别(是背景还是某个目标类别)和偏移量(和真实边界框的差距)。

  3. 训练的时候,根据真实边界框给每个锚框标注类别和偏移量,然后计算损失,更新模型参数。

四、SSD 模型的实现(PyTorch 版)

1. 定义类别预测层和边界框预测层

python
import torch
from torch import nn
from d2l import torch as d2l

def cls_predictor(num_inputs, num_anchors, num_classes):
    """类别预测层"""
    return nn.Conv2d(num_inputs, num_anchors * (num_classes + 1), kernel_size=3, padding=1)

def bbox_predictor(num_inputs, num_anchors):
    """边界框预测层"""
    return nn.Conv2d(num_inputs, num_anchors * 4, kernel_size=3, padding=1)

2. 定义多尺度特征块

我们可以定义一个下采样块,用来缩小特征图的尺寸:

python
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))
    return nn.Sequential(*blk)

3. 定义完整的 SSD 模型

我们可以定义一个简单的 SSD 模型,包含基础网络和多尺度特征块:

python
class TinySSD(nn.Module):
    """简单的SSD模型"""
    def __init__(self, num_classes, **kwargs):
        super(TinySSD, self).__init__(**kwargs)
        self.num_classes = num_classes
        # 基础网络
        self.blk_0 = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )
        # 多尺度特征块
        self.blk_1 = down_sample_blk(64, 128)
        self.blk_2 = down_sample_blk(128, 128)
        self.blk_3 = down_sample_blk(128, 128)
        self.blk_4 = nn.AdaptiveMaxPool2d((1,1))
        # 类别预测层
        self.cls_0 = cls_predictor(64, 4, num_classes)
        self.cls_1 = cls_predictor(128, 4, num_classes)
        self.cls_2 = cls_predictor(128, 4, num_classes)
        self.cls_3 = cls_predictor(128, 4, num_classes)
        self.cls_4 = cls_predictor(128, 4, num_classes)
        # 边界框预测层
        self.bbox_0 = bbox_predictor(64, 4)
        self.bbox_1 = bbox_predictor(128, 4)
        self.bbox_2 = bbox_predictor(128, 4)
        self.bbox_3 = bbox_predictor(128, 4)
        self.bbox_4 = bbox_predictor(128, 4)

    def forward(self, X):
        anchors, cls_preds, bbox_preds = [None] * 5, [None] * 5, [None] * 5
        # 基础网络
        X = self.blk_0(X)
        anchors[0] = d2l.multibox_prior(X, sizes=[0.2, 0.272], ratios=[1, 2, 0.5])
        cls_preds[0] = self.cls_0(X)
        bbox_preds[0] = self.bbox_0(X)
        # 多尺度特征块
        for i in range(1, 4):
            X = getattr(self, f'blk_{i}')(X)
            anchors[i] = d2l.multibox_prior(X, sizes=[0.37, 0.447], ratios=[1, 2, 0.5])
            cls_preds[i] = getattr(self, f'cls_{i}')(X)
            bbox_preds[i] = getattr(self, f'bbox_{i}')(X)
        # 最后一个特征块
        X = self.blk_4(X)
        anchors[4] = d2l.multibox_prior(X, sizes=[0.88, 1.05], ratios=[1, 2, 0.5])
        cls_preds[4] = self.cls_4(X)
        bbox_preds[4] = self.bbox_4(X)

        # 连接所有预测结果
        anchors = torch.cat(anchors, dim=1)
        cls_preds = self.concat_preds(cls_preds)
        cls_preds = cls_preds.reshape(cls_preds.shape[0], -1, self.num_classes + 1)
        bbox_preds = self.concat_preds(bbox_preds)
        return anchors, cls_preds, bbox_preds

    def concat_preds(self, preds):
        """连接不同尺度的预测结果"""
        return torch.cat([pred.permute(0, 2, 3, 1).reshape(pred.shape[0], -1) for pred in preds], dim=1)

4. 测试模型

我们可以创建一个模型实例,测试一下输入输出的形状:

python
net = TinySSD(num_classes=1)
X = torch.zeros((32, 3, 256, 256))
anchors, cls_preds, bbox_preds = net(X)
print('锚框形状:', anchors.shape)
print('类别预测形状:', cls_preds.shape)
print('边界框预测形状:', bbox_preds.shape)

运行结果:

Plain
锚框形状: torch.Size([1, 5444, 4])
类别预测形状: torch.Size([32, 5444, 2])
边界框预测形状: torch.Size([32, 21776])

四、SSD 模型的训练

1. 读取数据集

我们用香蕉检测数据集来训练模型,这个数据集只有 1 个目标类别(香蕉)。

python
batch_size = 32
train_iter, val_iter = d2l.load_data_bananas(batch_size)

2. 初始化模型和优化器

python
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
net = TinySSD(num_classes=1).to(device)
optimizer = torch.optim.SGD(net.parameters(), lr=0.2, weight_decay=5e-4)

3. 定义损失函数和评价函数

SSD 的损失有两部分:

  1. 类别损失:用交叉熵损失函数,计算锚框类别预测和真实类别的损失。

  2. 边界框损失:用 L1 范数损失(也叫绝对误差损失),计算正类锚框偏移量预测和真实偏移量的损失。

python
def calc_loss(cls_preds, cls_labels, bbox_preds, bbox_labels, bbox_masks):
    """计算损失"""
    # 类别损失
    cls_loss = nn.CrossEntropyLoss(reduction='none')(cls_preds, cls_labels)
    cls_loss = cls_loss.sum()
    # 边界框损失
    bbox_loss = nn.L1Loss(reduction='none')(bbox_preds * bbox_masks, bbox_labels * bbox_masks)
    bbox_loss = bbox_loss.sum()
    return cls_loss + bbox_loss

def cls_eval(cls_preds, cls_labels):
    """评价类别预测的准确率"""
    return float((cls_preds.argmax(dim=-1) == cls_labels).sum())

def bbox_eval(bbox_preds, bbox_labels, bbox_masks):
    """评价边界框预测的平均绝对误差"""
    return float((torch.abs((bbox_preds - bbox_labels) * bbox_masks)).sum())

4. 训练模型

python
num_epochs = 20
animator = d2l.Animator(xlabel='epoch', xlim=[1, num_epochs], legend=['class error', 'bbox mae'])
net.train()
for epoch in range(num_epochs):
    metric = d2l.Accumulator(4)
    for features, target in train_iter:
        X = features.to(device)
        Y = target.to(device)
        optimizer.zero_grad()
        # 前向传播
        anchors, cls_preds, bbox_preds = net(X)
        # 标注锚框
        bbox_labels, bbox_masks, cls_labels = d2l.multibox_target(anchors, Y)
        # 计算损失
        l = calc_loss(cls_preds, cls_labels, bbox_preds, bbox_labels, bbox_masks)
        l.backward()
        optimizer.step()
        # 记录指标
        metric.add(cls_eval(cls_preds, cls_labels), cls_labels.numel(),
                   bbox_eval(bbox_preds, bbox_labels, bbox_masks), bbox_labels.numel())
    # 计算误差
    cls_err = 1 - metric[0] / metric[1]
    bbox_mae = metric[2] / metric[3]
    animator.add(epoch + 1, (cls_err, bbox_mae))
    print(f'epoch {epoch+1}, class error {cls_err:.2e}, bbox mae {bbox_mae:.2e}')

五、SSD 模型的预测

训练完成后,我们可以用模型来预测图片里的目标:

python
def predict(X):
    """预测图片里的目标"""
    net.eval()
    anchors, cls_preds, bbox_preds = net(X.to(device))
    cls_probs = nn.Softmax(dim=-1)(cls_preds).permute(0, 2, 1)
    output = d2l.multibox_detection(cls_probs, bbox_preds, anchors)
    idx = [i for i, row in enumerate(output[0]) if row[0] != -1]
    return output[0, idx]

# 加载图片
img = d2l.plt.imread('../img/catdog.jpg')
h, w = img.shape[:2]
X = torch.from_numpy(d2l.resize(img, (256, 256))).permute(2, 0, 1).float()
X = X.unsqueeze(0)
# 预测
output = predict(X)
# 显示预测结果
fig = d2l.plt.imshow(img)
for row in output:
    score = float(row[1])
    if score < 0.9:
        continue
    hh, ww = img.shape[0:2]
    bbox = [row[2:6] * torch.tensor((w, h, w, h), device=row.device)]
    d2l.show_bboxes(fig.axes, bbox, f'banana {score:.2f}', 'w')

六、小结

  1. SSD 是一种单发多框检测模型,只用一个神经网络就可以完成目标检测,它通过多尺度特征块生成不同大小的锚框,检测不同大小的目标。

  2. SSD 的模型结构包括基础网络和多尺度特征块,每个特征块都有类别预测层和边界框预测层。

  3. SSD 的损失包括类别损失和边界框损失,类别损失用交叉熵损失,边界框损失用 L1 范数损失(也叫绝对误差损失)。

  4. SSD 的训练和预测都比较简单,适合快速测试和部署。

(注:文档部分内容可能由 AI 生成) 源地址

京ICP备2024093538号-1