首先先放下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一下，谢谢啦。
