pytorch实现ssd之loss函数的设计

对于Loss函数的设计，我现在是越来越感觉到比模型搭建更难，虽然在理论上，就是几个loss函数的计算，比如交叉熵啥的，理论看起来很简单，但是对于ssd，yolo这些，真正要设计出来loss的计算，里面的细节太多也太杂了，让我很头大，但是没法，只能硬着头皮啃呗， 其实在现在这个时候，也就是我敲字的时候，我对于整个loss函数也只能说懂了百分之八十，还有一些细节，我依旧还是比较模糊，这也是写这篇文章的目的，在梳理我自己的思路的同时，也希望能帮到大家，开题先感谢一下ssd解释很好的作者，大家有兴趣或者看我的有些不明白，可以转向上面这个链接，我的整个思路的梳理，这个作者对我帮助很大，我也只是站在他的肩膀上写出我对整个loss的理解。

在上一篇文章中，我们知道整个网络最后搭建出来的模型最后的输出数据维度是类似[2, 38, 38, 84]这种，不是很清楚的可以看看上篇文章我们先看看整个流程图：

我们先看看第一步，第一步的原理如下：
代码如下：

'''this code was desinged by nike hu 生成先验框'''
first_box = [30, 60, 111, 162, 213, 264] # 这个是默认的六个尺度的最初边框，是根据计算来的，见图片可知过程
second_box = [60, 111, 162, 213, 264, 315] # 这个相当于first_box[k+1]
ar = [[2], [2, 3], [2, 3], [2, 3], [2], [2]] # 这个是针对不同尺度的featuremap来进行比例计算，详情见图片
featuremap_w = [38, 19, 10, 5, 3, 1] # 这是输入的特征图的尺寸
imagew = 300 # 这个是最初的照片的尺寸

class PriBox(nn.Module):
    def __init__(self):
        super(PriBox, self).__init__()

    def forward(self):
        box = []
        for i in range(len(featuremap_w)): # 这就是对每个尺度的特征图就行遍历
            for x, y in product(range(featuremap_w[i]), repeat=2):
                '''
                # product(A, B)函数，返回A、B中的元素的笛卡尔积的元组iterables是可迭代对象,repeat指定iterable重复几次,即:
                # product(A,repeat=3)等价于product(A,A,A)
                # product('ABCD', 'xy') --> Ax Ay Bx By Cx Cy Dx Dy
                # product(range(2), repeat=3) --> 000 001 010 011 100 101 110 111
                '''
                x = (x + 0.5) / featuremap_w[i] # 将x数值缩放到0到1之间，其他人这里不是这样操作的，他们是根据特征图相对于最初图的缩小倍数来计算的，
                # 但是我觉得不行，因为第四层是缩小了四倍，最后尺寸按照理论是37.5，但是最后尺寸是38，当然这里差距应该不会太大
                y = (y + 0.5) / featuremap_w[i] # 将y数值缩放到0到1之间
                '''对应w和h，有四种情况和六种情况长宽比的，详情看图片'''
                '''对每个模型输出的特征图，有两个默认的box长宽比都为1，但是计算方法不一样，详情将图片'''
                w1 = first_box[i] / imagew * 1 # 等比例缩放,本来firstbox里面的值就是相对于最初的图片的尺寸得到的，现在将这个值放到0到1之间
                h1 = w1
                w2 = math.sqrt(first_box[i] * second_box[i]) / imagew  * 1
                h2 = w2
                '''剩下的几种情况，如下'''
                box += [x, y, w1, h1]
                box += [x, y, w2, h2]
                for w in ar[i]:
                    box += [x, y, w1 * math.sqrt(w), w1 / math.sqrt(w)]
                    box += [x, y, w1 / math.sqrt(w), w1 * math.sqrt(w)] # 这里跟上面这行反过来，是因为长宽比反过来了
        box = torch.tensor(box).view(-1, 4) # 最后的维度为[8732, 4]
        # print(box.shape)
        return box

好了，第一步我们完成了，这一步最后输出的维度是[8732, 4]，接下来我们开始进行第二步，在这一步中，我们得设计到一个解码的操作，这里得解码，是指的真实的标签框和先验框共同处理，得到一个转化的关系，原理如下图：
代码如下：

'''@author: nike hu'''
# 编码公式见图片,第一个参数是真实框，x1, x2, y1,y2,第二个参数是先验框x,y,w,h
def encode(box1, box2, thred):
    xy = (((box1[:, :2] + box1[:, 2:])) / 2 - box2[:, :2]) / (box2[:, 2:] * thred[0])
    #((box1[:, :2] + box1[:, 2:])) / 2 - box2[:, :2])是两点求中间值得公式，比如x1=1, x2=3,那么x1和x2的中间值就是（1+3）/2, 最后xy维度是-> [8732, 2]
    wh = torch.log(1e-10 + (box1[:, 2:] - box1[:, :2]) / box2[:, 2:]) / thred[1]
    loc = torch.cat((xy,wh), dim=1) # 维度为[8732, 4]
    return loc

处理上面这个encode，我们还得设计一个函数来求解iou，所谓的Iou就是针对两个边框而言的，iou就是两个边框重合部分占总面积的比例，原理如下：
代码如下：

def get_iou(box1, box2):
    '''我们这里批量处理box1和box2,处理方式为维度的扩充，box1指的是边框，维度是[n, 4]，box2指的是box, 维度是[m, 4],最后输出[n,m]
    @author: nike hu'''
    box1_num = box1.size(0) # 这是box1的个数，相当于有多少个边框
    box2_num = box2.size(0) # 这是box2的个数，相当于有多少个box，默认是8732个
    box_unit1 = torch.max(box1[:, :2].unsqueeze(1).expand(box1_num, box2_num, 2),
                        box2[:, :2].unsqueeze(0).expand(box1_num, box2_num, 2)) # 两边框相交的左上角
    box_unit2 = torch.min(box1[:, 2:].unsqueeze(1).expand(box1_num, box2_num, 2),
                        box2[:, 2:].unsqueeze(0).expand(box1_num, box2_num, 2)) # 两边框相交的右下角
    inter_area = box_unit2 - box_unit1 # 求出交叉面积的长和宽
    inter_area = inter_area.clamp(0) # 把小于0的数值全部变为0
    intersection_area = inter_area[:, :, 0] * inter_area[..., 1] # 这个就是交叉面积
    box1_area = ((box1[:, 2] - box1[:, 0]) * (box1[:, 3] - box1[:, 1])).unsqueeze(1).expand(box1_num, box2_num)
    box2_area = ((box2[:, 2] - box2[:, 0]) * (box2[:, 3] - box2[:, 1])).unsqueeze(0).expand(box1_num, box2_num)
    iou = intersection_area/ (box2_area + box1_area - intersection_area)
    return iou

知道这两个基本函数之后，我们就可以进行接下来的处理了，对先验框和标签框进行处理。需要注意的如下：

代码如下：

# 匹配标签框和最佳的box,最后返回的是每张照片每个box对应的校准后的最佳标签框，以及每个box的最可能的类别，如果这个值为0代表这个box没有类别对应是负样本
def match(thread, pri_box, truths, lables, loc_t, conf_t, idx):
    '''
    @author: nike hu
    :param thread:IOU阈值，小于阈值设为0
    :param pri_box:先验框， shape[M,4]
    :param truths:标签框的x1, y1, x2, y2, [N, 4]
    :param lables:图片的所有类别，shape[num_obj]
    :param variance:prior的方差,默认为(0.1, 0.2}
    :param loc_t:用于填充encoded loc 目标张量, shape[batchsize, M, 4]
    :param conf_t:用于填充encoded conf 目标张量,shape[batchsize, N, M]
    idx: 现在的batch index
    要注意，一个标签框可以对应多个box，但是一个box只能对应一个标签框
    '''
    ious = get_iou(truths, pri_box) # 先求得iou值，维度为[n, m]
    pri_iou, pri_id = ious.max(dim=1, keepdim=True)
    # 这里是每一行的最大值和下标，意思就是每个标签框对应的最佳的box的值和下标,维度变为[n, 1],数值对应0-8732
    truth_iou, truth_id = ious.max(dim=0, keepdim=True) # 这里得意思是每个box对应的最佳的下标,维度变为[1, m]，数值对应0-n
    pri_id = pri_id.squeeze(1) # [n, 1]=>[n]
    truth_id = truth_id.squeeze(0) # [1, m]=>[m]
    for i in range(pri_id.shape[0]):
        truth_id[pri_id[i]] = i # 每个box对应的最佳的下边框可能并不是最佳的，这里来纠正
    matches = truths[truth_id] # 这里是指每个box对应的最佳标签框，维度为[m, 4]
    lable = (lables[truth_id] + 1).squeeze(-1) # 这里是每个box对应的最佳边框的类别，维度为[m]
    lable[(truth_iou < thread)[0]] = 0 # 将不符合要求的，小于置信度概率的部分的box对应的类别设置为0，代表负样本
    loc = encode(matches, pri_box, (0.1, 0.2))
    # 这个函数就是将真实边框和预测边框再进行处理，第一个参数的意思是每个box对应的最佳边框的x1, y1, x2, y2,第二个参数是box的坐标x,y,w,h
    loc_t[idx] = loc # 对每张图片进行数据存储，每张照片里面每个box可能对应的x,y,w,h
    conf_t[idx] = lable # 每张照片每个box可能对应的类别

经过以上处理之后，我们就可以开始干第三步了，这一步，原理看图：
代码如下：

'''@author: nike hu'''
class Loss(nn.Module):
    def __init__(self):
        super(Loss, self).__init__()

    def forward(self, prediction, prio_box, target):
        '''这里得第一个参数是主干网络的输出，(batchsize, 38, 38, 6*4)这种
        prio_box就是预测框，维度为[8732, 4], target就是标签框的矩阵，维度为[batchsie, n, 5]'''
        pred_loc, pred_conf = prediction[0], prediction[1] # 将输出的位置特征图和类别特征图区分出来
        pred_loc = torch.cat([i.view(i.shape[0], -1, 4) for i in pred_loc], dim=1) # 将三个尺度的位置特征图进行合并并且维度改变=》[batchsize, 8732, 4]
        pred_conf = torch.cat([i.view(i.shape[0], -1, class_num) for i in pred_conf], dim=1) # 将三个尺度的类别特征图进行合并并且维度改变=》[batchsize, 8732, num_class]
        # print(pred_loc.shape, pred_conf.shape)
        target_loc = target[:,:, :4] # ->[b, n, 4]
        target_conf = target[:, :, -1] # ->[b, n, 1]
        '''接下来先将先验框和标签框进行处理'''
        batches = pred_conf.size(0) # 总的照片个数
        truth_loc = torch.zeros_like(pred_loc, requires_grad=False) # 这个用来接收先验框和标签框处理后的位置数据
        truth_conf = torch.zeros(batches, pred_loc.size(1), requires_grad=False) # # 这个用来接收先验框和标签框处理后的类别数据,维度为[batchsize,M]
        for i in range(batches):
            target_loc_now = target_loc[i] # 第i张照片的标签位置
            target_conf_now = target_conf[i] # 第i张照片对应的类别
            match(0.5, prio_box, target_loc_now, target_conf_now, truth_loc, truth_conf, i)
        truth_conf = truth_conf.long() # 数据类型转化为整形，因为这里数值的含义是类别
        '''经过以上操作，我们求得了每个box对应的最佳的标签框后的x,y,w,h，以及每个box对应的最佳类别，如果这里面数值为0代表是负样本，类别为背景,
        接下来我们将预测的Predic和上面操作得到的数值进行操作'''
        pos = truth_conf > 0 # 这个代表正样本，维度为[batchsize, M]
        pos_id = pos.unsqueeze(-1).expand_as(pred_loc) # 维度进行扩充，便于后面进行处理->[batchsize, M, 4]
        pred_loc_use = pred_loc[pos_id].view(-1, 4) # 取出正样本的预测坐标
        truth_loc_use = truth_loc[pos_id].view(-1, 4) # 取出正样本的box坐标
        loss_loc = torch.nn.functional.smooth_l1_loss(pred_loc_use, truth_loc_use) # 计算坐标之间的Loss
        '''计算了loc_los之后，我们接下来计算conf的loss操作如下'''
        pred_conf_use = pred_conf.view(-1, class_num) # 将所有照片的预测值进行处理，维度为[bachsize*8732, class_num]
        pred_conf_really = log_sum_exp(pred_conf_use) - torch.gather(pred_conf_use, dim=1, index=truth_conf.view(-1, 1))
        # 上面这个才是将每个box得到的置信度求出来,维度为[bachsize*8732, 1]
        pred_conf_really = pred_conf_really.view(batches, -1) # ->[batchsize, 8732]
        pred_conf_really[pos] = 0 # 将正样本设置为0，其余全为负样本
        print(pred_conf_really.shape)
        sort_pred_idx = torch.sort(pred_conf_really, dim=1, descending=True)[1] # 从高到低进行排序
        sort_id = torch.sort(sort_pred_idx, dim=1)[1] # 对上面所得的值进行排序，就可以得到pred_conf_really每个数值对应的等级，维度为[b, 8732]
        print(sort_id)
        num_pos = pos.long().sum(dim=1, keepdim=True) # 每张照片对应的正样本的数量, 维度为[b, 1]
        # num_eng = torch.clamp(3*num_pos, max=pos.size(1)-1) # 这里得3的意思是正负样本按照1:3的比例
        num = sort_id < 3 * num_pos # 这里得3的意思是正负样本按照1:3的比例, 维度为[b, 8732]
        # print(num)
        pos_index = pos.unsqueeze(-1).expand_as(pred_conf) # [batchsize, M]=>[batchsize, M, num_class]
        num_index = num.unsqueeze(-1).expand_as(pred_conf) # [batchsize, M]=>[batchsize, M, num_class]
        conf_data = pred_conf[(pos_index + num_index).gt(0)].view(-1, class_num) # 将网络的预测值得正负样本都抽出来
        truth_conf_data = truth_conf[(pos + num).gt(0)] # 将先验框和标签框得到的每个box对应的最佳边框的类别正负样本取出来
        loss_conf = torch.nn.functional.cross_entropy(conf_data, truth_conf_data) # 计算多分类的loss
        num_sum = num_pos.data.sum() # 统计所有图片中的正样本的数量
        loss_conf /= num_sum
        loss_loc /= num_sum
        return loss_loc, loss_conf

# 下面这个函数是求置信度需要的公式，详情见图
def log_sum_exp(x):
    '''输入的维度是[bachsize*8732, class_num]'''
    x_max = x.detach().max() # torch中的max函数如果不限制维度的话，相当于将所有值中选一个最大值,detach相当于只是把矩阵中的数值取出来
    y = torch.log(torch.sum(torch.exp(x - x_max), dim=1, keepdim=True)) + x_max # [bachsize*8732, class_num]->[bachsize*8732, 1]
    return y

好了，以上原理的截图，均是我在最上面给的链接的作者的，那作者写的实在是太好了，关于原理部分我实在找不到补充了，只能在代码里面尽量为大家多注释一些，多写一点我的理解，希望能帮到各位。代码有些多，需要耐心的去看，没办法，我也是搞了好久才大致的将整个过程清楚了，没有什么简单的捷径方法，光是看看原理就简单，但是真实要代码复现出来，很复杂，里面很多细节。感谢能看到最后的朋友的支持，我相信，大多数看见这么多代码，会直接退出去，正常正常，因为我也这样，不到不得已，也不愿意看其他人写的代码，更不要说一行一行代码去理解了。

2020 4.19