睿智的目标检测23——Pytorch搭建SSD目标检测平台

  • 学习前言
  • 什么是SSD目标检测算法
  • 源码下载
  • SSD实现思路
  • 一、预测部分
  • 1、主干网络介绍
  • 2、从特征获取预测结果
  • 3、预测结果的解码
  • 4、在原图上进行绘制
  • 二、训练部分
  • 1、真实框的处理
  • a、找到真实框对应的先验框
  • b、真实框的编码
  • 2、利用处理完的真实框与对应图片的预测结果计算loss
  • 训练自己的SSD模型
  • 一、数据集的准备
  • 二、数据集的处理
  • 三、开始网络训练
  • 四、训练结果预测


学习前言

一起来看看SSD的Pytorch实现吧,顺便训练一下自己的数据。

目标检测 封装 pytorch pytorch目标检测教程_目标检测 封装 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、主干网络介绍

目标检测 封装 pytorch pytorch目标检测教程_深度学习_02


SSD采用的主干网络是VGG网络,关于VGG的介绍大家可以看我的另外一篇博客,这里的VGG网络相比普通的VGG网络有一定的修改,主要修改的地方就是:

1、将VGG16的FC6和FC7层转化为卷积层。
2、去掉所有的Dropout层和FC8层;
3、新增了Conv6、Conv7、Conv8、Conv9。

目标检测 封装 pytorch pytorch目标检测教程_ssd_03


如图所示,输入的图片经过了改进的VGG网络(Conv1->fc7)和几个另加的卷积层(Conv6->Conv9),进行特征提取:

a、输入一张图片后,被resize到300x300的shape

b、conv1,经过两次[3,3]卷积网络,输出的特征层为64,输出为(300,300,64),再2X2最大池化,该最大池化步长为2,输出net为(150,150,64)。

c、conv2,经过两次[3,3]卷积网络,输出的特征层为128,输出net为(150,150,128),再2X2最大池化,该最大池化步长为2,输出net为(75,75,128)。

d、conv3,经过三次[3,3]卷积网络,输出的特征层为256,输出net为(75,75,256),再2X2最大池化,该最大池化步长为2,输出net为(38,38,256)。

e、conv4,经过三次[3,3]卷积网络,输出的特征层为512,输出net为(38,38,512),再2X2最大池化,该最大池化步长为2,输出net为(19,19,512)。

f、conv5,经过三次[3,3]卷积网络,输出的特征层为512,输出net为(19,19,512),再3X3最大池化,该最大池化步长为1,输出net为(19,19,512)。

g、利用卷积代替全连接层,进行了一次[3,3]卷积网络和一次[1,1]卷积网络,分别为fc6和fc7,输出的通道数为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)。

实现代码:

import torch.nn as nn
from torchvision.models.utils import load_state_dict_from_url


'''
该代码用于获得VGG主干特征提取网络的输出。
输入变量i代表的是输入图片的通道数,通常为3。

300, 300, 3 -> 300, 300, 64 -> 300, 300, 64 -> 150, 150, 64 -> 150, 150, 128 -> 150, 150, 128 -> 75, 75, 128 ->
75, 75, 256 -> 75, 75, 256 -> 75, 75, 256 -> 38, 38, 256 -> 38, 38, 512 -> 38, 38, 512 -> 38, 38, 512 -> 19, 19, 512 ->
19, 19, 512 -> 19, 19, 512 -> 19, 19, 512 -> 19, 19, 512 -> 19, 19, 1024 -> 19, 19, 1024

38, 38, 512的序号是22
19, 19, 1024的序号是34
'''
base = [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'C', 512, 512, 512, 'M',
            512, 512, 512]

def vgg(pretrained = False):
    layers = []
    in_channels = 3
    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
    # 19, 19, 512 -> 19, 19, 512 
    pool5 = nn.MaxPool2d(kernel_size=3, stride=1, padding=1)
    # 19, 19, 512 -> 19, 19, 1024
    conv6 = nn.Conv2d(512, 1024, kernel_size=3, padding=6, dilation=6)
    # 19, 19, 1024 -> 19, 19, 1024
    conv7 = nn.Conv2d(1024, 1024, kernel_size=1)
    layers += [pool5, conv6,
               nn.ReLU(inplace=True), conv7, nn.ReLU(inplace=True)]

    model = nn.ModuleList(layers)
    if pretrained:
        state_dict = load_state_dict_from_url("https://download.pytorch.org/models/vgg16-397923af.pth", model_dir="./model_data")
        state_dict = {k.replace('features.', '') : v for k, v in state_dict.items()}
        model.load_state_dict(state_dict, strict = False)
    return model

def add_extras(in_channels, backbone_name):
    layers = []
    if backbone_name == 'vgg':
        # Block 6
        # 19,19,1024 -> 19,19,256 -> 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 -> 10,10,128 -> 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 -> 5,5,128 -> 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 -> 3,3,128 -> 1,1,256
        layers += [nn.Conv2d(256, 128, kernel_size=1, stride=1)]
        layers += [nn.Conv2d(128, 256, kernel_size=3, stride=1)]
    else:
        layers += [InvertedResidual(in_channels, 512, stride=2, expand_ratio=0.2)]
        layers += [InvertedResidual(512, 256, stride=2, expand_ratio=0.25)]
        layers += [InvertedResidual(256, 256, stride=2, expand_ratio=0.5)]
        layers += [InvertedResidual(256, 64, stride=2, expand_ratio=0.25)]
        
    return nn.ModuleList(layers)

2、从特征获取预测结果

目标检测 封装 pytorch pytorch目标检测教程_深度学习_02


由上图可知,我们分别取:

  • conv4的第三次卷积的特征;
  • fc7卷积的特征;
  • conv6的第二次卷积的特征;
  • conv7的第二次卷积的特征;
  • conv8的第二次卷积的特征;
  • conv9的第二次卷积的特征。

共六个特征层进行下一步的处理。为了和普通特征层区分,我们称之为有效特征层,来获取预测结果。

对获取到的每一个有效特征层,我们都需要对其做两个操作,分别是:

  • 一次num_anchors x 4的卷积
  • 一次num_anchors x num_classes的卷积

num_anchors指的是该特征层每一个特征点所拥有的先验框数量。上述提到的六个特征层,每个特征层的每个特征点对应的先验框数量分别为4、6、6、6、4、4。

上述操作分别对应的对象为:

  • num_anchors x 4的卷积 用于预测 该特征层上 每一个网格点上 每一个先验框的变化情况。**
  • num_anchors x num_classes的卷积 用于预测 该特征层上 每一个网格点上 每一个预测对应的种类。

所有的特征层对应的预测结果的shape如下:

目标检测 封装 pytorch pytorch目标检测教程_Pytorch_05


实现代码为:

class SSD300(nn.Module):
    def __init__(self, num_classes, backbone_name, pretrained = False):
        super(SSD300, self).__init__()
        self.num_classes    = num_classes
        if backbone_name    == "vgg":
            self.vgg        = add_vgg(pretrained)
            self.extras     = add_extras(1024, backbone_name)
            self.L2Norm     = L2Norm(512, 20)
            mbox            = [4, 6, 6, 6, 4, 4]
            
            loc_layers      = []
            conf_layers     = []
            backbone_source = [21, -2]
            #---------------------------------------------------#
            #   在add_vgg获得的特征层里
            #   第21层和-2层可以用来进行回归预测和分类预测。
            #   分别是conv4-3(38,38,512)和conv7(19,19,1024)的输出
            #---------------------------------------------------#
            for k, v in enumerate(backbone_source):
                loc_layers  += [nn.Conv2d(self.vgg[v].out_channels, mbox[k] * 4, kernel_size = 3, padding = 1)]
                conf_layers += [nn.Conv2d(self.vgg[v].out_channels, mbox[k] * num_classes, kernel_size = 3, padding = 1)]
            #-------------------------------------------------------------#
            #   在add_extras获得的特征层里
            #   第1层、第3层、第5层、第7层可以用来进行回归预测和分类预测。
            #   shape分别为(10,10,512), (5,5,256), (3,3,256), (1,1,256)
            #-------------------------------------------------------------#  
            for k, v in enumerate(self.extras[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)]
        else:
            self.mobilenet  = mobilenet_v2(pretrained).features
            self.extras     = add_extras(1280, backbone_name)
            self.L2Norm     = L2Norm(96, 20)
            mbox            = [6, 6, 6, 6, 6, 6]

            loc_layers      = []
            conf_layers     = []
            backbone_source = [13, -1]
            for k, v in enumerate(backbone_source):
                loc_layers  += [nn.Conv2d(self.mobilenet[v].out_channels, mbox[k] * 4, kernel_size = 3, padding = 1)]
                conf_layers += [nn.Conv2d(self.mobilenet[v].out_channels, mbox[k] * num_classes, kernel_size = 3, padding = 1)]
            for k, v in enumerate(self.extras, 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)]

        self.loc            = nn.ModuleList(loc_layers)
        self.conf           = nn.ModuleList(conf_layers)
        self.backbone_name  = backbone_name
        
    def forward(self, x):
        #---------------------------#
        #   x是300,300,3
        #---------------------------#
        sources = list()
        loc     = list()
        conf    = list()

        #---------------------------#
        #   获得conv4_3的内容
        #   shape为38,38,512
        #---------------------------#
        if self.backbone_name == "vgg":
            for k in range(23):
                x = self.vgg[k](x)
        else:
            for k in range(14):
                x = self.mobilenet[k](x)
        #---------------------------#
        #   conv4_3的内容
        #   需要进行L2标准化
        #---------------------------#
        s = self.L2Norm(x)
        sources.append(s)

        #---------------------------#
        #   获得conv7的内容
        #   shape为19,19,1024
        #---------------------------#
        if self.backbone_name == "vgg":
            for k in range(23, len(self.vgg)):
                x = self.vgg[k](x)
        else:
            for k in range(14, len(self.mobilenet)):
                x = self.mobilenet[k](x)

        sources.append(x)
        #-------------------------------------------------------------#
        #   在add_extras获得的特征层里
        #   第1层、第3层、第5层、第7层可以用来进行回归预测和分类预测。
        #   shape分别为(10,10,512), (5,5,256), (3,3,256), (1,1,256)
        #-------------------------------------------------------------#      
        for k, v in enumerate(self.extras):
            x = F.relu(v(x), inplace=True)
            if self.backbone_name == "vgg":
                if k % 2 == 1:
                    sources.append(x)
            else:
                sources.append(x)

        #-------------------------------------------------------------#
        #   为获得的6个有效特征层添加回归预测和分类预测
        #-------------------------------------------------------------#      
        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())

        #-------------------------------------------------------------#
        #   进行reshape方便堆叠
        #-------------------------------------------------------------#  
        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)
        #-------------------------------------------------------------#
        #   loc会reshape到batch_size, num_anchors, 4
        #   conf会reshap到batch_size, num_anchors, self.num_classes
        #-------------------------------------------------------------#     
        output = (
            loc.view(loc.size(0), -1, 4),
            conf.view(conf.size(0), -1, self.num_classes),
        )
        return output

3、预测结果的解码

利用SSD的主干网络我们可以获得六个有效特征层,分别是:

  • conv4的第三次卷积的特征;
  • fc7卷积的特征;
  • conv6的第二次卷积的特征;
  • conv7的第二次卷积的特征;
  • conv8的第二次卷积的特征;
  • conv9的第二次卷积的特征。

通过对每一个特征层的处理,我们可以获得每个特征层对应的两个内容,分别是:

  • num_anchors x 4的卷积 用于预测 该特征层上 每一个网格点上 每一个先验框的变化情况。**
  • num_anchors x num_classes的卷积 用于预测 该特征层上 每一个网格点上 每一个预测对应的种类。

我们利用 num_anchors x 4的卷积每一个有效特征层对应的先验框 进行调整获得 预测框

在这里我们简单了解一下每个特征层到底做了什么:
每一个有效特征层将整个图片分成与其长宽对应的网格如conv4-3的特征层就是将整个图像分成38x38个网格;然后从每个网格中心建立多个先验框,对于conv4-3的特征层来说,它的每个特征点分别建立了4个先验框;

因此,对于conv4-3整个特征层来讲,整个图片被分成38x38个网格,每个网格中心对应4个先验框,一共建立了38x38x4个,5776个先验框。这些框密密麻麻的遍布在整个图片上网络的预测结果会对这些框进行调整获得预测框。

目标检测 封装 pytorch pytorch目标检测教程_深度学习_06


先验框虽然可以代表一定的框的位置信息与框的大小信息,但是其是有限的,无法表示任意情况,因此还需要调整,ssd利用num_anchors x 4的卷积的结果对先验框进行调整。

num_anchors x 4中的num_anchors表示了这个网格点所包含的先验框数量,其中的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、利用框的位置和得分进行非极大抑制。
实现代码如下:

import numpy as np
import torch
from torch import nn
from torchvision.ops import nms


class BBoxUtility(object):
    def __init__(self, num_classes):
        self.num_classes    = num_classes

    def ssd_correct_boxes(self, box_xy, box_wh, input_shape, image_shape, letterbox_image):
        #-----------------------------------------------------------------#
        #   把y轴放前面是因为方便预测框和图像的宽高进行相乘
        #-----------------------------------------------------------------#
        box_yx = box_xy[..., ::-1]
        box_hw = box_wh[..., ::-1]
        input_shape = np.array(input_shape)
        image_shape = np.array(image_shape)

        if letterbox_image:
            #-----------------------------------------------------------------#
            #   这里求出来的offset是图像有效区域相对于图像左上角的偏移情况
            #   new_shape指的是宽高缩放情况
            #-----------------------------------------------------------------#
            new_shape = np.round(image_shape * np.min(input_shape/image_shape))
            offset  = (input_shape - new_shape)/2./input_shape
            scale   = input_shape/new_shape

            box_yx  = (box_yx - offset) * scale
            box_hw *= scale

        box_mins    = box_yx - (box_hw / 2.)
        box_maxes   = box_yx + (box_hw / 2.)
        boxes  = np.concatenate([box_mins[..., 0:1], box_mins[..., 1:2], box_maxes[..., 0:1], box_maxes[..., 1:2]], axis=-1)
        boxes *= np.concatenate([image_shape, image_shape], axis=-1)
        return boxes

    def decode_boxes(self, mbox_loc, anchors, variances):
        # 获得先验框的宽与高
        anchor_width     = anchors[:, 2] - anchors[:, 0]
        anchor_height    = anchors[:, 3] - anchors[:, 1]
        # 获得先验框的中心点
        anchor_center_x  = 0.5 * (anchors[:, 2] + anchors[:, 0])
        anchor_center_y  = 0.5 * (anchors[:, 3] + anchors[:, 1])

        # 真实框距离先验框中心的xy轴偏移情况
        decode_bbox_center_x = mbox_loc[:, 0] * anchor_width * variances[0]
        decode_bbox_center_x += anchor_center_x
        decode_bbox_center_y = mbox_loc[:, 1] * anchor_height * variances[0]
        decode_bbox_center_y += anchor_center_y
        
        # 真实框的宽与高的求取
        decode_bbox_width   = torch.exp(mbox_loc[:, 2] * variances[1])
        decode_bbox_width   *= anchor_width
        decode_bbox_height  = torch.exp(mbox_loc[:, 3] * variances[1])
        decode_bbox_height  *= anchor_height

        # 获取真实框的左上角与右下角
        decode_bbox_xmin = decode_bbox_center_x - 0.5 * decode_bbox_width
        decode_bbox_ymin = decode_bbox_center_y - 0.5 * decode_bbox_height
        decode_bbox_xmax = decode_bbox_center_x + 0.5 * decode_bbox_width
        decode_bbox_ymax = decode_bbox_center_y + 0.5 * decode_bbox_height

        # 真实框的左上角与右下角进行堆叠
        decode_bbox = torch.cat((decode_bbox_xmin[:, None],
                                      decode_bbox_ymin[:, None],
                                      decode_bbox_xmax[:, None],
                                      decode_bbox_ymax[:, None]), dim=-1)
        # 防止超出0与1
        decode_bbox = torch.min(torch.max(decode_bbox, torch.zeros_like(decode_bbox)), torch.ones_like(decode_bbox))
        return decode_bbox

    def decode_box(self, predictions, anchors, image_shape, input_shape, letterbox_image, variances = [0.1, 0.2], nms_iou = 0.3, confidence = 0.5):
        #---------------------------------------------------#
        #   :4是回归预测结果
        #---------------------------------------------------#
        mbox_loc        = predictions[0]
        #---------------------------------------------------#
        #   获得种类的置信度
        #---------------------------------------------------#
        mbox_conf       = nn.Softmax(-1)(predictions[1])

        results = []
        #----------------------------------------------------------------------------------------------------------------#
        #   对每一张图片进行处理,由于在predict.py的时候,我们只输入一张图片,所以for i in range(len(mbox_loc))只进行一次
        #----------------------------------------------------------------------------------------------------------------#
        for i in range(len(mbox_loc)):
            results.append([])
            #--------------------------------#
            #   利用回归结果对先验框进行解码
            #--------------------------------#
            decode_bbox = self.decode_boxes(mbox_loc[i], anchors, variances)

            for c in range(1, self.num_classes):
                #--------------------------------#
                #   取出属于该类的所有框的置信度
                #   判断是否大于门限
                #--------------------------------#
                c_confs     = mbox_conf[i, :, c]
                c_confs_m   = c_confs > confidence
                if len(c_confs[c_confs_m]) > 0:
                    #-----------------------------------------#
                    #   取出得分高于confidence的框
                    #-----------------------------------------#
                    boxes_to_process = decode_bbox[c_confs_m]
                    confs_to_process = c_confs[c_confs_m]

                    keep = nms(
                        boxes_to_process,
                        confs_to_process,
                        nms_iou
                    )
                    #-----------------------------------------#
                    #   取出在非极大抑制中效果较好的内容
                    #-----------------------------------------#
                    good_boxes  = boxes_to_process[keep]
                    confs       = confs_to_process[keep][:, None]
                    labels      = (c - 1) * torch.ones((len(keep), 1)).cuda() if confs.is_cuda else torch.ones((len(keep), 1))
                    #-----------------------------------------#
                    #   将label、置信度、框的位置进行堆叠。
                    #-----------------------------------------#
                    c_pred      = torch.cat((good_boxes, labels, confs), dim=1).cpu().numpy()
                    # 添加进result里
                    results[-1].extend(c_pred)

            if len(results[-1]) > 0:
                results[-1] = np.array(results[-1])
                box_xy, box_wh = (results[-1][:, 0:2] + results[-1][:, 2:4])/2, results[-1][:, 2:4] - results[-1][:, 0:2]
                results[-1][:, :4] = self.ssd_correct_boxes(box_xy, box_wh, input_shape, image_shape, letterbox_image)

        return results

4、在原图上进行绘制

通过第三步,我们可以获得预测框在原图上的位置,而且这些预测框都是经过筛选的。这些筛选后的框可以直接绘制在图片上,就可以获得结果了。

二、训练部分

1、真实框的处理

真实框的处理可以分为两个部分,分别是:

  • 找到真实框对应的先验框
  • 对真实框进行编码。
a、找到真实框对应的先验框

在这一步中,我们需要找到真实框所对应的先验框,代表这个真实框由某个先验框进行预测。

我们首先需要将每一个的真实框和所有的先验框进行一个iou计算,这一步做的工作是计算每一个真实框和所有的先验框的重合程度。

在获得每一个真实框和所有的先验框的重合程度之后,选出和每一个真实框重合程度大于一定门限的先验框。代表这个真实框由这些先验框负责预测。

由于一个先验框只能负责一个真实框的预测,所以如果某个先验框和多个真实框的重合度较大,那么这个先验框只负责与其iou最大的真实框的预测。

在这一步后,我们可以找到每一个先验框所负责预测的真实框,在下一步中,我们需要根据这些真实框和先验框获得网络的预测结果。

实现代码如下:

def iou(self, box):
    #---------------------------------------------#
    #   计算出每个真实框与所有的先验框的iou
    #   判断真实框与先验框的重合情况
    #---------------------------------------------#
    inter_upleft    = np.maximum(self.anchors[:, :2], box[:2])
    inter_botright  = np.minimum(self.anchors[:, 2:4], box[2:])

    inter_wh    = inter_botright - inter_upleft
    inter_wh    = np.maximum(inter_wh, 0)
    inter       = inter_wh[:, 0] * inter_wh[:, 1]
    #---------------------------------------------# 
    #   真实框的面积
    #---------------------------------------------#
    area_true = (box[2] - box[0]) * (box[3] - box[1])
    #---------------------------------------------#
    #   先验框的面积
    #---------------------------------------------#
    area_gt = (self.anchors[:, 2] - self.anchors[:, 0])*(self.anchors[:, 3] - self.anchors[:, 1])
    #---------------------------------------------#
    #   计算iou
    #---------------------------------------------#
    union = area_true + area_gt - inter

    iou = inter / union
    return iou

def encode_box(self, box, return_iou=True, variances = [0.1, 0.1, 0.2, 0.2]):
    #---------------------------------------------#
    #   计算当前真实框和先验框的重合情况
    #   iou [self.num_anchors]
    #   encoded_box [self.num_anchors, 5]
    #---------------------------------------------#
    iou = self.iou(box)
    encoded_box = np.zeros((self.num_anchors, 4 + return_iou))
    
    #---------------------------------------------#
    #   找到每一个真实框,重合程度较高的先验框
    #   真实框可以由这个先验框来负责预测
    #---------------------------------------------#
    assign_mask = iou > self.overlap_threshold

    #---------------------------------------------#
    #   如果没有一个先验框重合度大于self.overlap_threshold
    #   则选择重合度最大的为正样本
    #---------------------------------------------#
    if not assign_mask.any():
        assign_mask[iou.argmax()] = True
    
    #---------------------------------------------#
    #   利用iou进行赋值 
    #---------------------------------------------#
    if return_iou:
        encoded_box[:, -1][assign_mask] = iou[assign_mask]
    
    #---------------------------------------------#
    #   找到对应的先验框
    #---------------------------------------------#
    assigned_anchors = self.anchors[assign_mask]
b、真实框的编码

利用SSD的主干网络我们可以获得六个有效特征层,分别是:

  • conv4的第三次卷积的特征;
  • fc7卷积的特征;
  • conv6的第二次卷积的特征;
  • conv7的第二次卷积的特征;
  • conv8的第二次卷积的特征;
  • conv9的第二次卷积的特征。

通过对每一个特征层的处理,我们可以获得每个特征层对应的三个内容,分别是:

  • num_anchors x 4的卷积 用于预测 该特征层上 每一个网格点上 每一个先验框的变化情况。**
  • num_anchors x num_classes的卷积 用于预测 该特征层上 每一个网格点上 每一个预测对应的种类。
  • 每一个特征层每一个特征点上对应的若干个先验框

因此,我们直接利用ssd网络预测到的结果,并不是预测框在图片上的真实位置,需要解码才能得到真实位置。

所以在训练的时候,如果我们需要计算loss函数,这个loss函数是相对于ssd网络的预测结果的。因此我们需要对真实框的信息进行处理,使得它的结构和预测结果的格式是一样的,这样的一个过程我们称之为编码(encode)。用一句话概括编码的过程就是将真实框的位置信息格式转化为ssd预测结果的格式信息

也就是,我们需要找到 每一张用于训练的图片每一个真实框对应的先验框,并求出如果想要得到这样一个真实框,我们的预测结果应该是怎么样的。

从预测结果获得真实框的过程被称作解码,而从真实框获得预测结果的过程就是编码的过程。因此我们只需要将解码过程逆过来就是编码过程了。

因此我们可以利用真实框和先验框进行编码,获得该特征点对应的该先验框应该有的预测结果。预测结果分为两个,分别是:

  • num_anchors x 4的卷积
  • num_anchors x num_classes的卷积

前者对应先验框的调整参数,后者对应先验框的种类,这个都可以通过先验框对应的真实框获得。

实现代码如下:

#---------------------------------------------#
#   逆向编码,将真实框转化为ssd预测结果的格式
#   先计算真实框的中心与长宽
#---------------------------------------------#
box_center  = 0.5 * (box[:2] + box[2:])
box_wh      = box[2:] - box[:2]
#---------------------------------------------#
#   再计算重合度较高的先验框的中心与长宽
#---------------------------------------------#
assigned_anchors_center = (assigned_anchors[:, 0:2] + assigned_anchors[:, 2:4]) * 0.5
assigned_anchors_wh     = (assigned_anchors[:, 2:4] - assigned_anchors[:, 0:2])

#------------------------------------------------#
#   逆向求取ssd应该有的预测结果
#   先求取中心的预测结果,再求取宽高的预测结果
#   存在改变数量级的参数,默认为[0.1,0.1,0.2,0.2]
#------------------------------------------------#
encoded_box[:, :2][assign_mask] = box_center - assigned_anchors_center
encoded_box[:, :2][assign_mask] /= assigned_anchors_wh
encoded_box[:, :2][assign_mask] /= np.array(variances)[:2]

encoded_box[:, 2:4][assign_mask] = np.log(box_wh / assigned_anchors_wh)
encoded_box[:, 2:4][assign_mask] /= np.array(variances)[2:4]
return encoded_box.ravel()

2、利用处理完的真实框与对应图片的预测结果计算loss

loss的计算分为三个部分:
1、获取所有正标签的框的预测结果的回归loss。
2、获取所有正标签的种类的预测结果的交叉熵loss。
3、获取一定负标签的种类的预测结果的交叉熵loss。

由于在ssd的训练过程中,正负样本极其不平衡,即 存在对应真实框的先验框可能只有若干个,但是不存在对应真实框的负样本却有几千个,这就会导致负样本的loss值极大,因此我们可以考虑减少负样本的选取,对于ssd的训练来讲,常见的情况是取三倍正样本数量的负样本用于训练。这个三倍呢,也可以修改,调整成自己喜欢的数字。

实现代码如下:

import torch
import torch.nn as nn
import torch.nn.functional as F

class MultiboxLoss(nn.Module):
    def __init__(self, num_classes, alpha=1.0, neg_pos_ratio=3.0,
                 background_label_id=0, negatives_for_hard=100.0):
        self.num_classes = num_classes
        self.alpha = alpha
        self.neg_pos_ratio = neg_pos_ratio
        if background_label_id != 0:
            raise Exception('Only 0 as background label id is supported')
        self.background_label_id = background_label_id
        self.negatives_for_hard = negatives_for_hard

    def _l1_smooth_loss(self, y_true, y_pred):
        abs_loss = torch.abs(y_true - y_pred)
        sq_loss = 0.5 * (y_true - y_pred)**2
        l1_loss = torch.where(abs_loss < 1.0, sq_loss, abs_loss - 0.5)
        return torch.sum(l1_loss, -1)

    def _softmax_loss(self, y_true, y_pred):
        y_pred = torch.clamp(y_pred, min = 1e-7)
        softmax_loss = -torch.sum(y_true * torch.log(y_pred),
                                      axis=-1)
        return softmax_loss

    def forward(self, y_true, y_pred):
        # --------------------------------------------- #
        #   y_true batch_size, 8732, 4 + self.num_classes + 1
        #   y_pred batch_size, 8732, 4 + self.num_classes
        # --------------------------------------------- #
        num_boxes       = y_true.size()[1]
        y_pred          = torch.cat([y_pred[0], nn.Softmax(-1)(y_pred[1])], dim = -1)

        # --------------------------------------------- #
        #   分类的loss
        #   batch_size,8732,21 -> batch_size,8732
        # --------------------------------------------- #
        conf_loss = self._softmax_loss(y_true[:, :, 4:-1], y_pred[:, :, 4:])
        
        # --------------------------------------------- #
        #   框的位置的loss
        #   batch_size,8732,4 -> batch_size,8732
        # --------------------------------------------- #
        loc_loss = self._l1_smooth_loss(y_true[:, :, :4],
                                        y_pred[:, :, :4])

        # --------------------------------------------- #
        #   获取所有的正标签的loss
        # --------------------------------------------- #
        pos_loc_loss = torch.sum(loc_loss * y_true[:, :, -1],
                                     axis=1)
        pos_conf_loss = torch.sum(conf_loss * y_true[:, :, -1],
                                      axis=1)

        # --------------------------------------------- #
        #   每一张图的正样本的个数
        #   num_pos     [batch_size,]
        # --------------------------------------------- #
        num_pos = torch.sum(y_true[:, :, -1], axis=-1)

        # --------------------------------------------- #
        #   每一张图的负样本的个数
        #   num_neg     [batch_size,]
        # --------------------------------------------- #
        num_neg = torch.min(self.neg_pos_ratio * num_pos, num_boxes - num_pos)
        # 找到了哪些值是大于0的
        pos_num_neg_mask = num_neg > 0
        # --------------------------------------------- #
        #   如果所有的图,正样本的数量均为0
        #   那么则默认选取100个先验框作为负样本
        # --------------------------------------------- #
        has_min = torch.sum(pos_num_neg_mask)
        
        # --------------------------------------------- #
        #   从这里往后,与视频中看到的代码有些许不同。
        #   由于以前的负样本选取方式存在一些问题,
        #   我对该部分代码进行重构。
        #   求整个batch应该的负样本数量总和
        # --------------------------------------------- #
        num_neg_batch = torch.sum(num_neg) if has_min > 0 else self.negatives_for_hard

        # --------------------------------------------- #
        #   对预测结果进行判断,如果该先验框没有包含物体
        #   那么它的不属于背景的预测概率过大的话
        #   就是难分类样本
        # --------------------------------------------- #
        confs_start = 4 + self.background_label_id + 1
        confs_end   = confs_start + self.num_classes - 1

        # --------------------------------------------- #
        #   batch_size,8732
        #   把不是背景的概率求和,求和后的概率越大
        #   代表越难分类。
        # --------------------------------------------- #
        max_confs = torch.sum(y_pred[:, :, confs_start:confs_end], dim=2)

        # --------------------------------------------------- #
        #   只有没有包含物体的先验框才得到保留
        #   我们在整个batch里面选取最难分类的num_neg_batch个
        #   先验框作为负样本。
        # --------------------------------------------------- #
        max_confs   = (max_confs * (1 - y_true[:, :, -1])).view([-1])

        _, indices  = torch.topk(max_confs, k = int(num_neg_batch.cpu().numpy().tolist()))

        neg_conf_loss = torch.gather(conf_loss.view([-1]), 0, indices)

        # 进行归一化
        num_pos     = torch.where(num_pos != 0, num_pos, torch.ones_like(num_pos))
        total_loss  = torch.sum(pos_conf_loss) + torch.sum(neg_conf_loss) + torch.sum(self.alpha * pos_loc_loss)
        total_loss  = total_loss / torch.sum(num_pos)
        return total_loss

训练自己的SSD模型

首先前往Github下载对应的仓库,下载完后利用解压软件解压,之后用编程软件打开文件夹。
注意打开的根目录必须正确,否则相对目录不正确的情况下,代码将无法运行。

一定要注意打开后的根目录是文件存放的目录。

目标检测 封装 pytorch pytorch目标检测教程_人工智能_07

一、数据集的准备

本文使用VOC格式进行训练,训练前需要自己制作好数据集,如果没有自己的数据集,可以通过Github连接下载VOC12+07的数据集尝试下。

训练前将标签文件放在VOCdevkit文件夹下的VOC2007文件夹下的Annotation中。

目标检测 封装 pytorch pytorch目标检测教程_目标检测 封装 pytorch_08


训练前将图片文件放在VOCdevkit文件夹下的VOC2007文件夹下的JPEGImages中。

目标检测 封装 pytorch pytorch目标检测教程_深度学习_09


此时数据集的摆放已经结束。

二、数据集的处理

在完成数据集的摆放之后,我们需要对数据集进行下一步的处理,目的是获得训练用的2007_train.txt以及2007_val.txt,需要用到根目录下的voc_annotation.py。

voc_annotation.py里面有一些参数需要设置。
分别是annotation_mode、classes_path、trainval_percent、train_percent、VOCdevkit_path,第一次训练可以仅修改classes_path

'''
annotation_mode用于指定该文件运行时计算的内容
annotation_mode为0代表整个标签处理过程,包括获得VOCdevkit/VOC2007/ImageSets里面的txt以及训练用的2007_train.txt、2007_val.txt
annotation_mode为1代表获得VOCdevkit/VOC2007/ImageSets里面的txt
annotation_mode为2代表获得训练用的2007_train.txt、2007_val.txt
'''
annotation_mode     = 0
'''
必须要修改,用于生成2007_train.txt、2007_val.txt的目标信息
与训练和预测所用的classes_path一致即可
如果生成的2007_train.txt里面没有目标信息
那么就是因为classes没有设定正确
仅在annotation_mode为0和2的时候有效
'''
classes_path        = 'model_data/voc_classes.txt'
'''
trainval_percent用于指定(训练集+验证集)与测试集的比例,默认情况下 (训练集+验证集):测试集 = 9:1
train_percent用于指定(训练集+验证集)中训练集与验证集的比例,默认情况下 训练集:验证集 = 9:1
仅在annotation_mode为0和1的时候有效
'''
trainval_percent    = 0.9
train_percent       = 0.9
'''
指向VOC数据集所在的文件夹
默认指向根目录下的VOC数据集
'''
VOCdevkit_path  = 'VOCdevkit'

classes_path用于指向检测类别所对应的txt,以voc数据集为例,我们用的txt为:

目标检测 封装 pytorch pytorch目标检测教程_深度学习_10

训练自己的数据集时,可以自己建立一个cls_classes.txt,里面写自己所需要区分的类别。

三、开始网络训练

通过voc_annotation.py我们已经生成了2007_train.txt以及2007_val.txt,此时我们可以开始训练了。
训练的参数较多,大家可以在下载库后仔细看注释,其中最重要的部分依然是train.py里的classes_path。

classes_path用于指向检测类别所对应的txt,这个txt和voc_annotation.py里面的txt一样!训练自己的数据集必须要修改!

目标检测 封装 pytorch pytorch目标检测教程_深度学习_11

修改完classes_path后就可以运行train.py开始训练了,在训练多个epoch后,权值会生成在logs文件夹中。
其它参数的作用如下:

#-------------------------------#
#   是否使用Cuda
#   没有GPU可以设置成False
#-------------------------------#
Cuda = True
#--------------------------------------------------------#
#   训练前一定要修改classes_path,使其对应自己的数据集
#--------------------------------------------------------#
classes_path    = 'model_data/voc_classes.txt'
#----------------------------------------------------------------------------------------------------------------------------#
#   权值文件请看README,百度网盘下载。数据的预训练权重对不同数据集是通用的,因为特征是通用的。
#   预训练权重对于99%的情况都必须要用,不用的话权值太过随机,特征提取效果不明显,网络训练的结果也不会好。
#
#   如果想要断点续练就将model_path设置成logs文件夹下已经训练的权值文件。 
#   当model_path = ''的时候不加载整个模型的权值。
#
#   此处使用的是整个模型的权重,因此是在train.py进行加载的,下面的pretrain不影响此处的权值加载。
#   如果想要让模型从主干的预训练权值开始训练,则设置model_path = '',下面的pretrain = True,此时仅加载主干。
#   如果想要让模型从0开始训练,则设置model_path = '',下面的pretrain = Fasle,Freeze_Train = Fasle,此时从0开始训练,且没有冻结主干的过程。
#   一般来讲,从0开始训练效果会很差,因为权值太过随机,特征提取效果不明显。
#----------------------------------------------------------------------------------------------------------------------------#
model_path      = 'model_data/ssd_weights.pth'
#------------------------------------------------------#
#   输入的shape大小
#------------------------------------------------------#
input_shape     = [300, 300]
#--------------------------------------------#
#   vgg或者mobilenetv2
#---------------------------------------------#
backbone        = "vgg"
#----------------------------------------------------------------------------------------------------------------------------#
#   是否使用主干网络的预训练权重,此处使用的是主干的权重,因此是在模型构建的时候进行加载的。
#   如果设置了model_path,则主干的权值无需加载,pretrained的值无意义。
#   如果不设置model_path,pretrained = True,此时仅加载主干开始训练。
#   如果不设置model_path,pretrained = False,Freeze_Train = Fasle,此时从0开始训练,且没有冻结主干的过程。
#----------------------------------------------------------------------------------------------------------------------------#
pretrained      = False
#----------------------------------------------------#
#   可用于设定先验框的大小,默认的anchors_size
#   是根据voc数据集设定的,大多数情况下都是通用的!
#   如果想要检测小物体,可以修改anchors_size
#   一般调小浅层先验框的大小就行了!因为浅层负责小物体检测!
#   比如anchors_size = [21, 45, 99, 153, 207, 261, 315]
#----------------------------------------------------#
anchors_size    = [30, 60, 111, 162, 213, 264, 315]

#----------------------------------------------------#
#   训练分为两个阶段,分别是冻结阶段和解冻阶段。
#   显存不足与数据集大小无关,提示显存不足请调小batch_size。
#   受到BatchNorm层影响,batch_size最小为2,不能为1。
#----------------------------------------------------#
#----------------------------------------------------#
#   冻结阶段训练参数
#   此时模型的主干被冻结了,特征提取网络不发生改变
#   占用的显存较小,仅对网络进行微调
#----------------------------------------------------#
Init_Epoch          = 0
Freeze_Epoch        = 50
Freeze_batch_size   = 16
Freeze_lr           = 5e-4
#----------------------------------------------------#
#   解冻阶段训练参数
#   此时模型的主干不被冻结了,特征提取网络会发生改变
#   占用的显存较大,网络所有的参数都会发生改变
#----------------------------------------------------#
UnFreeze_Epoch      = 100
Unfreeze_batch_size = 8
Unfreeze_lr         = 1e-4
#------------------------------------------------------#
#   是否进行冻结训练,默认先冻结主干训练后解冻训练。
#------------------------------------------------------#
Freeze_Train        = True
#------------------------------------------------------#
#   用于设置是否使用多线程读取数据
#   开启后会加快数据读取速度,但是会占用更多内存
#   内存较小的电脑可以设置为2或者0  
#------------------------------------------------------#
num_workers         = 4
#----------------------------------------------------#
#   获得图片路径和标签
#----------------------------------------------------#
train_annotation_path   = '2007_train.txt'
val_annotation_path     = '2007_val.txt'

四、训练结果预测

训练结果预测需要用到两个文件,分别是yolo.py和predict.py。
我们首先需要去yolo.py里面修改model_path以及classes_path,这两个参数必须要修改。

model_path指向训练好的权值文件,在logs文件夹里。
classes_path指向检测类别所对应的txt。

目标检测 封装 pytorch pytorch目标检测教程_Pytorch_12


完成修改后就可以运行predict.py进行检测了。运行后输入图片路径即可检测。