睿智的目標檢測23——Pytorch搭建SSD目標檢測平臺

學習前言

一起來看看SSD的Pytorch實現吧，順便訓練一下自己的數據。

什麼是SSD目標檢測算法

SSD是一種非常優秀的one-stage目標檢測方法，one-stage算法就是目標檢測和分類是同時完成的，其主要思路是利用CNN提取特徵後，均勻地在圖片的不同位置進行密集抽樣，抽樣時可以採用不同尺度和長寬比，物體分類與預測框的迴歸同時進行，整個過程只需要一步，所以其優勢是速度快。
但是均勻的密集採樣的一個重要缺點是訓練比較困難，這主要是因爲正樣本與負樣本（背景）極其不均衡（參見Focal Loss），導致模型準確度稍低。
SSD的英文全名是Single Shot MultiBox Detector，Single shot說明SSD算法屬於one-stage方法，MultiBox說明SSD算法基於多框預測。

源碼下載

https://github.com/bubbliiiing/ssd-pytorch
喜歡的可以點個star噢。

SSD實現思路

一、預測部分

1、主幹網絡介紹

SSD採用的主幹網絡是VGG網絡，關於VGG的介紹大家可以看我的另外一篇博客https://blog.csdn.net/weixin_44791964/article/details/102779878，這裏的VGG網絡相比普通的VGG網絡有一定的修改，主要修改的地方就是：
1、將VGG16的FC6和FC7層轉化爲卷積層。
2、去掉所有的Dropout層和FC8層；
3、新增了Conv6、Conv7、Conv8、Conv9。

如圖所示，輸入的圖片經過了改進的VGG網絡（Conv1->fc7）和幾個另加的卷積層（Conv6->Conv9），進行特徵提取：
a、輸入一張圖片後，被resize到300x300的shape

b、conv1，經過兩次[3,3]卷積網絡，輸出的特徵層爲64，輸出爲(300,300,64)，再2X2最大池化，輸出net爲(150,150,64)。

c、conv2，經過兩次[3,3]卷積網絡，輸出的特徵層爲128，輸出net爲(150,150,128)，再2X2最大池化，輸出net爲(75,75,128)。

d、conv3，經過三次[3,3]卷積網絡，輸出的特徵層爲256，輸出net爲(75,75,256)，再2X2最大池化，輸出net爲(38,38,256)。

e、conv4，經過三次[3,3]卷積網絡，輸出的特徵層爲512，輸出net爲(38,38,512)，再2X2最大池化，輸出net爲(19,19,512)。

f、conv5，經過三次[3,3]卷積網絡，輸出的特徵層爲512，輸出net爲(19,19,512)，再2X2最大池化，輸出net爲(19,19,512)。

g、利用卷積代替全連接層，進行了兩次[3,3]卷積網絡，輸出的特徵層爲1024，因此輸出的net爲(19,19,1024)。（從這裏往前都是VGG的結構）

h、conv6，經過一次[1,1]卷積網絡，調整通道數，一次步長爲2的[3,3]卷積網絡，輸出的特徵層爲512，因此輸出的net爲(10,10,512)。

i、conv7，經過一次[1,1]卷積網絡，調整通道數，一次步長爲2的[3,3]卷積網絡，輸出的特徵層爲256，因此輸出的net爲(5,5,256)。

j、conv8，經過一次[1,1]卷積網絡，調整通道數，一次padding爲valid的[3,3]卷積網絡，輸出的特徵層爲256，因此輸出的net爲(3,3,256)。

k、conv9，經過一次[1,1]卷積網絡，調整通道數，一次padding爲valid的[3,3]卷積網絡，輸出的特徵層爲256，因此輸出的net爲(1,1,256)。

實現代碼：

base = [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'C', 512, 512, 512, 'M',
            512, 512, 512]

def vgg(i):
    layers = []
    in_channels = i
    for v in base:
        if v == 'M':
            layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
        elif v == 'C':
            layers += [nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True)]
        else:
            conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
            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.Conv2d(1024, 1024, kernel_size=1)
    layers += [pool5, conv6,
               nn.ReLU(inplace=True), conv7, nn.ReLU(inplace=True)]
    return layers

def add_extras(i, batch_norm=False):
    # Extra layers added to VGG for feature scaling
    layers = []
    in_channels = i

    # Block 6
    # 19,19,1024 -> 10,10,512
    layers += [nn.Conv2d(in_channels, 256, kernel_size=1, stride=1)]
    layers += [nn.Conv2d(256, 512, kernel_size=3, stride=2, padding=1)]

    # Block 7
    # 10,10,512 -> 5,5,256
    layers += [nn.Conv2d(512, 128, kernel_size=1, stride=1)]
    layers += [nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1)]

    # Block 8
    # 5,5,256 -> 3,3,256
    layers += [nn.Conv2d(256, 128, kernel_size=1, stride=1)]
    layers += [nn.Conv2d(128, 256, kernel_size=3, stride=1)]
    
    # Block 9
    # 3,3,256 -> 1,1,256
    layers += [nn.Conv2d(256, 128, kernel_size=1, stride=1)]
    layers += [nn.Conv2d(128, 256, kernel_size=3, stride=1)]

    return layers

2、從特徵獲取預測結果

由上圖我們可以知道，我們分別取conv4的第三次卷積的特徵、fc7的特徵、conv6的第二次卷積的特徵、conv7的第二次卷積的特徵、conv8的第二次卷積的特徵、conv9的第二次卷積的特徵，爲了和普通特徵層區分，我們稱之爲有效特徵層，來獲取預測結果。

對獲取到的每一個有效特徵層，我們分別對其進行一次num_priors x 4的卷積、一次num_priors x num_classes的卷積、並需要計算每一個有效特徵層對應的先驗框。而num_priors指的是該特徵層所擁有的先驗框數量。

其中：
num_priors x 4的卷積 用於預測 該特徵層上 每一個網格點上 每一個先驗框的變化情況。（爲什麼說是變化情況呢，這是因爲ssd的預測結果需要結合先驗框獲得預測框，預測結果就是先驗框的變化情況。）

num_priors x num_classes的卷積 用於預測 該特徵層上 每一個網格點上 每一個預測框對應的種類。

每一個有效特徵層對應的先驗框對應着該特徵層上 每一個網格點上 預先設定好的多個框。

所有的特徵層對應的預測結果的shape如下：

實現代碼爲：

class SSD(nn.Module):
    def __init__(self, phase, base, extras, head, num_classes):
        super(SSD, self).__init__()
        self.phase = phase
        self.num_classes = num_classes
        self.cfg = Config
        self.vgg = nn.ModuleList(base)
        self.L2Norm = L2Norm(512, 20)
        self.extras = nn.ModuleList(extras)
        self.priorbox = PriorBox(self.cfg)
        with torch.no_grad():
            self.priors = Variable(self.priorbox.forward())
        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, 0, 200, 0.01, 0.45)
    def forward(self, x):
        sources = list()
        loc = list()
        conf = list()

        # 獲得conv4_3的內容
        for k in range(23):
            x = self.vgg[k](x)

        s = self.L2Norm(x)
        sources.append(s)

        # 獲得fc7的內容
        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)



        # 添加回歸層和分類層
        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())

        # 進行resize
        loc = torch.cat([o.view(o.size(0), -1) for o in loc], 1)
        conf = torch.cat([o.view(o.size(0), -1) for o in conf], 1)
        if self.phase == "test":
            # loc會resize到batch_size,num_anchors,4
            # conf會resize到batch_size,num_anchors,
            output = self.detect(
                loc.view(loc.size(0), -1, 4),                   # loc preds
                self.softmax(conf.view(conf.size(0), -1,
                             self.num_classes)),                # conf preds
                self.priors              
            )
        else:
            output = (
                loc.view(loc.size(0), -1, 4),
                conf.view(conf.size(0), -1, self.num_classes),
                self.priors
            )
        return output

mbox = [4, 6, 6, 6, 4, 4]

def get_ssd(phase,num_classes):

    vgg, extra_layers = add_vgg(3), add_extras(1024)

    loc_layers = []
    conf_layers = []
    vgg_source = [21, -2]
    for k, v in enumerate(vgg_source):
        loc_layers += [nn.Conv2d(vgg[v].out_channels,
                                 mbox[k] * 4, kernel_size=3, padding=1)]
        conf_layers += [nn.Conv2d(vgg[v].out_channels,
                        mbox[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, mbox[k]
                                 * 4, kernel_size=3, padding=1)]
        conf_layers += [nn.Conv2d(v.out_channels, mbox[k]
                                  * num_classes, kernel_size=3, padding=1)]

    SSD_MODEL = SSD(phase, vgg, extra_layers, (loc_layers, conf_layers), num_classes)
    return SSD_MODEL

3、預測結果的解碼

我們通過對每一個特徵層的處理，可以獲得三個內容，分別是：

num_priors x 4的卷積 用於預測 該特徵層上 每一個網格點上 每一個先驗框的變化情況。**

num_priors x num_classes的卷積 用於預測 該特徵層上 每一個網格點上 每一個預測框對應的種類。

每一個有效特徵層對應的先驗框對應着該特徵層上 每一個網格點上 預先設定好的多個框。

我們利用 num_priors x 4的卷積 與 每一個有效特徵層對應的先驗框 獲得框的真實位置。

每一個有效特徵層對應的先驗框就是，如圖所示的作用：
每一個有效特徵層將整個圖片分成與其長寬對應的網格，如conv4-3的特徵層就是將整個圖像分成38x38個網格；然後從每個網格中心建立多個先驗框，如conv4-3的特徵層就是建立了4個先驗框；對於conv4-3的特徵層來講，整個圖片被分成38x38個網格，每個網格中心對應4個先驗框，一共包含了，38x38x4個，5776個先驗框。

先驗框雖然可以代表一定的框的位置信息與框的大小信息，但是其是有限的，無法表示任意情況，因此還需要調整，ssd利用num_priors x 4的卷積的結果對先驗框進行調整。

num_priors x 4中的num_priors表示了這個網格點所包含的先驗框數量，其中的4表示了x_offset、y_offset、h和w的調整情況。

x_offset與y_offset代表了真實框距離先驗框中心的xy軸偏移情況。
h和w代表了真實框的寬與高相對於先驗框的變化情況。

SSD解碼過程就是將每個網格的中心點加上它對應的x_offset和y_offset，加完後的結果就是預測框的中心，然後再利用 先驗框和h、w結合 計算出預測框的長和寬。這樣就能得到整個預測框的位置了。

當然得到最終的預測結構後還要進行得分排序與非極大抑制篩選這一部分基本上是所有目標檢測通用的部分。
1、取出每一類得分大於self.obj_threshold的框和得分。
2、利用框的位置和得分進行非極大抑制。
實現代碼如下：

# Adapted from https://github.com/Hakuyume/chainer-ssd
def decode(loc, priors, variances):
    boxes = torch.cat((
        priors[:, :2] + loc[:, :2] * variances[0] * priors[:, 2:],
        priors[:, 2:] * torch.exp(loc[:, 2:] * variances[1])), 1)
    boxes[:, :2] -= boxes[:, 2:] / 2
    boxes[:, 2:] += boxes[:, :2]
    return boxes

class Detect(Function):
    def __init__(self, num_classes, bkg_label, top_k, conf_thresh, nms_thresh):
        self.num_classes = num_classes
        self.background_label = bkg_label
        self.top_k = top_k
        self.nms_thresh = nms_thresh
        if nms_thresh <= 0:
            raise ValueError('nms_threshold must be non negative.')
        self.conf_thresh = conf_thresh
        self.variance = Config['variance']

    def forward(self, loc_data, conf_data, prior_data):
        loc_data = loc_data.cpu()
        conf_data = conf_data.cpu()
        num = loc_data.size(0)  # batch size
        num_priors = prior_data.size(0)
        output = torch.zeros(num, self.num_classes, self.top_k, 5)
        conf_preds = conf_data.view(num, num_priors,
                                    self.num_classes).transpose(2, 1)
        # 對每一張圖片進行處理
        for i in range(num):
            # 對先驗框解碼獲得預測框
            decoded_boxes = decode(loc_data[i], prior_data, self.variance)
            conf_scores = conf_preds[i].clone()

            for cl in range(1, self.num_classes):
                # 對每一類進行非極大抑制
                c_mask = conf_scores[cl].gt(self.conf_thresh)
                scores = conf_scores[cl][c_mask]
                if scores.size(0) == 0:
                    continue
                l_mask = c_mask.unsqueeze(1).expand_as(decoded_boxes)
                boxes = decoded_boxes[l_mask].view(-1, 4)
                # 進行非極大抑制
                ids, count = nms(boxes, scores, self.nms_thresh, self.top_k)
                output[i, cl, :count] = \
                    torch.cat((scores[ids[:count]].unsqueeze(1),
                               boxes[ids[:count]]), 1)
        flt = output.contiguous().view(num, -1, 5)
        _, idx = flt[:, :, 0].sort(1, descending=True)
        _, rank = idx.sort(1)
        flt[(rank < self.top_k).unsqueeze(-1).expand_as(flt)].fill_(0)
        return output

4、在原圖上進行繪製

通過第三步，我們可以獲得預測框在原圖上的位置，而且這些預測框都是經過篩選的。這些篩選後的框可以直接繪製在圖片上，就可以獲得結果了。

二、訓練部分

1、真實框的處理

從預測部分我們知道，每個特徵層的預測結果，num_priors x 4的卷積 用於預測 該特徵層上 每一個網格點上 每一個先驗框的變化情況。

也就是說，我們直接利用ssd網絡預測到的結果，並不是預測框在圖片上的真實位置，需要解碼才能得到真實位置。

而在訓練的時候，我們需要計算loss函數，這個loss函數是相對於ssd網絡的預測結果的。我們需要把圖片輸入到當前的ssd網絡中，得到預測結果；同時還需要把真實框的信息，進行編碼，這個編碼是把真實框的位置信息格式轉化爲ssd預測結果的格式信息。

也就是，我們需要找到 每一張用於訓練的圖片的每一個真實框對應的先驗框，並求出如果想要得到這樣一個真實框，我們的預測結果應該是怎麼樣的。

從預測結果獲得真實框的過程被稱作解碼，而從真實框獲得預測結果的過程就是編碼的過程。

因此我們只需要將解碼過程逆過來就是編碼過程了。

實現代碼如下：

def encode(matched, priors, variances):
    g_cxcy = (matched[:, :2] + matched[:, 2:])/2 - priors[:, :2]
    g_cxcy /= (variances[0] * priors[:, 2:])
    g_wh = (matched[:, 2:] - matched[:, :2]) / priors[:, 2:]
    g_wh = torch.log(g_wh) / variances[1]
    return torch.cat([g_cxcy, g_wh], 1) 

在訓練的時候我們只需要選擇iou最大的先驗框就行了，這個iou最大的先驗框就是我們用來預測這個真實框所用的先驗框。

因此我們還要經過一次篩選，將上述代碼獲得的真實框對應的所有的iou較大先驗框的預測結果中，iou最大的那個篩選出來。

實現代碼如下：

def match(threshold, truths, priors, variances, labels, loc_t, conf_t, idx):
    # 計算所有的先驗框和真實框的重合程度
    overlaps = jaccard(
        truths,
        point_form(priors)
    )
    # 所有真實框和先驗框的最好重合程度
    # [truth_box,1]
    best_prior_overlap, best_prior_idx = overlaps.max(1, keepdim=True)
    best_prior_idx.squeeze_(1)
    best_prior_overlap.squeeze_(1)
    # 所有先驗框和真實框的最好重合程度
    # [1,prior]
    best_truth_overlap, best_truth_idx = overlaps.max(0, keepdim=True)
    best_truth_idx.squeeze_(0)
    best_truth_overlap.squeeze_(0)
    # 找到與真實框重合程度最好的先驗框，用於保證每個真實框都要有對應的一個先驗框
    best_truth_overlap.index_fill_(0, best_prior_idx, 2)
    # 對best_truth_idx內容進行設置
    for j in range(best_prior_idx.size(0)):
        best_truth_idx[best_prior_idx[j]] = j
    
    # 找到每個先驗框重合程度最好的真實框
    matches = truths[best_truth_idx]          # Shape: [num_priors,4]
    conf = labels[best_truth_idx] + 1         # Shape: [num_priors]
    # 如果重合程度小於threhold則認爲是背景
    conf[best_truth_overlap < threshold] = 0  # label as background
    loc = encode(matches, priors, variances)
    loc_t[idx] = loc    # [num_priors,4] encoded offsets to learn
    conf_t[idx] = conf  # [num_priors] top class label for each prior

2、利用處理完的真實框與對應圖片的預測結果計算loss

loss的計算分爲三個部分：
1、獲取所有正標籤的框的預測結果的迴歸loss。
2、獲取所有正標籤的種類的預測結果的交叉熵loss。
3、獲取一定負標籤的種類的預測結果的交叉熵loss。

由於在ssd的訓練過程中，正負樣本極其不平衡，即 存在對應真實框的先驗框可能只有十來個，但是不存在對應真實框的負樣本卻有幾千個，這就會導致負樣本的loss值極大，因此我們可以考慮減少負樣本的選取，對於ssd的訓練來講，常見的情況是取三倍正樣本數量的負樣本用於訓練。這個三倍呢，也可以修改，調整成自己喜歡的數字。

實現代碼如下：

class MultiBoxLoss(nn.Module):
    def __init__(self, num_classes, overlap_thresh, prior_for_matching,
                 bkg_label, neg_mining, neg_pos, neg_overlap, encode_target,
                 use_gpu=True):
        super(MultiBoxLoss, self).__init__()
        self.use_gpu = use_gpu
        self.num_classes = num_classes
        self.threshold = overlap_thresh
        self.background_label = bkg_label
        self.encode_target = encode_target
        self.use_prior_for_matching = prior_for_matching
        self.do_neg_mining = neg_mining
        self.negpos_ratio = neg_pos
        self.neg_overlap = neg_overlap
        self.variance = Config['variance']

    def forward(self, predictions, targets):
        # 迴歸信息，置信度，先驗框
        loc_data, conf_data, priors = predictions
        # 計算出batch_size
        num = loc_data.size(0)
        # 取出所有的先驗框
        priors = priors[:loc_data.size(1), :]
        # 先驗框的數量
        num_priors = (priors.size(0))
        num_classes = self.num_classes
        # 創建一個tensor進行處理
        loc_t = torch.Tensor(num, num_priors, 4)
        conf_t = torch.LongTensor(num, num_priors)
        for idx in range(num):
            # 獲得框
            truths = targets[idx][:, :-1].data
            # 獲得標籤
            labels = targets[idx][:, -1].data
            # 獲得先驗框
            defaults = priors.data
            # 找到標籤對應的先驗框
            match(self.threshold, truths, defaults, self.variance, labels,
                  loc_t, conf_t, idx)
        if self.use_gpu:
            loc_t = loc_t.cuda()
            conf_t = conf_t.cuda()
            
        # 轉化成Variable
        loc_t = Variable(loc_t, requires_grad=False)
        conf_t = Variable(conf_t, requires_grad=False)

        # 所有conf_t>0的地方，代表內部包含物體
        pos = conf_t > 0
        # 求和得到每一個圖片內部有多少正樣本
        num_pos = pos.sum(dim=1, keepdim=True)
        # 計算迴歸loss
        pos_idx = pos.unsqueeze(pos.dim()).expand_as(loc_data)
        loc_p = loc_data[pos_idx].view(-1, 4)
        loc_t = loc_t[pos_idx].view(-1, 4)
        loss_l = F.smooth_l1_loss(loc_p, loc_t, size_average=False)

        # 轉化形式
        batch_conf = conf_data.view(-1, self.num_classes)
        # 你可以把softmax函數看成一種接受任何數字並轉換爲概率分佈的非線性方法
        # 獲得每個框預測到真實框的類的概率
        loss_c = log_sum_exp(batch_conf) - batch_conf.gather(1, conf_t.view(-1, 1))
        loss_c = loss_c.view(num, -1)

        loss_c[pos] = 0 
        # 獲得每一張圖新的softmax的結果
        _, loss_idx = loss_c.sort(1, descending=True)
        _, idx_rank = loss_idx.sort(1)
        # 計算每一張圖的正樣本數量
        num_pos = pos.long().sum(1, keepdim=True)
        # 限制負樣本數量
        num_neg = torch.clamp(self.negpos_ratio*num_pos, max=pos.size(1)-1)
        neg = idx_rank < num_neg.expand_as(idx_rank)

        # 計算正樣本的loss和負樣本的loss
        pos_idx = pos.unsqueeze(2).expand_as(conf_data)
        neg_idx = neg.unsqueeze(2).expand_as(conf_data)
        conf_p = conf_data[(pos_idx+neg_idx).gt(0)].view(-1, self.num_classes)
        targets_weighted = conf_t[(pos+neg).gt(0)]
        loss_c = F.cross_entropy(conf_p, targets_weighted, size_average=False)

        # Sum of losses: L(x,c,l,g) = (Lconf(x, c) + αLloc(x,l,g)) / N

        N = num_pos.data.sum()
        loss_l /= N
        loss_c /= N
        return loss_l, loss_c

訓練自己的ssd模型

ssd整體的文件夾構架如下：

本文使用VOC格式進行訓練。
訓練前將標籤文件放在VOCdevkit文件夾下的VOC2007文件夾下的Annotation中。

訓練前將圖片文件放在VOCdevkit文件夾下的VOC2007文件夾下的JPEGImages中。

在訓練前利用voc2ssd.py文件生成對應的txt。

再運行根目錄下的voc_annotation.py，運行前需要將classes改成你自己的classes。

classes = ["aeroplane", "bicycle", "bird", "boat", "bottle", "bus", "car", "cat", "chair", "cow", "diningtable", "dog", "horse", "motorbike", "person", "pottedplant", "sheep", "sofa", "train", "tvmonitor"]

就會生成對應的2007_train.txt，每一行對應其圖片位置及其真實框的位置。

在訓練前需要修改model_data裏面的voc_classes.txt文件，需要將classes改成你自己的classes。

還有config文件下面的Num_Classes，修改成分類的數量+1。

運行train.py即可開始訓練。