ssd算法的pytorch实现与解读

首先先放下github地址：https://github.com/acm5656/ssd_pytorch

然后放上参考的代码的github地址：https://github.com/amdegroot/ssd.pytorch

为什么要使用pytorch复现呢，因为好多大佬的代码对于萌新真的不友好，看半天看不懂，所以笔者本着学习和练手的目的，尝试复现下，并分享出来帮助其他萌新学习，大佬有兴趣看了后可以提些建议~

然后对ssd原理感兴趣的同学可以参考我的这篇博客https://www.cnblogs.com/cmai/p/10076050.html，主要对SSD模型进行了讲解。在这就主要讲解代码实现上的内容了，就不再讲原理了。

首先看下项目目录：

VOCdevkit：存放训练数据

weights     ：存放权重文件

Config.py  ：默认的一些配置

Test.py      ：测试单张照片的识别

Train.py    ：训练的py文件

augmentation.py：data augmentation的py文件，主要功能是扩大训练数据

detection.py：对识别的结果的数据进行部分筛选，传送给Test.py文件，供其调用使用

l2norm.py：进行l2正则化

loss_function.py：计算损失函数

ssd_net_vgg.py：ssd模型的实现

utils.py：工具类

voc0712.py：重写dataset类，提取voc的数据并规则化

模型搭建

模型搭建在ssd_net_vgg.py中，这个类只需要将一点，即vgg的网络需要注意，必须采用笔者的方式搭建，否则pre-train的model加载出错，具体的原因不在这里阐述。

模型的实现过程，将loc和conf的提取分开进行了，这个不影响正常的使用，只是在计算损失函数时，能够方便编程而已。

default box计算

代码在utils.py文件下，代码如下：

def default_prior_box():
    mean_layer = []
    for k,f in enumerate(Config.feature_map):
        mean = []
        for i,j in product(range(f),repeat=2):
            f_k = Config.image_size/Config.steps[k]
            cx = (j+0.5)/f_k
            cy = (i+0.5)/f_k

            s_k = Config.sk[k]/Config.image_size
            mean += [cx,cy,s_k,s_k]

            s_k_prime = sqrt(s_k * Config.sk[k+1]/Config.image_size)
            mean += [cx,cy,s_k_prime,s_k_prime]
            for ar in Config.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)]
        if Config.use_cuda:
            mean = torch.Tensor(mean).cuda().view(Config.feature_map[k], Config.feature_map[k], -1).contiguous()
        else:
            mean = torch.Tensor(mean).view( Config.feature_map[k],Config.feature_map[k],-1).contiguous()
        mean.clamp_(max=1, min=0)
        mean_layer.append(mean)

    return mean_layer

该函数则是生成box，与论文中的数量对应，最后的输出是6个list，每个list对应一个特征层输出的default box数，具体数量参考上一篇ssd论文解读的博客。计算公式同参考上篇博客。

Loss函数计算

loss函数的功能实现在loss_function.py中，具体核心代码如下：

class LossFun(nn.Module):
    def __init__(self):
        super(LossFun,self).__init__()
    def forward(self, prediction,targets,priors_boxes):
        loc_data , conf_data = prediction
        loc_data = torch.cat([o.view(o.size(0),-1,4) for o in loc_data] ,1)
        conf_data = torch.cat([o.view(o.size(0),-1,21) for o in conf_data],1)
        priors_boxes = torch.cat([o.view(-1,4) for o in priors_boxes],0)
        if Config.use_cuda:
            loc_data = loc_data.cuda()
            conf_data = conf_data.cuda()
            priors_boxes = priors_boxes.cuda()
        # batch_size
        batch_num = loc_data.size(0)
        # default_box数量
        box_num = loc_data.size(1)
        # 存储targets根据每一个prior_box变换后的数据
        target_loc = torch.Tensor(batch_num,box_num,4)
        target_loc.requires_grad_(requires_grad=False)
        # 存储每一个default_box预测的种类
        target_conf = torch.LongTensor(batch_num,box_num)
        target_conf.requires_grad_(requires_grad=False)
        if Config.use_cuda:
            target_loc = target_loc.cuda()
            target_conf = target_conf.cuda()
        # 因为一次batch可能有多个图，每次循环计算出一个图中的box，即8732个box的loc和conf，存放在target_loc和target_conf中
        for batch_id in range(batch_num):
            target_truths = targets[batch_id][:,:-1].data
            target_labels = targets[batch_id][:,-1].data
            if Config.use_cuda:
                target_truths = target_truths.cuda()
                target_labels = target_labels.cuda()
            # 计算box函数，即公式中loc损失函数的计算公式
            utils.match(0.5,target_truths,priors_boxes,target_labels,target_loc,target_conf,batch_id)
        pos = target_conf > 0
        pos_idx = pos.unsqueeze(pos.dim()).expand_as(loc_data)
        # 相当于论文中L1损失函数乘xij的操作
        pre_loc_xij = loc_data[pos_idx].view(-1,4)
        tar_loc_xij = target_loc[pos_idx].view(-1,4)
        # 将计算好的loc和预测进行smooth_li损失函数
        loss_loc = F.smooth_l1_loss(pre_loc_xij,tar_loc_xij,size_average=False)

        batch_conf = conf_data.view(-1,21)

        # 参照论文中conf计算方式，求出ci
        loss_c = utils.log_sum_exp(batch_conf) - batch_conf.gather(1, target_conf.view(-1, 1))

        loss_c = loss_c.view(batch_num, -1)
        # 将正样本设定为0
        loss_c[pos] = 0

        # 将剩下的负样本排序，选出目标数量的负样本
        _, 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(3*num_pos, max=pos.size(1)-1)

        # 提取出正负样本
        neg = idx_rank < num_neg.expand_as(idx_rank)
        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, 21)
        targets_weighted = target_conf[(pos+neg).gt(0)]
        loss_c = F.cross_entropy(conf_p, targets_weighted, size_average=False)

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

其中较为复杂的是match函数，其具体的代码如下：

def match(threshold, truths, priors, variances, labels, loc_t, conf_t, idx):
    """计算default box和实际位置的jaccard比，计算出每个box的最大jaccard比的种类和每个种类的最大jaccard比的box
    Args:
        threshold: (float) jaccard比的阈值.
        truths: (tensor) 实际位置.
        priors: (tensor) default box
        variances: (tensor) 这个数据含义暂时不清楚，笔者测试过，如果不使用同样可以训练.
        labels: (tensor) 一个图片实际包含的类别数.
        loc_t: (tensor) 需要存储每个box不同类别中的最大jaccard比.
        conf_t: (tensor) 存储每个box的最大jaccard比的类别.
        idx: (int) 当前的批次
    """
    # 计算jaccard比
    overlaps = jaccard(
        truths,
        # 转换priors，转换为x_min,y_min,x_max和y_max
        point_form(priors)
    )
    # [1,num_objects] best prior for each ground truth
    # 实际包含的类别对应box中jaccarb最大的box和对应的索引值，即每个类别最优box
    best_prior_overlap, best_prior_idx = overlaps.max(1, keepdim=True)
    # [1,num_priors] best ground truth for each prior
    # 每一个box,在实际类别中最大的jaccard比的类别，即每个box最优类别
    best_truth_overlap, best_truth_idx = overlaps.max(0, keepdim=True)
    best_truth_idx.squeeze_(0)
    best_truth_overlap.squeeze_(0)
    best_prior_idx.squeeze_(1)
    best_prior_overlap.squeeze_(1)
    # 将每个类别中的最大box设置为2，确保不影响后边操作
    best_truth_overlap.index_fill_(0, best_prior_idx, 2)

    # 计算每一个box的最优类别，和每个类别的最优loc
    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]
    conf[best_truth_overlap < threshold] = 0  # label as background
    # 实现loc的转换，具体的转换公式参照论文中的loc的loss函数的计算公式
    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

代码已经添加了比较详细的注释了，因此不再做过多的解释了。

个人认为比较难的部分代码就是上述的几块，希望读者有时间可以debug调试测试一下，再配合注释，应该能够理解具体的内容，代码中data augumentation 部分没有做详细的解释，这部分笔者也没搞得太明白，只是知道其功能是对数据集进行了扩大，即扩大图像尺寸或者裁剪其中一部分内容等功能。

注：

这个代码有一个bug，训练的时候loss值有一定的概率会变为nan，个人在训练时候的经验是在Config.py文件中，要修改batch_size大小，越大出现的概率越小，原因应该是部分训练集特征比较分散，导致预测结果得分相差较大，在计算损失函数有一个计算e的次方，导致溢出，这是个人看法，不清楚是否正确。

以上是个人的理解，如果帮到你了，希望能够在github上star一下，谢谢啦。