CornerNet代码解析——损失函数

CornerNet代码解析——损失函数

前言

今天要解析的是CornerNet的Loss层源码，论文中Loss的解析在这：CornerNet的损失函数原理

总体损失

总体的损失函数如下图所示，三个输出分别对应三部分损失，每部分损失有着对应的权重。接下来分别讲述每一块的损失。

源码中将Loss写成一个类：class AELoss，在CornerNet\models\py_utils\kp.py中.

class AELoss(nn.Module):
    
    def __init__(self, pull_weight=1, push_weight=1, regr_weight=1, focal_loss=_neg_loss):
        super(AELoss, self).__init__()
        # pull_weight = α
        self.pull_weight = pull_weight
        # push_weight = β
        self.push_weight = push_weight
        # regr_weight = γ
        self.regr_weight = regr_weight
        # 这其实就是heatmap的loss
        self.focal_loss  = focal_loss
        # 这其实就是embedding的loss
        self.ae_loss     = _ae_loss
        # 这其实就是offset的loss
        self.regr_loss   = _regr_loss

    def forward(self, outs, targets):
        stride = 6
        # ::跳着选
        '''
        首先明确两个输入：outs和targets
        outs：这是网络的预测结果，outs是一个列表，列表维度为12，outs[0::stride]这些是表示列表的
        切片操作，意思是隔stride(6)个跳着选。举个例子outs = [1,2,3,4,5,6,7,8,9,10,11,12]，
        outs[0::6]=[1, 7],其实这12个事6个两两成对，也就是左上角的heatmap有两个，右下角的heatmap有两个
        左上角的embedding有两个，右下角的embedding有两个，左上角的offset有两个，右下角的offset有两个，
        共12个，为什么要两份？应该跟上面的nstack有关，上述的nstack=2，所以循环出来outs不是6，而是12，
        映射到论文就是跟这句话：we also add intermediate supervision in training。这是中继监督，具体是啥
        我也还在看。也就是说下面的6个都是列表，每个列表里面都含有两个tensor，具体维度如下：
        '''
        # 两个都是[batch_size, 类别数, 128, 128]
        tl_heats = outs[0::stride]
        # 两个都是[batch_size, 类别数, 128, 128]
        br_heats = outs[1::stride]
        # 两个都是[batch_size, 128, 1]
        tl_tags  = outs[2::stride]
        # 两个都是[batch_size, 128, 1]
        br_tags  = outs[3::stride]
        # 两个都是[batch_size, 128, 2]
        tl_regrs = outs[4::stride]
        # 两个都是[batch_size, 128, 2]
        br_regrs = outs[5::stride]
        '''
        targets是gt，标准答案,也是个列表，但就只有下面5个，没有两份
        具体维度如下
        '''
        # [batch_size, 类别数, 128, 128]
        gt_tl_heat = targets[0]
        # [batch_size, 类别数, 128, 128]
        gt_br_heat = targets[1]
        # [3, 128]
        gt_mask    = targets[2]
        # [3, 128, 2]
        gt_tl_regr = targets[3]
        # [3, 128, 2]
        gt_br_regr = targets[4]

上述就是传入的预测值和真实值，Loss也就是计算预测的和真实之间的误差，当Loss值越小，那么说明网络预测的结果越好。接下去有了预测和真实值，具体分析三个部分的Loss。

1、Heatmap的损失

Heatmap损失的理论理解在这，接下来是源码理解：

这部分代码在CornerNet\models\py_utils\kp.py中

# focal loss
        focal_loss = 0
        # 到这里将heatmap经过sigmoid，将值映射到0-1之间，变成keypoint的响应值，还是列表，
        # 维度还是[batch_size, 类别数, 128, 128]
        tl_heats = [_sigmoid(t) for t in tl_heats]
        br_heats = [_sigmoid(b) for b in br_heats]
        # 在CornerNet\models\py_utils\kp_utils.py中详细讲述了focal_loss，这个focal loss就是_neg_loss，形参有体现
        focal_loss += self.focal_loss(tl_heats, gt_tl_heat)
        focal_loss += self.focal_loss(br_heats, gt_br_heat)

接着去到CornerNet\models\py_utils\kp_utils.py中详细讲述focal_loss：

'''
首先清楚函数的输入：
preds是列表：(2,)，表示一个列表中含两个tensor，每个tensor的维度是(batch_size, 类别数, 128, 128)
gt是tensor：(batch_size, 类别数, 128, 128)
'''
def _neg_loss(preds, gt):
    # pos_inds是0、1tensor，维度[3,7,128,128]。
    # eq函数是遍历gt这个tensor每个element，和1比较，如果等于1，则返回1，否则返回0
    pos_inds = gt.eq(1)
    # otherwise则是表明ycij第c个通道的（i,j）坐标上值不为1
    # 遍历gt这个tensor每个element，和1比较，如果小于1，则返回1，否则返回0
    neg_inds = gt.lt(1)
    # 总结下上面两个变量：上面这两个0-1位置互补
    # 回头看这两个变量，再结合公式1，公式1后面有两个判断条件：if ycij=1 and otherwise
    # 这里就是那两个判断条件，ycij=1表示第c个通道的（i,j）坐标上值为1，也即是gt中这个位置有目标
    # 也就是pos_inds是ycij=1，neg_inds是otherwise

    # torch.pow是次幂函数,其中gt[neg_inds]表示取出neg_inds中值为1的gt的值
    # 所以gt[neg_inds]就变成一个向量了，那么维度就等于neg_inds中有多少为1的
    # 可以neg_inds.sum()看看，1 - gt[neg_inds]就是单纯的用1减去每个element，
    # 然后每个element开4次方，就成了neg_weights，这个neg_weights是一维向量
    # 把gt中每个小于1的数字取出来，然后用1减去，在开方，那不是更小了，
    # 就是原来就很小，现在又降权。
    # gt[neg_inds]就是公式(1)中的Ycij
    # neg_weights就是公式(1)中的(1-ycij)^β，β就是4
    neg_weights = torch.pow(1 - gt[neg_inds], 4)

    loss = 0
    # 循环2次，因为preds是一个列表，有2部分，每部分放着一个tensor，每个tensor的
    # 维度为[batch_size，类别数，128，128]，也就是pred维度为[batch_size，类别数，128，128]
    for pred in preds:
        # 首先记住pos_inds中的1就是gt中有目标的地方，neg_inds中的1是gt中没有目标的地方
        # 将gt认为有目标的地方，pred也按这个地方取出数值，变成向量，pos_inds有多少个1，
        # pos_pred就多少维(一行向量)
        pos_pred = pred[pos_inds]
        # 将gt认为没有目标的地方，pred也按这个地方取出数值，变成向量，neg_inds有多少个1，
        # neg_pred就多少维(一行向量)
        neg_pred = pred[neg_inds]
        # 以上出现的pos_xxx, neg_xxx，命名的意思就是正样本positive和负样本negative


        # 这里对应的是论文中的公式(1)，也就是heatmap的loss
        # 可以先根据公式把相应的变量确认下：pos_pred就是公式中的Pcij。
        # neg_pred就是公式中的要经过二维高斯的Pcij，neg_weights就是(1-ycij)^β
        pos_loss = torch.log(pos_pred) * torch.pow(1 - pos_pred, 2)
        neg_loss = torch.log(1 - neg_pred) * torch.pow(neg_pred, 2) * neg_weights


        # gt的那个tensor中，值为1的个数，num_pos对应公式(1)中的N
        num_pos  = pos_inds.float().sum()
        # 累加
        pos_loss = pos_loss.sum()
        neg_loss = neg_loss.sum()
        # pos_pred是一维的。统计pos_pred中的元素个数，单纯的数个数而已，
        # 就算pos_pred中值为0的，也算一个
        if pos_pred.nelement() == 0:
            loss = loss - neg_loss
        else:
            # 用减号体现公式(1)中的-1
            loss = loss - (pos_loss + neg_loss) / num_pos
    # 返回最终的heatmap的loss
    return loss

2、Embedding的损失

Heatmap损失的理论理解在这，接下来是源码理解：

接着回到CornerNet\models\py_utils\kp.py，看怎么调用embedding的loss：

# tag loss
        # 初始化为0
        pull_loss = 0
        push_loss = 0
        # tl_tags、br_tags是列表，里面有两个tensor，每个tensor的维度为[batch_size, 128, 1]
        # 论文中说到的embedding是一维向量。也就是说，维度表示：一个batch_size一张图，用128*1的矩阵表示？？
        # 那么这个for循环，循环2次，每次进去的是[batch_size, 128, 1]的tl_tag, br_tag
        for tl_tag, br_tag in zip(tl_tags, br_tags):
            pull, push = self.ae_loss(tl_tag, br_tag, gt_mask)
            pull_loss += pull
            push_loss += push
        # 算出来的loss乘以相应的权重
        pull_loss = self.pull_weight * pull_loss
        push_loss = self.push_weight * push_loss

接着去到CornerNet\models\py_utils\kp_utils.py中详细讲述ae_loss：

'''
embedding的损失
输入：tag0、tag1为左上右下各一个[batch_size, 128, 1]的tensor，再来一个gt中的mask，这个mask是
0、1矩阵，维度[batch_size, 128],也就是一张图用128维来表示？？？？？？
'''
def _ae_loss(tag0, tag1, mask):
    # mask是[batch_size, 128]，这个就是第一维全部相加(sum)，就是把每个batch的128个数字相加，所以num的
    # 维度是[batch_size, 1]，1是128个数字的值相加变成一个数字，而mask还是0-1矩阵，所以这个num代表了
    # 每张图有多少个1.这个num代表公式(4)和(5)中的N
    num  = mask.sum(dim=1, keepdim=True).float()
    # 先看torch.squeeze() 这个函数主要对数据的维度进行压缩，去掉维数为1的的维度
    # 所以tag0和tag1的维度变成了[batch_size, 128]，和mask一样
    # tag0就是公式(4)中的etk
    tag0 = tag0.squeeze()
    # tag0就是公式(4)中的ebk
    tag1 = tag1.squeeze()
    # 单纯的求平均而已，这个tag_mean对应公式(4)和(5)中的ek,维度不变
    tag_mean = (tag0 + tag1) / 2
    
    # 这里能够体现是同类别的，因为累加只有一次，也就是Lpull用来缩小
    # 同类别左上右下角点的embedding vector的距离
    # 公式(4)前半段
    tag0 = torch.pow(tag0 - tag_mean, 2) / (num + 1e-4)
    # 这句能体现累加，这里tag0已经是单个数字
    tag0 = tag0[mask].sum()
    # 公式(4)后半段
    tag1 = torch.pow(tag1 - tag_mean, 2) / (num + 1e-4)
    # 这句能体现累加，这里tag1已经是单个数字
    tag1 = tag1[mask].sum()
    # 总的Lpull
    pull = tag0 + tag1
    
    # Lpush
    # 这里能够体现是不同类别的，因为累加有两次，公式(5)中的j不等于k,也就是Lpush用来扩大
    # 不同类别左上右下角点的embedding vector的距离
    # 这时候mask的维度由[3,128]-->[3,128,128]
    mask = mask.unsqueeze(1) + mask.unsqueeze(2)
    # 遍历mask这个tensor每个element，和2比较，如果等于2，则返回1，否则返回0，但为啥是2呢？
    mask = mask.eq(2)
    # num的维度[3, 1]-->[3, 1, 1]
    num  = num.unsqueeze(2)
    # num2的维度[3, 1, 1],num2表示公式(5)中的N(N-1)
    num2 = (num - 1) * num

    # dist是公式(5)中绝对值之间的运算
    # dist维度[3, 128, 128]=[3, 1, 128]-[3, 128, 1]
    dist = tag_mean.unsqueeze(1) - tag_mean.unsqueeze(2)
    # 1表示公式(5)三角形
    dist = 1 - torch.abs(dist)
    # 公式(5)就是relu，所以计算方式直接套relu
    dist = nn.functional.relu(dist, inplace=True)
    dist = dist - 1 / (num + 1e-4)
    dist = dist / (num2 + 1e-4)
    # 这时候mask的维度[3,128,128]，dist维度[3,128,128]
    dist = dist[mask]
    # sum之后就变成一个数字了
    push = dist.sum()
    # 返回两个loss，两个tensor的数字
    return pull, push

3、Offset的损失

Heatmap损失的理论理解在这，接下来是源码理解：

接着回到CornerNet\models\py_utils\kp.py，看怎么调用offset的loss：

# offsets loss
        regr_loss = 0
        # tl_regrs、br_regrs是列表，里面有两个tensor，每个tensor的维度为[batch_size, 128, 2]
        # 维度表示：一个batch_size一张图，用128*2的矩阵表示？？
        # 那么这个for循环，循环2次，每次进去的是[batch_size, 128, 2]的tl_regr, br_regr
        for tl_regr, br_regr in zip(tl_regrs, br_regrs):
            regr_loss += self.regr_loss(tl_regr, gt_tl_regr, gt_mask)
            regr_loss += self.regr_loss(br_regr, gt_br_regr, gt_mask)
        regr_loss = self.regr_weight * regr_loss

        # 总的loss
        loss = (focal_loss + pull_loss + push_loss + regr_loss) / len(tl_heats)
        # unsqueeze(i) 表示将第i维设置为1维
        return loss.unsqueeze(0)

接着去到CornerNet\models\py_utils\kp_utils.py中详细讲述regr_loss：

'''
输入：regr偏移量，维度[batch_size, 128, 2],gt_regr维度[batch_size, 128, 2]
mask维度[batch_size, 128]
'''
def _regr_loss(regr, gt_regr, mask):
    # 公式(3)的N
    num  = mask.float().sum()
    # mask.unsqueeze(2)维度[batch_size, 128, 1]
    # mask的维度[batch_size, 128, 2]
    mask = mask.unsqueeze(2).expand_as(gt_regr)
    # 取出mask中1对应的位置，然后在预测的偏移量和真实的偏移量中取出这些位置的值
    # 此时二者的维度变为一维向量
    regr    = regr[mask]
    gt_regr = gt_regr[mask]
    # 直接调用自带的SmoothL1Loss
    regr_loss = nn.functional.smooth_l1_loss(regr, gt_regr, size_average=False)
    # 最后除N
    regr_loss = regr_loss / (num + 1e-4)
    return regr_loss