之前，對SSD的論文進行了解讀，可以回顧之前的博客：https://www.cnblogs.com/dengshunge/p/11665929.html。
為了加深對SSD的理解，因此對SSD的源碼進行了復現，主要參考的github項目是ssd.pytorch。同時，我自己對該項目增加了大量注釋：https://github.com/Dengshunge/mySSD_pytorch
搭建SSD的項目，可以分成以下三個部分：
接下來，本篇博客重點分析網絡搭建。
該部分整體比較簡單，思路也很清晰。
首先，在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()函數，即能返回預測框的坐標和類別置信度，以此可以計算損失函數。