ssd网络结构_SSD 源码实现详细解析 (PyTorch)-模型定义篇

# ssd.py
class SSD(nn.Module): # 自定义SSD网络
    def __init__(self, phase, size, base, extras, head, num_classes):
        # ... SSD 模型初始化
    def forward(self, x):
        # ... 定义forward函数, 将设计好的layers和ops应用到输入图片 x 上
 
    def load_weights(self, base_file):
        # ... 加载参数权重值
def vgg(cfg, i, batch_norm=False):
    # ... 搭建vgg网络
def add_extras(cfg, i, batch_norm=False):
    # ... 向VGG网络中添加额外的层用于feature scaling
def multibox(vgg, extra_layers, cfg, num_classes):
    # ... 构建multibox结构
base = {
    ...} # vgg 网络结构参数
extras = {
    ...} # extras 层参数
mbox = {
    ...} # multibox 相关参数
def build_ssd(phase, size=300, num_classes=21):
    # ... 构建模型函数, 调用上面的函数进行构建

# ssd.py
class SSD(nn.Module): # 自定义SSD网络
    def __init__(self, phase, size, base, extras, head, num_classes):
        # ...
    def forward(self, x):
        # ...
    def load_weights(self, base_file):
        # ...
def vgg(cfg, i, batch_norm=False):
    # ... 搭建vgg网络
def add_extras(cfg, i, batch_norm=False):
    # ... 向VGG网络中添加额外的层用于feature scaling
def multibox(vgg, extra_layers, cfg, num_classes):
    # ... 构建multibox结构
base = {
     # vgg 网络结构参数
    '300': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'C', 512, 512, 512, 'M', 512, 512, 512],
    '500': []
}
extras = {
     # extras 层参数
    '300': [256, 'S', 512, 128, 'S', 256, 128, 256, 128, 256],
    '500': []
}
mbox = {
     # multibox 相关参数
    '300': [4, 6, 6, 6, 4, 4],
    '500': []
}
def build_ssd(phase, size=300, num_classes=21):
    # 构建模型函数, 调用上面的函数进行构建
    if phase != "test" and phase != "train": # 只能是训练或者预测阶段
        print("ERROR: Phase: " + phase + " not recognized")
        return
    if size != 300:
        print("ERROR: You specified size " + repr(size) + ". However, "+
                "currently only SSD300 is supported!") # 仅仅支持300size的SSD
        return
    base_, extras_, head_ = multibox(vgg(base[str(size)], 3),
            add_extras(extras[str(size), 1024),
            mbox[str(size)], num_classes )
    return SSD(phase, size, base_, extras_, head_, num_classes)

# ssd.py
 
def vgg(cfg, i, batch_norm = False):
    # cfg = base['300'] = [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'C', 512, 512, 512, 'M', 512, 512, 512],
    # i = 3
    layers = []
    in_channels = i
    for v in cfg:
        if v == 'M':
            layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
        if v == 'C':
            layers += [nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True)]
        else:
            conv2d = nn.Conv2d(in_channels=in_channels, out_channels=v, kernel_size=3, padding=1)
            if batch_norm:
                layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]
            else:
                layers += [conv2d, nn.ReLU(inplace=True)]
            in_channels = v
        pool5 = nn.MaxPool2d(kernel_size=3, stride=1, padding=1)
        conv6 = nn.Conv2d(512, 1024, kernel_size=3, padding=6, dilation=6)
        conv7 = nn.Con2d(1024, 1024, kernel_size=1)
        layers += [pool5, conv6, nn.ReLU(inplace=True), conv7, nn.ReLU(inplace=True)]
        return layers

# ssd.py
def add_extras(cfg, i, batch_norm=False):
    # cfg = [256, 'S', 512, 128, 'S', 256, 128, 256, 128, 256]
    # i = 1024
    layers = []
    in_channels = i
    flag = False
    for k, v in enumerate(cfg):
        if in_channels != 'S':
            if v == 'S': # (1,3)[True] = 3, (1,3)[False] = 1
                layers += [nn.Conv2d(in_channels=in_channels, out_channels=cfg[k+1],
                                    kernel_size=(1, 3)[flag], stride=2, padding=1)]
            else:
                layers += [nn.Conv2d(in_channels=in_channels, out_channels=v,
                                    kernel_size=(1, 3)[flag])]
            flag = not flag
        in_channels = v
    return layers

def add_extras():
    exts1_1 = nn.Conv2d(in_channels=1024, out_channels=256, kernel_size=1)
    exts1_2 = nn.Conv2d(in_channels=256, out_channels=512, kernel_size=3, stride=2, padding=1)
    exts2_1 = nn.Conv2d(512, 128, 1, 1, 0)
    exts2_2 = nn.Conv2d(128, 256, 3, 2, 1)
    exts3_1 = nn.Conv2d(256, 128, 1, 1, 0)
    exts3_2 = nn.Conv2d(128, 256, 3, 1, 0)
    exts4_1 = nn.Conv2d(256, 128, 1, 1, 0)
    exts4_2 = nn.Conv2d(128, 256, 3, 1, 0)
 
    return [exts1_1, exts1_2, exts2_1, exts2_2, exts3_1, exts3_2, exts4_1, exts4_2]

# ssd.py
def multibox(vgg, extra_layers, cfg, num_classes):
    # cfg = [4, 6, 6, 6, 4, 4]
    # num_classes = 21
    # ssd总共会选择6个卷积特征图谱进行预测, 分别为, vggnet的conv4_3, 以及extras_layers的5段卷积的输出(每段由两个卷积层组成, 具体可看extras_layers的实现).
    # 也就是说, loc_layers 和 conf_layers 分别具有6个预测层.
    loc_layers = []
    conf_layers = []
    vgg_source = [21, -2]
    for k, v in enumerate(vgg_source):
        loc_layers += [nn.Conv2d(vgg[v].out_channels, cfg[k]*4, kernel_size=3, padding=1]
        conf_layers += [nn.Conv2d(vgg[v].out_channels, cfg[k]*num_classes, kernel_size=3, padding=1)]
    for k, v in enumerate(extra_layers[1::2], 2):
        loc_layers += [nn.Conv2d(v.out_channels, cfg[k]*4, kernel_size=3, padding=1)]
        conf_layers += [nn.Conv2d(v.out_channels, cfg[k]*num_classes, kernel_size=3, padding=1)]
    return vgg, extra_layers, (loc_layers, conf_layers)

# ssd.py
def multibox(vgg, extras, num_classes):
    loc_layers = []
    conf_layers = []
    #vgg_source=[21, -2] # 21 denote conv4_3, -2 denote conv7
 
    # 定义6个坐标预测层, 输出的通道数就是每个像素点上会产生的 default box 的数量
    loc1 = nn.Conv2d(vgg[21].out_channels, 4*4, 3, 1, 1) # 利用conv4_3的特征图谱, 也就是 vgg 网络 List 中的第 21 个元素的输出(注意不是第21层, 因为这中间还包含了不带参数的池化层).
    loc2 = nn.Conv2d(vgg[-2].out_channels, 6*4, 3, 1, 1) # Conv7
    loc3 = nn.Conv2d(vg