SSD源碼解讀——網絡搭建

以前，對SSD的論文進行了解讀，能夠回顧以前的博客：http://www.javashuo.com/article/p-rrxantwu-cm.html。

爲了加深對SSD的理解，所以對SSD的源碼進行了復現，主要參考的github項目是ssd.pytorch。同時，我本身對該項目增長了大量註釋：https://github.com/Dengshunge/mySSD_pytorch

搭建SSD的項目，能夠分紅如下三個部分：

1. 數據讀取；
2. 網絡搭建；
3. 損失函數的構建。
4. 網絡測試。

接下來，本篇博客重點分析網絡搭建。

---

該部分總體比較簡單，思路也很清晰。

首先，在train.py中，網絡搭建的函數入口是函數build_ssd()，該函數須要傳入如下幾個參數："train"或者"test"字符串、圖片尺寸、類別數。其中，"train"或者"test"字符串用於區分該網絡是用於訓練仍是測試，這兩個階段的網絡有些許不一樣，本文主要將訓練階段的網絡；而類別數須要加上背景，對於VOC而言，有20個類別，加上1個背景，即類別數是21。

ssd_net = build_ssd('train', voc['min_dim'], voc['num_classes'])

這裏，先放一張SSD的網絡結構圖，能夠看出，SSD網絡是有3部分組成的，vgg主幹網絡，新增網絡（Conv6以後的層）和用於檢測的頭部網絡（Extra Feature Layers）。

接着，在ssd.py中，首先定了一個參數，以下所示。這裏主要以SSD300爲例。這些參數有什麼用呢？字典base的參數指的是用於搭建VGG主幹網絡輸出通道數，其中「M」表示須要進行maxpooling；字典extras的參數一樣表示新增層的輸出通道數，其中「S」表示須要stride=2的降採樣；字典mbo的參數表示用於特徵融合的層中，每一個層對應未知(x,y)的錨點框數量，在SSD300中，使用了6個層進行特徵融合，如Conv_4層中，每一個位置使用4個錨點框進行預測。

base = {
    '300': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M',
            512, 512, 512],  # M表示maxpolling
    '512': [],
}
extras = {
    '300': [256, 'S', 512, 128, 'S', 256, 128, 256, 128, 256],  # S表示stride=2
    '512': [],
}
mbox = {
    '300': [4, 6, 6, 6, 4, 4],  # 每一個特徵圖的每一個點，對應錨點框的數量
    '512': [],
}

當定義完須要使用到的參數後，能夠進行如具體搭建的環節。函數build_ssd()的定義以下所示。利用函數multibox()來構建SSD網絡的各個部分，分別是VGG主幹網絡，新增層和用於檢測的頭部網絡（或許能夠理解爲分類頭和迴歸頭）。而VGG主幹網絡是經過函數vgg()來實現，新增層是經過函數add_extras()來實現，而函數multibox()則搭建用於檢測的頭部網絡。最後用這些層來初始化類SSD。

def build_ssd(phase, size=300, num_class=21):
    if phase != 'test' and phase != 'train':
        raise ("ERROR: Phase: " + phase + " not recognized")
    base_, extras_, head_ = multibox(vgg(base[str(size)]),
                                     add_extras(extras[str(size)], in_channels=1024),
                                     mbox[str(size)],
                                     num_class)
    return SSD(phase, size, base_, extras_, head_, num_class)

咱們來看一下VGG主幹網絡是如何搭建的。函數vgg()須要將上述的base字典傳入進去，根據base字典，來搭建卷積層和池化層。做者對vgg網絡進行了改進，即將fc6和fc7更改爲conv6和conv7。值得留意的是，在conv6中，使用了空點卷積，dilation=6，增大感覺野。在SSD論文的最後，也討論了空洞卷積對結果有好的影響。最後，將這些卷積層和池化層放入list中，並返回這個list。

def vgg(cfg=base['300'], batch_norm=False):
    '''
    該函數來源於torchvision.models.vgg19()中的make_layers()
    '''
    layers = []
    in_channels = 3

    # vgg主體部分,到論文的conv6以前
    for v in cfg:
        if v == 'M':
            # ceil_mode是向上取整
            layers += [nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True)]
        else:
            conv2d = nn.Conv2d(in_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, ceil_mode=True)

    conv6 = nn.Conv2d(512, 1024, kernel_size=3, padding=6, dilation=6)
    conv7 = nn.Conv2d(1024, 1024, kernel_size=1)
    layers += [pool5, conv6, nn.ReLU(inplace=True), conv7, nn.ReLU(inplace=True)]

    return layers

接下來，咱們來了解一下SSD在vgg中新增層，即conv7以後的網絡層。一樣，函數add_extras()須要傳入字典extras，來構建網絡層。這裏能夠留意一下，kernel_size的寫法，(1,3)爲一個元祖tuple，flag來控制取哪一個值，便可變換使用3*3或者1*1的卷積核，減小代碼的冗餘。最後將新構建的層存入list中，並返回這個list。

def add_extras(cfg=extras['300'], in_channels=1024):
    '''
    完成SSD後半部分的網絡構建，即做者新加上去的網絡,從conv7以後到conv11_2
    '''
    layers = []
    flag = False  # 交替控制卷積核,使用1*1或者使用3*3
    for k, v in enumerate(cfg):
        if in_channels != 'S':
            if v == 'S':
                layers += [nn.Conv2d(in_channels, cfg[k + 1], kernel_size=(1, 3)[flag], stride=2, padding=1)]
            else:
                layers += [nn.Conv2d(in_channels, v, kernel_size=(1, 3)[flag])]
            flag = not flag
        in_channels = v
    return layers

當有了vgg的主幹網絡和新增層後，能夠將某些層進行特徵融合和預測了。這裏，就須要使用到函數multibox()。須要將vgg主幹網絡和新增層的list、字典mbox和類別數傳入函數中。首先，函數multibox()會建立兩個list，用於保存位置迴歸的層和置信度的層。對於每一個用於融合的特徵層，會分紅兩部分，一個用於迴歸，使用3*3的卷積，輸出通道數是cfg[k] * 4，其中cfg[k]表示每一個位置上錨點框的數量，4表示[x_min,y_min,x_max,y_max]；另一個用於類別的判斷，也是使用3*3的卷積，輸出通道數是cfg[k] * num_class，表示每一個錨點框判斷其屬於哪個類別，在voc中，num_class=21（包含背景）。能夠理解成將此特徵層分紅了分類頭和迴歸頭，每一個錨點框會輸出4個座標和21個類別置信度。最後將vgg主幹網絡、新增層、分類頭和迴歸頭返回。

def multibox(vgg, extra_layers, cfg, num_class):
    '''
    返回vgg網絡,新增網絡,位置網絡和置信度網絡
    '''
    loc_layers = []  # 判斷位置
    conf_layers = []  # 判斷置信度

    vgg_source = [21, -2]  # 21表示conv4_3的序號,-2表示conv7的序號
    for k, v in enumerate(vgg_source):
        # vgg[v]表示須要提取的特徵圖
        # cfg[k]表明該特徵圖下每一個點對應的錨點框數量
        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_class, kernel_size=3, padding=1)]

    for k, v in enumerate(extra_layers[1::2], 2):
        # [1::2]表示,從第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_class, kernel_size=3, padding=1)]

    return vgg, extra_layers, (loc_layers, conf_layers)

函數multibox()返回的各個層，用於初始化類SSD。首先，因爲由於「train」階段和「test」階段是有點區別的，本節依然主要將「train」階段，所以，須要傳入phase參數，參數只能是兩個值(train,test)。函數PriorBox()的做用是來建立先驗錨點框，返回的shape爲[8732,4]，其中具備8732個錨點框，4表示每一個錨點框的座標[中心點x，中心點y，寬，高]，這裏的座標值有點不太同樣。因爲傳入的網絡層是以list列表的形式，所以，用nn.ModuleList()將其轉換爲pytorch的網絡結構。

接下來看類SSD中的函數forward()，用於前向推理。按順序對輸入圖片進行處理，在conv4中，須要對特徵圖進行L2正則化。並將用於特徵融合的特徵圖存在放sources中。在獲得5個用於融合的特徵圖後，將這些特徵圖輸入到分類頭和迴歸頭中，每一個特徵圖對應各自的分類頭和迴歸頭。這裏注意一下，分類頭或者回歸頭卷積後，使用了permute()函數。該函數的做用是交換維度，本來的維度是[batch_size,channel,height,weight]，交換維度後變成了[batch_size,height,weight,channel]，這樣作的目的是方便後續的處理。將處理後的結果保存在loc和conf這兩個List中。後續接着對loc和conf進行變換，利用view()函數，最終，loc的shape爲[batch_size,8732*4]，conf的shape爲[batch_size,8732*21]。

最後，將loc和conf這兩個List又變換維度，返回出去，用於計算loss損失函數（感受這麼多變換，有點重複呀，應該能夠省略一部分）。"train"階段和"test"階段返回的結果相似，其中不一樣點是，在test階段，置信度須要通過softmax。

class SSD(nn.Module):
    '''
    構建SSD的主函數，將base(vgg)、新增網絡和位置網絡與置信度網絡組合起來
    '''

    def __init__(self, phase, size, base, extras, head, num_classes):
        super(SSD, self).__init__()
        self.phase = phase
        self.num_classes = num_classes
        self.priors = torch.Tensor(PriorBox(voc))
        self.size = size

        # SSD網絡
        self.vgg = nn.ModuleList(base)
        # 對conv4_3的特徵圖進行L2正則化
        self.L2Norm = L2Norm(512, 20)

        self.extras = nn.ModuleList(extras)
        self.loc = nn.ModuleList(head[0])
        self.conf = nn.ModuleList(head[1])

        if phase == 'test':
            self.softmax = nn.Softmax(dim=-1)
            self.detect = Detect(num_classes=self.num_classes, top_k=200,
                                 conf_thresh=0.01, nms_thresh=0.45)

    def forward(self, x):
        sources = []  # 保存特徵圖
        loc = []  # 保存每一個特徵圖進行位置網絡後的信息
        conf = []  # 保存每一個特徵圖進行置信度網絡後的信息

        # 處理輸入至conv4_3
        for k in range(23):
            x = self.vgg[k](x)

        # 對conv4_3進行L2正則化
        s = self.L2Norm(x)
        sources.append(s)

        # 完成vgg後續的處理
        for k in range(23, len(self.vgg)):
            x = self.vgg[k](x)
        sources.append(x)

        # 使用新增網絡進行處理
        for k, v in enumerate(self.extras):
            x = F.relu(v(x), inplace=True)
            if k % 2 == 1:
                sources.append(x)

        # 將特徵圖送入位置網絡和置信度網絡
        # l(x)或者c(x)的shape爲[batch_size,channel,height,weight],使用了permute後,變成[batch_size,height,weight,channel]
        # 這樣作應該是爲了方便後續處理
        for (x, l, c) in zip(sources, self.loc, self.conf):
            loc.append(l(x).permute(0, 2, 3, 1).contiguous())
            conf.append(c(x).permute(0, 2, 3, 1).contiguous())

        # 進行格式變換
        loc = torch.cat([o.view(o.size(0), -1) for o in loc], 1)  # [batch_size,34928],錨點框的數量8732*4=34928
        conf = torch.cat([o.view(o.size(0), -1) for o in conf], 1)

        if self.phase == 'train':
            output = (loc.view(loc.size(0), -1, 4),  # [batch_size,num_priors,4]
                      conf.view(conf.size(0), -1, self.num_classes),  # [batch_size,num_priors,21]
                      self.priors)  # [num_priors,4]
        else:  # Test
            output = self.detect(
                loc.view(loc.size(0), -1, 4),  # 位置預測
                self.softmax(conf.view((conf.size(0), -1, self.num_classes))),  # 置信度預測
                self.priors.cuda()  # 先驗錨點框
            )

        return output

在上面類SSD中，說起到了先驗錨點框的構建函數PriorBox()，這個函數在models/prior_box.py中。首先，根據用於融合的特徵圖尺寸和product()函數，生成一系列的點，如(0,0),(0,1),(0,2)等。而後根據這些像素點位置，偏移0.5做爲錨點框的中心點，即cx和cy，並將其歸一化。而後計算論文中的$s_k$和${s_k}'$，對應s_k和s_k_prime，先計算$a_r=1$的狀況，再計算其他$a_r$的狀況。此時，mean的shape爲[1,34928]，所以，須要使用view()函數，將其切割出來，變成[8732,4]。記得，這裏的錨點框的座標是[中心點x，中心點y，寬，高]。

def PriorBox(cfg):
    '''
    爲全部特徵圖生成預設的錨點框，返回全部生成的錨點框，尺寸爲[8732,4]，
    每行表示[中心點x，中心點y，寬，高]
    '''
    image_size = cfg['min_dim']  # 300
    feature_maps = cfg['feature_maps']  # [38, 19, 10, 5, 3, 1],特徵圖尺寸
    steps = cfg['steps']  # [8, 16, 32, 64, 100, 300]
    min_sizes = cfg['min_sizes']  # [30, 60, 111, 162, 213, 264]
    max_sizes = cfg['max_sizes']  # [60, 111, 162, 213, 264, 315]
    aspect_ratios = cfg['aspect_ratios']  # [[2], [2, 3], [2, 3], [2, 3], [2], [2]]

    mean = []
    # 爲全部特徵圖生成錨點框
    for k, f in enumerate(feature_maps):
        # product(list1,list2)的做用是依次取出list1中的每1個元素，與list2中的每1個元素，
        # 組成元組，而後，將全部的元組組成一個列表，返回
        # 而這裏使用了repeat,說明1個list重複2次
        for i, j in product(range(f), repeat=2):
            f_k = image_size / steps[k]
            # 計算中心點,這裏的j是沿x方向變化的
            cx = (j + 0.5) / f_k
            cy = (i + 0.5) / f_k

            # aspect_ratio=1有兩種狀況,s_k=s_k,s_k=sqrt(s_k*s_(k+1))
            s_k = min_sizes[k] / image_size
            mean += [cx, cy, s_k, s_k]

            s_k_prime = sqrt(s_k * (max_sizes[k] / image_size))
            mean += [cx, cy, s_k_prime, s_k_prime]

            # 剩餘的aspect_ratio
            for ar in aspect_ratios[k]:
                mean += [cx, cy, s_k * sqrt(ar), s_k / sqrt(ar)]
                mean += [cx, cy, s_k / sqrt(ar), s_k * sqrt(ar)]

    # 此時的mean是1*34928的list，要4個數就分割出來，因此須要用view,從而變成[8732,4],即有8732個錨點框
    output = torch.Tensor(mean).view(-1, 4)
    if cfg['clip']:
        # 對每一個元素進行截斷限制,限制爲[0,1]之間
        output.clamp_(min=0, max=1)
    return output

最後，類SSD中還對conv4的特徵層使用了L2正則化，該函數在models/l2norm.py中。在函數forwand()中，按每一個通道對其值進行L2正則化，即除以通道的平方根來實現歸一化。

class L2Norm(nn.Module):
    '''
    對conv4_3進行l2歸一化
    '''

    def __init__(self, n_channels, scale):
        super(L2Norm, self).__init__()
        self.n_channels = n_channels
        self.gamma = scale
        self.eps = 1e-10
        self.weight = nn.Parameter(torch.Tensor(self.n_channels))  # n_channels個隨機數
        self.reset_parameters()

    def reset_parameters(self):
        # 使用gamma來填充weight的每一個值
        nn.init.constant_(self.weight, self.gamma)

    def forward(self, x):
        # 按通道進行求值
        norm = x.pow(2).sum(dim=1, keepdim=True).sqrt() + self.eps  # [1,1,38,38]
        x = torch.div(x, norm)
        # 將weight經過3個unsqueeze展開成[1,512,1,1],而後經過expand_as進行擴展,造成[1,512,38,38]
        out = self.weight.unsqueeze(0).unsqueeze(2).unsqueeze(3).expand_as(x) * x
        return out

 

---

 

至此，SSD的網絡搭建過程已經完成了，經過類SSD的forward()函數，即能返回預測框的座標和類別置信度，以此能夠計算損失函數。