学习李沐动手学深度学习中的实现ssd

代码已注释，运行时出现小问题在代码后说明。

import torch
import torchvision
from torch import nn
from torch.nn import functional as F
from d2l import torch as d2l
import matplotlib.pyplot as plt


def cls_predictor(num_inputs, num_anchors, num_classes):
    # 输入通道数：num_inputs,
    # 输出通道数：锚框的个数num_anchors×(类别数num_classes+背景类1),
    # 此时高宽没有归一化，所以卷积核3×3，填充1，保证输出高宽和输入高宽一样。
    return nn.Conv2d(num_inputs,
                     num_anchors * (num_classes + 1),
                     kernel_size=3, padding=1)


def bbox_predictor(num_inputs, num_anchors):
    # 每个锚框4个偏移量值故乘4。
    return nn.Conv2d(num_inputs, num_anchors * 4, kernel_size=3, padding=1)


def forward(x, block):
    # 返回block块输出。
    return block(x)


# 测试张量维度
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))
print(Y1.shape, Y2.shape)
# 结果torch.Size([2, 55, 20, 20]) torch.Size([2, 33, 10, 10])
# 特征图尺度改变除了批量之外，其他都发生变化，55与33是预测输出通道个数。（批量大小，通道数，高度，宽度）


def flatten_pred(pred):
    # 利用permute函数进行换序操作，把通道数放在最后。
    # start_dim=1沿维度1拉成：批量数×(高×宽×通道数)的二维张量，为了下面拼接。
    # 换序操作避免类别预测在flatten后相距较远。
    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)


# 测试沿一维度拼接结果55 * 20 * 20 + 33 * 10 * 10 = 25300,
# 结果：torch.Size([2, 25300])
print(concat_preds([Y1, Y2]).shape)


# 定义一个简单的CNN网络,输入维度in_channels，输出out_channels，高宽减半。
def down_sample_blk(in_channels, out_channels):
    blk = []
    # 卷积，BN层，ReLU激活函数，重复两次。
    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
    # 池化层默认步长等于2，形成高宽减半的效果。
    blk.append(nn.MaxPool2d(2))
    # 放在Sequential中。加星：*args：接收若干个位置参数，转换成元组tuple形式。
    return nn.Sequential(*blk)


print(forward(torch.zeros((2, 3, 20, 20)), down_sample_blk(3, 10)).shape)
# 结果：torch.Size([2, 10, 10, 10])，高和宽减半块会更改输入通道的数量，并将输入特征图的高度和宽度减半。


def base_net():
    blk = []
    num_filters = [3, 16, 32, 64]