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
扫一扫
2万+
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
