SSD目标检测

SSD目标检测

本次所使用的数据集是香蕉数据集，当然是参考
https://zh.d2l.ai/chapter_computer-vision/ssd.html
动手学深度学习这本书的

首先是香蕉数据集的读取，这个我没看明白，但应该有他自己的逻辑，其实不必太纠结，我们只需要能把数据读进来就可以了
target 有五个值，标签和框左上角以及右下角的坐标

import os 
import pandas as pd
import torch
import torchvision
from d2l import torch as d2l
from torch import nn
from torch.nn import functional as F
import numpy as np
#@save
d2l.DATA_HUB['banana-detection'] = (
    d2l.DATA_URL + 'banana-detection.zip',
    '5de26c8fce5ccdea9f91267273464dc968d20d72')
def read_data_bananas(is_train=True):
    data_dir = d2l.download_extract("banana-detection")
    csv_fname = os.path.join(data_dir, "bananas_train" if is_train else "bananas_val", "label.csv")
    csv_data = pd.read_csv(csv_fname)
    csv_data = csv_data.set_index("img_name")
    images, targets = [], []
    for img_name, target in csv_data.iterrows():
        images.append(torchvision.io.read_image(
        os.path.join(data_dir, "bananas_train" if is_train else "bananas_val", "images", f'{img_name}')))
        targets.append(list(target))
    # target 有五个值，标签和框的坐标，每个图像都有香蕉
    return images, torch.tensor(targets).unsqueeze(1) / 256

接着我们要创建一个Dataset类，以便我们可以用DataLoader来读，这里就没什么好说的

#  read_data_bananas 读图像和标签，创建BananasDataset类来创建一个自定义Dataset类来加载香蕉数据集
class BananasDataset(torch.utils.data.Dataset):
    def __init__(self, is_train):
        self.features, self.labels = read_data_bananas(is_train)
        print('read' + str(len(self.features)) + (f'training examples' if is_train else 'validation examples'))

    def __getitem__(self, idx):
        return (self.features[idx].float(), self.labels[idx])
    
    def __len__(self):
        return len(self.features)

最后创建一个函数使用DataLoader来读取数据

def load_data_bananas(batch_size):
    train_iter = torch.utils.data.DataLoader(BananasDataset(is_train=True), batch_size, shuffle=True)
    val_iter = torch.utils.data.DataLoader(BananasDataset(is_train=False), batch_size)
    return train_iter, val_iter

接着我们就可以读一下数据啦

batch_size, edge_size = 32, 256
train_iter, _ = load_data_bananas(batch_size)
batch = next(iter(train_iter))
batch[0].shape, batch[1].shape #标签第二个索引含义是所有图像中可能出现的最大边框数，在这里就是１

(torch.Size([32, 3, 256, 256]), torch.Size([32, 1, 5]))
第一个图像的就不说了，第二个是target，他的第二个索引1代表所有图像中可能出现的最大边框数，这里为1

接着展示一下图像，需要注意的是permute代表着换维度，比如原来我们的shape是(32,3,256,256),permute(0,2,3,1)就是把二、三维度的数换到1、2维度，相当于把通道数放在最后了，变为(32,256,256,3)

imgs = (batch[0][0:10].permute(0,2,3,1)) / 255 # 相当于原来的第二、三维度换到1，2维度
axes = d2l.show_images(imgs, 2, 5, scale = 2)
for ax, label in zip(axes, batch[1][0:10]):
    #print(ax)
    d2l.show_bboxes(ax, [label[0][1:5] * edge_size], colors=['w'])

接下来，我们就可以开始真正的关于目标检测SSD的东西了，首先是类别预测层
假设目标类别的数量为q，那么锚框的类别就有q+1个，0为背景，设特征图的高和宽分别为w和h（不管他有多少层，因为生成锚框只看特征图的w和h），假设以每个像素点生成锚框的个数为a，那么需要对hwa个锚框进行分类，如果使用全连接层进行输出，参数太多了。可以用卷积来预测
我们只需要使用一个能保持输入高和宽的卷积层，输出通道数为a(q+1)，这样比如在(x,y)像素点上，i(q+1)+j的通道就代表索引为i的锚框预测成索引为j类别的概率。
下面，我们就定义一个这样的类别预测层，num_anchors和num_classes分别为a和q
PS：代码的注释是我当时写的时候的想法，以上是我写文章时候的想法，大家可以结合的看，可能会更明白一点。

# 类别预测层
# 之前说过，假设目标类别的数量为q，那么锚框有q+1个类别，其中0是背景，每个像素中心生产a个锚框
# 如果一张特征图高宽是h,w，以这张特征图生成原图像的锚框的话，我们需要
# 对hwa个锚框进行分类，用全连接层很容易导致参数过多，于是还是采用卷积通道的方法，输出的卷积层通道数为a(q+1)，在点x,y上，索引为i(q+1)+j
# 为索引为i的锚框对类别索引j的预测
def cls_predictor(num_inputs, num_anchors, num_classes):
    return nn.Conv2d(num_inputs, num_anchors*(num_classes + 1), kernel_size=3, padding=1)

边界框偏移量预测层，这个跟上面一样的，就是他是每个像素点需要预测a4个偏移量，所以输出通道是a4

# 边界框预测层
# 跟类别预测出类似，唯一不同的是，他需要预测的是a*4，即每个中心每个框的4个偏移量
def bbox_predictor(num_inputs, num_anchors):
    return nn.Conv2d(num_inputs, num_anchors*4, kernel_size=3, padding=1)

很明显可以看出，如果特征图的形状不同的话，那么输出的通道数和形状肯定不同，例如

# 不同形状的特征图所输出的预测形状不同
# 两个批次
def forward(x, block):
    return block(x)
# 假设Y1，Y2每个像素点分别生成5，3个锚框，目标类别数量为10
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

为了将两个预测输出（无论是预测类别还是偏移量）连接起来，就是把后面三维给堆叠起来，变成(batch_size,hwchannels)的形状，为什么要把通道数放在最后面，我是这样理解的，正常你展开的话，连在一起的会是一张图片上每个像素点的其中一个锚框预测为某一类的概率，把通道数放在最后面的话，连在一起的就是一个像素点分别在不同类别上预测的概率

# 可以看到除了第一维，其他的都不一样，为了方便将输出连接起来，我们将预测结果转为二维的(batch_size, 高*宽*通道数)，为什么需要把通道数变到最后一维
#  因为如果正常flatten的话，连在一起的会是一张图片每个像素点上的一个锚框预测某一类的概率，变化完之后，连在一起的就是一个像素点分别在
# 不同类别上预测的概率
def flatten_pred(pred):
    return torch.flatten(pred.permute(0, 2, 3, 1), start_dim=1)
def concat_preds(pred):
    return torch.cat([flatten_pred(p) for p in pred], dim=1)
concat_preds([Y1, Y2]).shape

说完这些后，接下来就要分析一下SSD的每个模块了
高宽减半块
每个高宽减半块由两个填充为1的卷积层，以及一个最大池化层组成，当然每个卷积层后面跟着BN和激活函数

# 高宽减半块，这个好理解，减半的主要是靠maxpooling，前面加了两个卷积
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)

基本网络块
基本网络块由三个高宽减半块构成，并逐步将通道数翻倍

# 定义基本网络块， 经过三个高宽减半块，并逐步将通道数翻倍
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)

完整模型
完整的模型由五个模块组成，每个块既生成特征图，又生成锚框，又用于预测这些锚框的类别和偏移量，第一个模块是基本网络块，后面三个是高宽减半块，最后一个模块用的是全局最大池化，将高宽缩为1

# 完整的模型，单发多框又五个模块组成，每个块生成的特征图既用于锚框，又用于预测这些锚框的类别和偏移量，第一个是网络基本块，第二到第四个是
# 高宽减半块，最后一个是全局最大池化，将高度和宽度都降到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

现在为每个块定义前向传播，与图像分类不同，此处的输出不仅包括特征图Y，还包括根据Y生成的锚框，基于预测这些锚框的类别和偏移量（基于Y）

# 现在为每个块定义前向传播，与图像分类不同，此处的输出不仅包括特征图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)

一个较接近顶部的多尺度特征块是用于检测较大目标的，因此，需要生成更大的锚框，即每个特征图生成的锚框的size都不一样

# 一个较接近顶部的多尺度特征块是用于检测较大目标的，因此需要更大的锚框
# 五个模块的锚框size都不一样
sizes = [[0.2, 0.272], [0.37, 0.447], [0.54, 0.619], [0.71, 0.79], [0.88, 0.961]] # 0.272 = sqrt(0.2, 0.37)
ratios = [[1, 2, 0.5]] * 5
num_anchors = len(sizes[0]) + len(ratios[0]) - 1

接下来就可以定义完整的TinyYOLO了，一些不太好懂的我就写了注释了

# 定义完整TinySSD(nn.Block)
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] # 每一模块输出的通道数
        # setattr(object, name, value)
        for i in range(5):
            setattr(self, f'blk_{i}', get_blk(i)) # 相当于self.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, f'blk_{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, axis=1)
        cls_preds = concat_preds(cls_preds)
        cls_preds = cls_preds.reshape(cls_preds.shape[0], -1, self.num_classes + 1) # 变一下形状，第一维是batch，第二维就是依次按照每个特征图上的锚框排了
        bbox_preds = concat_preds(bbox_preds)  # 上面变一下估计是为了后面好算，这个变也行感觉，就看后面代码怎么写了
        return anchors, cls_preds, bbox_preds

可以测试一下，输入一个256*256的图像，batch_size=32

net = TinySSD(num_classes=1)
X = torch.zeros((32, 3, 256, 256))  # batch_size = 32, 3通道，图像256*256
anchors, cls_preds, bbox_preds = net(X)
print(anchors.shape)       # 框的数量，第三维是位置，32个图像的锚框大小都是一样的
print(cls_preds.shape)     # 为所有框预测类别
print(bbox_preds.shape)    # 预测位置

最后就开始训练模型啦
读取数据集

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

然后定义一下loss，分别用类别的Loss和偏移量loss，类别就用交叉熵就好了，偏移量就用的是绝对误差
因为cls_preds的shape是(batch,锚框个数，类别个数嘛)，然后先把他reshape成（batch锚框个数，类别个数），同样的把cls_labels也弄成这样(cls_labels的shape是(batch，锚框个数))，想用交叉熵损失的话就得这样，最后再reshape回去，bbox可以不这样做，因为他本来就是(batch,锚框个数4)没有第三维

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
net = TinySSD(num_classes=1)
optimizer = torch.optim.SGD(net.parameters(), lr=0.2, weight_decay=5e-4)
cls_loss = nn.CrossEntropyLoss(reduction='none')
bbox_loss = nn.L1Loss(reduction='none')  # 偏移量是L1损失
# 掩码令负类锚框和填充锚框不参与计算
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) # 因为之前通过reshape，batch也被reshape了，要reshape回来
    bbox = bbox_loss(bbox_preds * bbox_masks, bbox_labels * bbox_masks).mean(dim=1)
    return cls + bbox

然后用准确率和平均误差来评价分类以及偏移量标准

# 用准确率和平均误差来评价分类和预测量标准
def cls_eval(cls_preds, cls_labels):
    return float((cls_preds.argmax(dim=-1).type(cls_labels.dtype) == cls_labels).sum())
def bbox_eval(bbox_preds, bbox_labels, bbox_masks):
    return float((torch.abs((bbox_labels - bbox_preds) * bbox_masks)).sum())

之后就可以开始训练我们的模型了，其中的animator应该就是个画图的，我没去深入了解，timer和metric也没了解，就是一个辅助手段把，不影响训练

num_epochs = 20
timer = d2l.Timer()
animator = d2l.Animator(xlabel='epoch', xlim=[1, num_epochs],
                        legend=['class error', 'bbox mae'])
net = net.to(device)
for epoch in range(num_epochs):
    # 训练精确度的和，训练精确度的和中的示例数
    # 绝对误差的和，绝对误差的和中的示例数
    metric = d2l.Accumulator(4)
    net.train()
    for features, target in train_iter:
        timer.start()
        optimizer.zero_grad()
        X, Y = features.to(device), target.to(device)
        # 生成多尺度的锚框，为每个锚框预测类别和偏移量
        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 = calc_loss(cls_preds, cls_labels, bbox_preds, bbox_labels,
                    bbox_masks)
        l.mean().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, bbox_mae = 1 - metric[0] / metric[1], metric[2] / metric[3]
    animator.add(epoch + 1, (cls_err, bbox_mae))
print(f'class err {cls_err:.2e}, bbox mae {bbox_mae:.2e}')
print(f'{len(train_iter.dataset) / timer.stop():.1f} examples/sec on '
      f'{str(device)}')

再最后就是推理预测了，我们先拿一个图像，然后令他的batch_size = 1，当然img是不需要这个1的，并且转化为img的时候记得把通道放在最后，需要注意的是，里面可能会有一些变换维度的操作，那是为了适应自己之前写的函数输入

X = torchvision.io.read_image('banana.jpg').unsqueeze(0).float()
img = X.squeeze(0).permute(1, 2, 0).long()
def predict(X):
    net.eval()
    anchors, cls_preds, bbox_preds = net(X.to(device))
    cls_probs = F.softmax(cls_preds, dim=2).permute(0, 2, 1) # 转化为bn,类别数+1， 锚框个数
    output = d2l.multibox_detection(cls_probs, bbox_preds, anchors)  # bn，锚框个数，每个锚框的信息
    idx = [i for i, row in enumerate(output[0]) if row[0] != -1]  # 锚框信息的第一维是类别
    return output[0, idx]  # bn=0，不是背景的类

output = predict(X)

def display(img, output, threshold):
    d2l.set_figsize((5, 5))
    fig = d2l.plt.imshow(img)
    for row in output:
        score = float(row[1])
        if score < threshold:
            continue
        h,w = img.shape[0:2]
        bbox = [row[2:6] * torch.tensor((w, h, w, h), device = row.device)]
        d2l.show_bboxes(fig.axes, bbox, '%.2f' % score, 'w')
display(img, output.cpu(), threshold=0.5)