个人博客：http://www.chenjianqu.com/
原文链接：http://www.chenjianqu.com/show-92.html
上篇博客<SSD源码阅读一>阅读了ssd.pytorch项目的部分代码，包括包括数据集、数据增强和构建网络结构的部分。在模型类class SSD(nn.Module)中，构造函数定义这个变量self.priorbox = PriorBox(self.cfg) ，PriorBox是用于构造先验框的类，代码如下：
layers/functions/prior_box.py
class PriorBox(object):
"""Compute priorbox coordinates in center-offset form for each source feature map.
"""
#cfg是定义在data/config.py里面，是模型的相关配置
def __init__(self, cfg):
super(PriorBox, self).__init__()
#输入数据分辨率，一般为300
self.image_size = cfg['min_dim']
#输出特征图的数量
self.num_priors = len(cfg['aspect_ratios'])
#尺度
self.variance = cfg['variance'] or [0.1]
#各个输出特征图的分辨率
self.feature_maps = cfg['feature_maps']
#计算先验框用到的Smin和Smax
self.min_sizes = cfg['min_sizes']
self.max_sizes = cfg['max_sizes']
#各个输出特征图每个像素对应到原图的大小
self.steps = cfg['steps']
#输出特征图用到的比例
self.aspect_ratios = cfg['aspect_ratios']
self.clip = cfg['clip']
#VOC或COCO
self.version = cfg['name']
for v in self.variance:
if v <= 0:
raise ValueError('Variances must be greater than 0')
def forward(self):
mean = []
for k, f in enumerate(self.feature_maps):#f是特征图的边长
for i, j in product(range(f), repeat=2):#product(range(f), repeat=2)求range(f)和range(f)的笛卡尔积，生成各个像素的坐标
'''
每个先验框中心的计算公式是：(cx,cy)=((i+0.5)/|fk|, (j+0.5)/|fk|)，其中|fk|是特征图的边长
得到的坐标范围是[0,1]
'''
f_k = self.image_size / self.steps[k]#steps是输出特征图的下采样率，则f_k是特征图的边长
# unit center x,y
cx = (j + 0.5) / f_k
cy = (i + 0.5) / f_k
’‘’
每个特征图使用的先验框大小：
Sk=Smin+(Smax-Smin)(k-1)/(m-1)，k的值[1,m]，m是输出特征图的数量，
Sk是先验框相对于整张图片的比例，Smin=0.2，Smax=0.95
先验框的宽度w_k_a=Sk*ar^0.5，高度h_k_a=Sk/ar^0.5，
其中ar是特征图的比例，有{3, 2, 1, 1/2, 1/3}
‘’‘
#下面计算各个比例的先验框的的宽高，并把它们放到list()里面
#注意这里的宽高是相对于原图的比例
# aspect_ratio: 1
# rel size: min_size
s_k = self.min_sizes[k]/self.image_size
mean += [cx, cy, s_k, s_k]
# aspect_ratio: 1
# rel size: sqrt(s_k * s_(k+1))
s_k_prime = sqrt(s_k * (self.max_sizes[k]/self.image_size))
mean += [cx, cy, s_k_prime, s_k_prime]
# rest of aspect ratios
for ar in self.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)]
#reshape先验框
output = torch.Tensor(mean).view(-1, 4)
#clamp_ 将output张量每个元素的夹紧到区间 [min,max]，并返回结果到一个新张量。其实就是我们常用的截断函数（分段阈值）。
if self.clip:
output.clamp_(max=1, min=0)
return output
上面代码的代码严格按照论文<SSD论文笔记>中先验框的计算，forward()函数的中mean变量保存了一张图片中各个先验框的中心坐标和长宽(均是相对值)。先验框主要是计算损失函数的时候用到，前向传播的时候，直接对每张图片拷贝mean张量即可。
回到SSD网络结构，根据原论文的说法，与其他层相比，conv4_3具有不同的特征尺度，因此使用ParseNet中介绍的L2 normalization技术将特征图中每个位置的feature norm缩放到20，并在反向传播期间学习尺度。因此这里也实现了L2 normalization，如下：
layers/modules/l2norm.py
class L2Norm(nn.Module): #参数：输入特征图的通道数，缩放像素值到达的范围 def __init__(self,n_channels, scale): super(L2Norm,self).__init__() self.n_channels = n_channels self.gamma = scale or None self.eps = 1e-10 self.weight = nn.Parameter(torch.Tensor(self.n_channels)) self.reset_parameters() def reset_parameters(self): init.constant_(self.weight,self.gamma) def forward(self, x): norm = x.pow(2).sum(dim=1, keepdim=True).sqrt()+self.eps #x /= norm x = torch.div(x,norm) out = self.weight.unsqueeze(0).unsqueeze(2).unsqueeze(3).expand_as(x) * x return out
到此，网络结构部分基本完成。接下来，回到训练的主程序train.py train()里面，在构建网络之后，因为作者有多个GPU，因此它这里设置数据并行：
train.py train()
if args.cuda: net = torch.nn.DataParallel(ssd_net)#数据并行，多个GPU的时候使用 cudnn.benchmark = True
对于cudnn.benchmark，大部分情况下，设置这个 flag 可以让内置的 cuDNN 的 auto-tuner 自动寻找最适合当前配置的高效算法，来达到优化运行效率的问题。如果网络的输入数据维度或类型上变化不大，设置 torch.backends.cudnn.benchmark = true 可以增加运行效率；如果网络的输入数据在每次 iteration 都变化的话，会导致 cnDNN 每次都会去寻找一遍最优配置，这样反而会降低运行效率。
接下来，程序判断是从零开始训练还是恢复训练，如果是恢复训练则直接加载整个模型的权重，否则加载vgg的权重，然后将模型转换为GPU模式：
train.py train()
#权重加载
if args.resume:
print('Resuming training, loading {}...'.format(args.resume))
ssd_net.load_weights(args.resume)
else:
vgg_weights = torch.load(args.save_folder + args.basenet)
print('Loading base network...')
ssd_net.vgg.load_state_dict(vgg_weights)
#使用GPU模式
if args.cuda:
net = net.cuda()
如果是从零开始训练，那么对新增层的权重进行xavier初始化：
train.py train()
#权重初始化
if not args.resume:
print('Initializing weights...')
#使用xavier方法进行初始化，apply(fn)：将fn函数递归地应用到网络模型的每个子模型中进行参数的初始化。
ssd_net.extras.apply(weights_init)
ssd_net.loc.apply(weights_init)
ssd_net.conf.apply(weights_init)
然后设置优化器、设置损失函数：
train.py train()
#使用SGD优化器 optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) #定义损失函数 criterion = MultiBoxLoss(cfg['num_classes'], 0.5, True, 0, True, 3, 0.5,False, args.cuda)
根据论文原文<SSD论文笔记>，我们知道SSD的损失函数由定义损失和置信度损失组成，这里的代码实现如下，这里回顾以下SSD损失函数的计算：
其中定位损失各项的计算：
要计算上式中的损失函数，首先要确定xijk，即要将gt框与先验框进行匹配，匹配的策略：策略1：找到与gt bbox的IOU最大的先验包围框，这样保证每个gt bbox至少对应一个先验框。但是一个图片中gt bbox非常少，而先验框却很多，这样正负样本极其不平衡。策略2：对于其它的先验包围框，若与gt bbox的 IOU 大于某个阈值（一般是0.5），那么该先验框也与这个ground truth进行匹配。
首先知道IoU的实现：
layers/box_utils.py jaccard()
#注意这里的包围框格式是：[xmin,ymin,xmax,ymax]，即包围框左上角和右下角的坐标 #计算两组框中两两的交集，box_a.shape=[A,4]，box_b.shape=[B,4] def intersect(box_a, box_b): #box_a中框的数量 A = box_a.size(0) #box_b中框的数量 B = box_b.size(0) #box_a[:,2:]原本的维度是(A,2),经过unsqueeze(1)函数变为(A,1,2)，再经过expand(A,B,2)变为(A,B,2),该操作不会分配内存 t1=box_a[:, 2:].unsqueeze(1).expand(A, B, 2) t2=box_b[:, 2:].unsqueeze(0).expand(A, B, 2) max_xy = torch.min(t1,t2)#两个[xmax,ymax]之间的比较，求他们两点中小的点 #两个[xmin,ymin]之间的比较，求他们两个的大的点 n1=box_a[:, :2].unsqueeze(1).expand(A, B, 2) n2=box_b[:, :2].unsqueeze(0).expand(A, B, 2) min_xy = torch.max(n1,n2) #求出交集区域的宽高,inter.shape=[A,B,2] inter = torch.clamp((max_xy - min_xy), min=0) #求出交集的面积，返回.shape=[A,B] return inter[:, :, 0] * inter[:, :, 1] #计算两组包围框中两两包围框的IoU，box_a.shape=[A,4]，box_b.shape=[B,4] def jaccard(box_a, box_b): #计算两组包围框中两两包围框的交集面积，inter.shape=[A,B] inter = intersect(box_a, box_b) area_a = (box_a[:, 2]-box_a[:, 0]) * (box_a[:, 3]-box_a[:, 1]) #box_a中所有框的区域面积 area_a = area_a.unsqueeze(1).expand_as(inter) #将area_a的维度由[A,B]拓展为[A,B] area_b = (box_b[:, 2]-box_b[:, 0]) * (box_b[:, 3]-box_b[:, 1])#box_b中所有框的区域面积 area_b = area_b.unsqueeze(0).expand_as(inter) # [A,B] #IoU=A ∩ B / A ∪ B = A ∩ B / (area(A) + area(B) - A ∩ B) union = area_a + area_b - inter return inter / union #维度:[A,B], #返回矩阵的第i行、第j列元素的值代表box_a中的第i个包围框与box_b中的第j个包围框的IoU
然后就是匹配算法的实现：
layers/box_utils.py match()
#将box的 [框中心坐标的x,框中心坐标的y,框的宽,框的高]， #转换为 [框左上角坐标的x,框左上角坐标的y,框右下角坐标的x，框右下角坐标的y] def point_form(boxes): #[cx,cy]-[w/2,h/2]=[cx - w/2 , cy - h/2]=[xmin,ymin] return torch.cat( ( boxes[:, :2] - boxes[:, 2:]/2, # xmin, ymin boxes[:, :2] + boxes[:, 2:]/2 # xmax, ymax ), 1) #在列进行拼接 #GT框和先验框匹配算法实现 def match(threshold,#挑选正样本的阈值 truths, #当前图片中的gt框，shape:[num_truths,4] priors, #当前图片中的先验框,shape:[num_priors,4] variances, # labels, #当前图片中每个gt框的标签 loc_t,#匹配好的gt框 conf_t, #匹配好的置信度 idx#当前图片在batch的索引 ): #计算IOU,overlaps的维度是[A,B]，其中A是truths的数量，B是priors的数量 overlaps = jaccard( truths, point_form(priors)) #point_form(priors)将priors的格式[框中心坐标的x,框中心坐标的y,框的宽,框的高]， #转换为 [框左上角坐标的x,框左上角坐标的y,框右下角坐标的x，框右下角坐标的y] #计算每个gt框最匹配的先验框。即overlaps中每行数据中的最大值，一行代表一个gt box与所有先验框的IOU #shape为[1,B]，两个返回值分别代表gt框最大的IoU值和对应的先验框的索引 best_prior_overlap, best_prior_idx = overlaps.max(1, keepdim=True) #计算每个先验框最匹配的gt框，即计算overlaps中每列数据中的最大值，一列代表一个先验框与所有gt box的IOU，后面用于计算最大是否大于阈值 #shape为[1,B]，代表的是所有先验框的最大IOU和对应的gt框的索引 best_truth_overlap, best_truth_idx = overlaps.max(0, keepdim=True) best_truth_idx.squeeze_(0) #将shape从[1,B]变为[B] best_truth_overlap.squeeze_(0) best_prior_idx.squeeze_(1) #将shape从[A,1]变为[A] best_prior_overlap.squeeze_(1) ''' index_fill_(dim, index, val) 按照index，将val的值填充self的dim维度。如： >>> x = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=torch.float) >>> index = torch.tensor([0, 2]) >>> x.index_fill_(1, index, -1) tensor([[-1., 2., -1.], [-1., 5., -1.], [-1., 8., -1.]]) A是truths的数量，B是priors的数量 best_truth_overlap代表各个先验框与所有gt box的最大的IoU，shape=[B] best_prior_idx代表各个gt box的最大IoU匹配先验框的索引,shape=[A] 因此best_truth_overlap.index_fill_(0, best_prior_idx, 2)的效果是 设置每个gt框与最大IoU先验框的的IoU为2,这确保每个gt框匹配的prior框不会因为IoU太低，而被过滤掉 ''' best_truth_overlap.index_fill_(0, best_prior_idx, 2) # ensure best prior #确保每个gt框都能匹配到一个先验框。best_prior_idx.shape=[A]，里面的每个值代表一个gt框与最大IoU的先验框的索引。 for j in range(best_prior_idx.size(0)): idx=best_prior_idx[j] #第j个gt框 与各个先验框IoU最大 的 先验框索引 #best_truth_idx表示所有先验框与各个gt框IoU最大的gt框的索引 best_truth_idx[idx]=j #第j个gt框匹配的先验框设置其匹配的gt框索引为j #获取每个prior匹配的gt框 matches = truths[best_truth_idx] # Shape: [num_priors,4] #每个prior对应类别标签，加1是因为0是背景类,best_truth_idx.shape=[B] conf = labels[best_truth_idx] + 1 # Shape: [num_priors] #令IoU小于threshold的先验框的类别标签为0，即为背景类 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 #对位置进行编码，参数：每个先验框匹配的gt框，先验框，没搞懂variances这个参数 def encode(matched, priors, variances): #下面计算的变量均来自 定位损失函数里面的 #匹配好的gt框的中心和先验框中心的x、y的距离，即计算公式中的g_cxcy g_cxcy = (matched[:, :2] + matched[:, 2:])/2 - priors[:, :2] #计算公式中的g_cxcy_^ g_cxcy /= (variances[0] * priors[:, 2:]) #计算公式中的g_w和g_h g_wh = (matched[:, 2:] - matched[:, :2]) / priors[:, 2:] g_wh = torch.log(g_wh) / variances[1] return torch.cat([g_cxcy, g_wh], 1) # [num_priors,4]
上面就是先验框和gt框匹配的过程，经过这个步骤，就可以计算损失函数了，损失函数的计算放在layers/modules/multibox_loss.py里面的MultiBoxLoss类里面，来看损失函数计算的前半部分：
layers/modules/multibox_loss.py forward()
def forward(self, predictions, targets): ''' 预测的数据 loc_data：预测框位置点矩阵 size: (batch_size,num_priors,4) -> (32,8732,4) conf_data：预测框置信度矩阵 size:(batch_size,num_priors,num_classes) -> (32,8732,3) priors：先验框矩阵 size:(num_priors,4) -> (8732,4) ''' loc_data, conf_data, priors = predictions num = loc_data.size(0) #batch_size num_priors=loc_data.size(1) #每张图片预测的先验框数量 priors = priors[:num_priors, :] num_priors = (priors.size(0))#先验框的数量 num_classes = self.num_classes #类别数 #匹配先验框和gt框，gt框的格式为[:,xmin,ymin,xmax,ymax] 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 #第idx张图片的所有gt框 labels = targets[idx][:, -1].data #第idx张图片的所有gt框对应的标签，为整数 defaults = priors.data #先验框 #先验框匹配 match(self.threshold, truths, defaults, self.variance, labels,loc_t, conf_t, idx) #经过匹配，loc_t是每个先验框对应的gt框，维度是[32,8732,4]， #conf_t是先验框的标签,维度是[32,8732] if self.use_gpu: loc_t = loc_t.cuda() conf_t = conf_t.cuda() # wrap targets loc_t = Variable(loc_t, requires_grad=False) conf_t = Variable(conf_t, requires_grad=False) #如果某个先验框的标签>0，则设置pos标志位为true。pos的维度是[32,8732] pos = conf_t > 0 #每张图片中非背景类的先验框数，num_pos的维度[32,1] num_pos = pos.sum(dim=1, keepdim=True) # Shape:[batch,num_priors,4],即将pos维度扩展为[32,8732,4] pos_idx = pos.unsqueeze(pos.dim()).expand_as(loc_data) #选出先验框类别非0 对应的预测框和gt框 loc_p = loc_data[pos_idx].view(-1, 4) loc_t = loc_t[pos_idx].view(-1, 4) #定位损失使用smooth_L1 loss_l = F.smooth_l1_loss(loc_p, loc_t, size_average=False)
上面是计算定位损失的过程，需要注意的是，网络预测的是相对于包围框中心和宽度和高度相对于先验框的偏移值，因此也需要将先验框对应的gt框转换为gt框相对于先验框的偏移值，也就是box_utils.py encode()函数的作用。
定位损失只用到了positive输出框用于计算，而计算置信度损失的时候，不仅需要计算positive输出框作为正例，还需要其它输出框作为负例。但是剩余的预测框的数量太多，会导致正负样本极度不平衡，这会导致得到的模型偏向性十分严重，甚至不可用。因此我们需要从剩余的输出框中挑选出高置信度的作为负样本，这就是难例挖掘，SSD使用的正负样本比例为1:3。难例挖掘和置信度损失的代码如下：
layers/modules/multibox_loss.py forward()
#将[batch_size,8732,21] reshape为[batch_size*8732,21] batch_conf = conf_data.view(-1, self.num_classes) ''' batch_conf是该批次的所有预测框的类别置信度，[batch_size*8732, 21] conf_t是该批次的所有先验框的匹配的标签，[batch_size*8732, 1] gather函数的作用是沿着定轴dim(1),按照Index(conf_t.view(-1, 1))取出元素 batch_conf.gather(1, conf_t.view(-1, 1))得到矩阵的维度是[batch_size*8732,1],代表每个预测框的最匹配类的置信度 计算所有预测框的置信度损失 ''' loss_c = log_sum_exp(batch_conf) - batch_conf.gather(1, conf_t.view(-1, 1)) #难例挖掘 loss_c[pos] = 0 #过滤掉positive预测框 loss_c = loss_c.view(num, -1) #reshape为[batch_size,8732,1] _, loss_idx = loss_c.sort(1, descending=True) #将每张图片中所有的先验框 按照置信度从高到低排序 _, idx_rank = loss_idx.sort(1) num_pos = pos.long().sum(1, keepdim=True) #pos是bool，先将其转换为long，在求和，结果是每张图片中positive预测框的数量 ''' torch.clamp(input, min, max, out=None) → Tensor 将张量的每个元素夹紧到区间[min,max]内， | min, if x_i < min y_i = | x_i, if min <= x_i <= max | max, if x_i > max self.negpos_ratio*num_pos的shape是[batch_size,1]，下面语句的意思是将每张图片负样本的数量限制在pos.size(1)-1 ''' num_neg = torch.clamp(self.negpos_ratio*num_pos, max=pos.size(1)-1) #neg是难例挖掘最终得到的负样本索引， 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) ''' gt()函数是比较两个同维度张量的对应元素，比如C=A.gt(B)，若A[i,j]>B[i,j]，则C[i,j]==0 (pos_idx+neg_idx).gt(0)表示该样本是否是训练样本 ''' #挑出预测结果的正负样本 conf_p = conf_data[(pos_idx+neg_idx).gt(0)].view(-1, self.num_classes) #挑出匹配的gt标签 targets_weighted = conf_t[(pos+neg).gt(0)] #置信度损失使用交叉熵 #交叉熵的计算公式为loss=sum(label*log(predict)) 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() #batch中所有positive框的总数 loss_l /= N loss_c /= N return loss_l, loss_c
回到训练代码，前面已经定义了数据集和损失，接着就差不多可以开始训练了，剩余的训练代码如下：
train.py train()
net.train()#训练模式
# loss counters
loc_loss = 0
conf_loss = 0
epoch = 0
print('Loading the dataset...')
epoch_size = len(dataset) // args.batch_size
print('Training SSD on:', dataset.name)
print('Using the specified args:')
print(args)
step_index = 0
if args.visdom:
vis_title = 'SSD.PyTorch on ' + dataset.name
vis_legend = ['Loc Loss', 'Conf Loss', 'Total Loss']
iter_plot = create_vis_plot('Iteration', 'Loss', vis_title, vis_legend)
epoch_plot = create_vis_plot('Epoch', 'Loss', vis_title, vis_legend)
#定义数据迭代器
data_loader = data.DataLoader(dataset,
args.batch_size,
num_workers=args.num_workers,#处理数据集的线程数
shuffle=True,
collate_fn=detection_collate,#将一个list的sample组成一个mini-batch的函数
pin_memory=True #如果设置为True，那么data loader将会在返回它们之前，将tensors拷贝到CUDA中的固定内存（CUDA pinned memory）中.
)
# create batch iterator
batch_iterator = iter(data_loader)
#迭代训练
for iteration in range(args.start_iter, cfg['max_iter']):
#添加数据到visdom
if args.visdom and iteration != 0 and (iteration % epoch_size == 0):
update_vis_plot(epoch, loc_loss, conf_loss, epoch_plot, None,'append', epoch_size)
# reset epoch loss counters
loc_loss = 0
conf_loss = 0
epoch += 1
#学习率调整
if iteration in cfg['lr_steps']:
step_index += 1
adjust_learning_rate(optimizer, args.gamma, step_index)
#加在训练数据
images, targets = next(batch_iterator)
if args.cuda:
images = Variable(images.cuda())
targets = [Variable(ann.cuda(), volatile=True) for ann in targets]
else:
images = Variable(images)
targets = [Variable(ann, volatile=True) for ann in targets]
t0 = time.time()
#前向传播
out = net(images)
#反向传播
optimizer.zero_grad()
#计算损失函数
loss_l, loss_c = criterion(out, targets)
loss = loss_l + loss_c
loss.backward()
optimizer.step()
t1 = time.time()
loc_loss += loss_l.data[0]
conf_loss += loss_c.data[0]
if iteration % 10 == 0:
print('timer: %.4f sec.' % (t1 - t0))
print('iter ' + repr(iteration) + ' || Loss: %.4f ||' % (loss.data[0]), end=' ')
if args.visdom:
update_vis_plot(iteration, loss_l.data[0], loss_c.data[0],iter_plot, epoch_plot, 'append')
if iteration != 0 and iteration % 5000 == 0:
print('Saving state, iter:', iteration)
torch.save(ssd_net.state_dict(), 'weights/ssd300_COCO_' +repr(iteration) + '.pth')
torch.save(ssd_net.state_dict(),args.save_folder + '' + args.dataset + '.pth')
def adjust_learning_rate(optimizer, gamma, step):
"""Sets the learning rate to the initial LR decayed by 10 at every
specified step
# Adapted from PyTorch Imagenet example:
# https://github.com/pytorch/examples/blob/master/imagenet/main.py
"""
lr = args.lr * (gamma ** (step))
for param_group in optimizer.param_groups:
param_group['lr'] = lr
这部分代码没有什么太多需要解释的。
本篇博客详细解释了先验框计算、损失函数计算和模型的训练。
