【6PACK代码注解】损失函数

文章目录

• 前言
• 一、init
• 二、forward
• • 输入输出
  • 逐句解析
  • • 2.1 注意力网络损失函数
    • 2.2 多视图一致性损失&平移损失
    • • change_to_ver()
    • 2.3 旋转损失
    • • estimate_rotation()
    • 2.4 轮廓损失(表面损失)**L~sil~**
    • 2.5 分离损失**L~sep~**
• 三、eval中位姿预测
• • eval()
  • estimate_pose()

---

前言

【6PACK全记录】6-PACK论文学习及复现记录

一、init

train.py中criterion = Loss(opt.num_kp, opt.num_cates)部分将调用__init__函数，进行初始化。
输入参数：

 - num_key：关键点数
 - num_cate：类别数

代码结构：

def __init__(self, num_key, num_cate):
        super(Loss, self).__init__(True)
        self.num_key = num_key
        self.num_cate = num_cate

        self.oneone = Variable(torch.ones(1)).cuda() #将1作为1个结点，默认requires_grad = False

        self.normal = tdist.Normal(torch.tensor([0.0]), torch.tensor([0.0005]))#定义正态分布，均值0，标准差0.0005

        self.pconf = torch.ones(num_key) / num_key#shape为[8]的张量，各元素为0.125
        self.pconf = Variable(self.pconf).cuda()

        self.sym_axis = Variable(torch.from_numpy(np.array([0, 1, 0]).astype(np.float32))).cuda().view(1, 3, 1)#【1，3，1】--[[0],[1],[0]]
        self.threezero = Variable(torch.from_numpy(np.array([0, 0, 0]).astype(np.float32))).cuda()#[3]---[0,0,0]

        self.zeros = torch.FloatTensor([0.0 for j in range(num_key-1) for i in range(num_key)]).cuda()#56个0的tensor

        self.select1 = torch.tensor([i for j in range(num_key-1) for i in range(num_key)]).cuda()
        #先遍历i再遍历j---7组0，1，2....，7共56个元素
        self.select2 = torch.tensor([(i%num_key) for j in range(1, num_key) for i in range(j, j+num_key)]).cuda()
        #7组1,2...7,0每组轮转

        self.knn = KNearestNeighbor(1)#找最近的1个邻点

二、forward

train.py中_, item_score = criterion(Kp_fr, Kp_to, anc_fr, anc_to, att_fr, att_to, r_fr, t_fr, r_to, t_to, mesh, scale, cate)将调用forward函数

输入输出

输入参数：

 - Kp_fr:初始帧的Keypoints坐标
 - Kp_to：当前帧的Keypoints坐标
 - anc_fr：初始帧的anchor坐标
 - anc_to：当前帧的anchor坐标
 - att_fr：初始帧的置信分数
 - att_to：当前帧的置信分数
 - r_fr：初始帧的位姿R
 - t_fr：初始帧的位姿T
 - r_to：当前帧的位姿R
 - t_to:当前帧的位姿T
 - mesh:bbox网格点
 - scale:归一化系数
 - cate：当前训练类别的idx

输出：

 - loss：总损失函数，是6种损失函数的加权和(是variable，有梯度)
 - score：4种重要损失函数的加权和(是数值，无梯度)

---

逐句解析

def forward(self, Kp_fr, Kp_to, anc_fr, anc_to, att_fr, att_to, r_fr, t_fr, r_to, t_to, mesh, scale, cate):
        if cate.view(-1).item() in [2, 4, 5]:
            sym_or_not = False #2，4，5类没有旋转对称性---采用直角坐标系
        else:
            sym_or_not = True #0，1，3有旋转对称性---采用柱坐标系

        num_kp = self.num_key
        num_anc = len(anc_fr[0])

对于没有旋转对称性的类别，世界坐标系取直角坐标系；对于有旋转对称性的类别，世界坐标系转为柱坐标系。

2.1 注意力网络损失函数

############ Attention Loss
gt_t_fr = t_fr.view(1, 1, 3).repeat(1, num_anc, 1)#【1，125，3】
min_fr = torch.min(torch.norm(anc_fr - gt_t_fr, dim=2).view(-1))#前帧中离质心最近的距离beta
loss_att_fr = torch.sum(((torch.norm(anc_fr - gt_t_fr, dim=2).view(1, num_anc) - min_fr) * att_fr).contiguous().view(-1))

gt_t_to = t_to.view(1, 1, 3).repeat(1, num_anc, 1)
min_to = torch.min(torch.norm(anc_to - gt_t_to, dim=2).view(-1))
loss_att_to = torch.sum(((torch.norm(anc_to - gt_t_to, dim=2).view(1, num_anc) - min_to) * att_to).contiguous().view(-1))

loss_att = (loss_att_fr + loss_att_to).contiguous() / 2.0

2层MLP做注意力网络选择目标锚点时的损失函数L_anc

其中min_fr、min_to为锚点离质心最近的距离 β beta β

---

2.2 多视图一致性损失&平移损失

说明

符号	定义
Δ Delta Δtgt	groud truth 平移矩阵变化，通过两帧t的差值定义
Δ Delta Δt	预测平移矩阵变化，通过keypoints坐标均值的变化定义
Δ Delta ΔRgt	groud truth 旋转矩阵变化，通过当前帧相对位姿和初始帧实际位姿的矩阵乘Rto·RfrT定义
Δ Delta ΔR	预测旋转矩阵变化，通过svd分解求预测旋转矩阵

############# Different View Loss
#torch.bmm--矩阵乘
gt_Kp_fr = torch.bmm(Kp_fr - t_fr, r_fr).contiguous()#世界坐标系下的关键点坐标
gt_Kp_to = torch.bmm(Kp_to - t_to, r_to).contiguous()

if sym_or_not:#有旋转对称性
    ver_Kp_fr, cent_fr = self.change_to_ver(gt_Kp_fr)
    ver_Kp_to, cent_to = self.change_to_ver(gt_Kp_to)
    Kp_dis = torch.mean(torch.norm((ver_Kp_fr - ver_Kp_to), dim=2), dim=1)
    #一个数，两帧相应kp的距离的均值
    #多视图一致性损失
    Kp_cent_dis = (torch.norm(cent_fr - self.threezero) + torch.norm(cent_to - self.threezero)) / 2.0
    #一个数，两帧kp平均坐标的间距
    #平移损失
else:#无旋转对称性
    Kp_dis = torch.mean(torch.norm((gt_Kp_fr - gt_Kp_to), dim=2), dim=1)
    cent_fr = torch.mean(gt_Kp_fr, dim=1).view(-1).contiguous()
    cent_to = torch.mean(gt_Kp_to, dim=1).view(-1).contiguous()
    Kp_cent_dis = (torch.norm(cent_fr - self.threezero) + torch.norm(cent_to - self.threezero)) / 2.0

前两句通过R、T变换得到两帧的关键点在世界坐标系中的坐标gt_Kp（ground truth）。
得到的损失函数中，Kp_dis为多视图一致性损失L_mvc，即相邻帧关键点距离的均值

其中，k_i^t为视频t时刻帧种关键点的坐标，减号后面的部分代表将t-1时刻关键点坐标通过真实的相对位姿变化推测出一个t时的坐标，多视图一致性损失即为推测出的关键点与实际关键点的平均距离。
Kp_cent_dis为平移损失L_tra，即相邻帧关键点平均坐标的间距

其中，对于有旋转对称性的实例，需要先将坐标系转换为(d,h, θ theta θ)坐标，调用change_to_ver()函数，输出的ver_Kp为柱坐标系下每个关键点的坐标，各通道依次为关键点的z坐标、半径r、与前一个关键点的夹角；输出的cent为8个关键点平均坐标。函数的实现如下：

change_to_ver()

def change_to_ver(self, Kp):
        pconf2 = self.pconf.view(1, self.num_key, 1)#【1，8，1】均为0.125
        cent0 = torch.sum(Kp * pconf2, dim=1).view(-1).contiguous()#kp中x,y,z3个维度坐标的均值
        
        num_kp = self.num_key
        ver_Kp_1 = Kp[:, :, 1].view(1, num_kp, 1).contiguous()#【1，8，1】z坐标

        kk_1 = Kp[:, :, 0].view(1, num_kp, 1).contiguous()#【1，8，1】x坐标
        kk_2 = Kp[:, :, 2].view(1, num_kp, 1).contiguous()#【1，8，1】y坐标
        rad = torch.cat((kk_1, kk_2), dim=2).contiguous()#【1，8，2】x、y坐标
        ver_Kp_2 = torch.norm(rad, dim=2).view(1, num_kp, 1).contiguous()#【1，8，1】sqrt(x^2+y^2)

        tmp_aim_0 = torch.cat((Kp[:, 1:, :], Kp[:, 0:1, :]), dim=1).contiguous()#轮换，将第一个关键点挪到最后
        aim_0_x = tmp_aim_0[:, :, 0].view(-1).contiguous()#【8】
        aim_0_y = tmp_aim_0[:, :, 2].view(-1).contiguous()#【8】

        aim_1_x = Kp[:, :, 0].view(-1).contiguous()
        aim_1_y = Kp[:, :, 2].view(-1).contiguous()

        angle = torch.atan2(aim_1_y, aim_1_x) - torch.atan2(aim_0_y, aim_0_x)#【8】相邻关键点的夹角
        angle[angle < 0] += 2 * math.pi
        ver_Kp_3 = angle.view(1, num_kp, 1).contiguous() * 0.01#【1，8，1】

        ver_Kp = torch.cat((ver_Kp_1, ver_Kp_2, ver_Kp_3), dim=2).contiguous()#【1，8，3】--新坐标系，通道分别为高h、半径r、与前一个关键点的角angle

        return ver_Kp, cent0

2.3 旋转损失

############# Pose Error Loss
        rot_Kp_fr = (Kp_fr - t_fr).contiguous()#中间坐标系的关键点坐标--R^T乘世界坐标系
        rot_Kp_to = (Kp_to - t_to).contiguous()#【1，8，3】
        rot = torch.bmm(r_to, r_fr.transpose(2, 1))#ground truth 的deta R？【1，3，3】

        if sym_or_not:
            rot = torch.bmm(rot, self.sym_axis).view(-1)#rot与[[0],[1],[0]]乘，【1，3，1】即【3】原rot的第2列
            pred_r = self.estimate_rotation(rot_Kp_fr, rot_Kp_to, sym_or_not)
            loss_rot = (torch.acos(torch.sum(pred_r * rot) / (torch.norm(pred_r) * torch.norm(rot)))).contiguous()
            loss_rot = loss_rot
        else:
            pred_r = self.estimate_rotation(rot_Kp_fr, rot_Kp_to, sym_or_not)
            frob_sqr = torch.sum(((pred_r - rot) * (pred_r - rot)).view(-1)).contiguous()
            frob = torch.sqrt(frob_sqr).unsqueeze(0).contiguous()
            cc = torch.cat([self.oneone, frob / (2 * math.sqrt(2))]).contiguous()
            loss_rot = 2.0 * torch.mean(torch.asin(torch.min(cc))).contiguous()

该部分求旋转损失L_rot

其中ΔR^gt为ground truth R的变化，通过rot = torch.bmm(r_to, r_fr.transpose(2, 1))语句求出。
ΔR为代码中的pred_r，需要调用estimate_rotation()函数，该函数通过SVD求预测的旋转矩阵。
此外，注意到有旋转对称性的实例中，rot和pred_r都只取第二列的数据，是因为在(d,h, θ theta θ)坐标系下，旋转损失定义为：

---

estimate_rotation()

该函数利用SVD奇异值分解求旋转矩阵，输入pt0, pt1分别为2组顺序一致的点， sym_or_not标志是否为虚拟实例。
输出预测的旋转矩阵pred_r。

该方法的理论知识，参照SVD计算旋转，平移矩阵
文章中下图能清晰地反映此函数的实现逻辑：

代码实现：

def estimate_rotation(self, pt0, pt1, sym_or_not):
        #pt0、pt1为相邻两帧的关键点坐标（中间坐标系下）
        pconf2 = self.pconf.view(1, self.num_key, 1)#【1，8，1】均为0.125
        cent0 = torch.sum(pt0 * pconf2, dim=1).repeat(1, self.num_key, 1).contiguous()#【1，8，3】将各维度均值复制8遍
        cent1 = torch.sum(pt1 * pconf2, dim=1).repeat(1, self.num_key, 1).contiguous()

        diag_mat = torch.diag(self.pconf).unsqueeze(0)#【1，8，8】对角为0.125的矩阵
        x = (pt0 - cent0).transpose(2, 1).contiguous()#【1，3，8】前一帧关键点偏离各维度平均值的距离（中间坐标系下的坐标）
        y = pt1 - cent1#【1，8，3】后一帧关键点偏离各维度平均值的距离

        pred_t = cent1 - cent0#【1，8，3】根据相邻帧预测的位姿T复制8次

        cov = torch.bmm(torch.bmm(x, diag_mat), y).contiguous().squeeze(0)#【1，3，3】cov=x diag y

        u, _, v = torch.svd(cov)#奇异值分解，cov=u s v，其中u、v是正交矩阵，s是对角阵
        #u、v【3，3】，

        u = u.transpose(1, 0).contiguous()#【3，3】，u->u^T
        d = torch.det(torch.mm(v, u)).contiguous().view(1, 1, 1).contiguous()#|vu^T|1个数
        #torch.mm 2维矩阵乘，torch.bmm 3维矩阵乘（加了一维batchsize）
        u = u.transpose(1, 0).contiguous().unsqueeze(0)#【3，3】

        ud = torch.cat((u[:, :, :-1], u[:, :, -1:] * d), dim=2)#前两列不变，最后一列乘d
        v = v.transpose(1, 0).contiguous().unsqueeze(0)#v^T

        pred_r = torch.bmm(ud, v).transpose(2, 1).contiguous()#预测的R矩阵

        if sym_or_not:
            pred_r = torch.bmm(pred_r, self.sym_axis).contiguous().view(-1).contiguous()#只取第二列，旋转只在x-y平面？
            
        return pred_r

---

2.4 轮廓损失(表面损失)L_sil

############# Close To Surface Loss
        #轮廓损失loss_surf
        bs = 1
        num_p = 1
        num_point_mesh = self.num_key
        
        #mesh【1，3895，3】bbox网格点
        target_fr = mesh[0].transpose(1, 0).contiguous().view(3, -1)#【3，3895】
        pred_fr = gt_Kp_fr.permute(2, 0, 1).contiguous().view(3, -1)#【3，8】且将kp坐标调整为[y,x,z]顺序
        inds = self.knn(target_fr.unsqueeze(0), pred_fr.unsqueeze(0))#【1，1，8】
        #在mesh中找离每个关键点最近的一个点，inds是他们的idx
        target_fr = torch.index_select(target_fr, 1, inds.view(-1) - 1)#根据idx找这些点【3，8】
        target_fr = target_fr.view(3, bs * num_p, num_point_mesh).permute(1, 2, 0).contiguous()#【1，8，3】
        pred_fr = pred_fr.view(3, bs * num_p, num_point_mesh).permute(1, 2, 0).contiguous()#【1，8，3】

        loss_surf_fr = torch.mean(torch.norm((pred_fr - target_fr), dim=2), dim=1)
        #各关键点与离他们最近的邻点的距离取平均

        target_to = mesh[0].transpose(1, 0).contiguous().view(3, -1)
        pred_to = gt_Kp_to.permute(2, 0, 1).contiguous().view(3, -1)
        inds = self.knn(target_to.unsqueeze(0), pred_to.unsqueeze(0))
        target_to = torch.index_select(target_to, 1, inds.view(-1) - 1)
        target_to = target_to.view(3, bs * num_p, num_point_mesh).permute(1, 2, 0).contiguous()
        pred_to = pred_to.view(3, bs * num_p, num_point_mesh).permute(1, 2, 0).contiguous()

        loss_surf_to = torch.mean(torch.norm((pred_to - target_to), dim=2), dim=1)

        loss_surf = (loss_surf_fr + loss_surf_to).contiguous() / 2.0

该部分计算轮廓一致性的损失，迫使关键点更接近物体表面，以提高可解释性。
定义为8个Keypoints与他们最接近的bbox网格点的距离的平均值。
其中，self.knn函数在给定范围中找距离目标点最近的一个点，返回它的idx。

2.5 分离损失L_sep

############# Separate Loss
        #分离损失loss_sep
        scale = scale.view(-1)#【3】
        max_rad = torch.norm(scale).item()#归一化系数的欧氏长度

        gt_Kp_fr_select1 = torch.index_select(gt_Kp_fr, 1, self.select1).contiguous()
        #7组0，1，2....，7为idx对应的kp，【1，56，3】
        gt_Kp_fr_select2 = torch.index_select(gt_Kp_fr, 1, self.select2).contiguous()
        #7组1,2...7,0每组轮转对应的kp，【1，56，3】
        loss_sep_fr = torch.norm((gt_Kp_fr_select1 - gt_Kp_fr_select2), dim=2).view(-1).contiguous()
        #不同kp的间距【56】
        #希望间距大一些，则希望max_rad/8.0 - loss_sep_fr小，所以用max_rad/8.0 - loss_sep_fr做损失
        loss_sep_fr = torch.max(self.zeros, max_rad/8.0 - loss_sep_fr).contiguous()
        loss_sep_fr = torch.mean(loss_sep_fr).contiguous()#取均值，至此成为1个数

        gt_Kp_to_select1 = torch.index_select(gt_Kp_to, 1, self.select1).contiguous()
        gt_Kp_to_select2 = torch.index_select(gt_Kp_to, 1, self.select2).contiguous()
        loss_sep_to = torch.norm((gt_Kp_to_select1 - gt_Kp_to_select2), dim=2).view(-1).contiguous()
        loss_sep_to = torch.max(self.zeros, max_rad/8.0 - loss_sep_to).contiguous()
        loss_sep_to = torch.mean(loss_sep_to).contiguous()

        loss_sep = (loss_sep_fr + loss_sep_to) / 2.0

该部分计算分离损失L_sep，迫使关键点之间保持一定的距离，以避免构型退化和改善姿态估计。
其中loss_sep_fr = torch.norm((gt_Kp_fr_select1 - gt_Kp_fr_select2), dim=2).view(-1).contiguous()语句求出所有Keypoints间的距离，为改善估计效果，我们不希望所有的关键点都挤在一起，即希望Keypoints间的距离大一些。
由于损失函数在训练过程中是逐次减小的，故取max_rad/8.0 - loss_sep_fr的均值作为分离损失。

---

至此，6种损失函数全部得到，训练中实际使用的损失函数为这6种函数的加权和。

########### SUM UP
        #实际使用的loss是几种loss的加权和
        loss = loss_att * 4.0 + Kp_dis * 3.0 + Kp_cent_dis + loss_rot * 0.2 + loss_surf * 3.0 + loss_sep
        #loss是variable类型的，需要梯度
        score = (loss_att * 4.0 + Kp_dis * 3.0 + Kp_cent_dis + loss_rot * 0.2).item()
        #通过.item()取出variable中的元素，是一个数，这里是除轮廓损失、分离损失外的损失加权和
        print(cate.view(-1).item(), loss_att.item(), Kp_dis.item(), Kp_cent_dis.item(), loss_rot.item(), loss_surf.item(), loss_sep)

        return loss, score

三、eval中位姿预测 eval()

该函数输入初始帧和当前帧的关键点、当前帧的目标锚点，输出两帧间的相对位姿和预测误差，代码实现如下：

def ev(self, Kp_fr, Kp_to, att_to):
        #Kp_fr：初始帧一组关键点
        #Kp_to：当前帧的一组关键点
        #att_to：当前帧的目标锚点
        ori_Kp_fr = Kp_fr
        ori_Kp_to = Kp_to

        new_r, new_t = self.estimate_pose(Kp_fr, Kp_to)#根据两组关键点预测位姿变化

        Kp_to = torch.bmm((ori_Kp_to - new_t), new_r)#由当前帧的坐标和RT反推初始帧坐标

        Kp_dis = torch.mean(torch.norm((Kp_fr - Kp_to), dim=2), dim=1)#反推结果与真实值的差距

        new_t *= 1000.0
        return ori_Kp_fr, new_r.detach().cpu().numpy()[0], new_t.detach().cpu().numpy()[0], Kp_dis.item(), att_to

注意，此部分生成的new_t、new_r只是当前帧与初始帧之间的相对位姿，并不是实际的物体6D位姿（世界系与相机系见的变换），实际的物体6D位姿还需要在eval.py中与初始位姿结合得到。故Kp_to = torch.bmm((ori_Kp_to - new_t), new_r)得到的Kp_to是由相对位姿反推的初始帧坐标，仍然在相机系下，Kp_dis = torch.mean(torch.norm((Kp_fr - Kp_to), dim=2), dim=1)为反推的初始帧坐标与实际初始帧坐标的平均距离，作为位姿的预测误差。

estimate_pose()

ev()中预测两帧的相对位姿的过程实际在estimate_pose()中实现，代码如下：

def estimate_pose(self, pt0, pt1):#输入两组关键点坐标
        pconf2 = self.pconf.view(1, self.num_key, 1)#8个0.25，用于计算平均值
        cent0 = torch.sum(pt0 * pconf2, dim=1).repeat(1, self.num_key, 1).contiguous()#pt0均值--关键点质心坐标
        cent1 = torch.sum(pt1 * pconf2, dim=1).repeat(1, self.num_key, 1).contiguous()#py1均值

        diag_mat = torch.diag(self.pconf).unsqueeze(0)
        x = (pt0 - cent0).transpose(2, 1).contiguous()#pt0在中间坐标系下的坐标(与世界系同原点，与相机系同基)
        y = pt1 - cent1

        pred_t = cent1 - cent0#通过质心的移动预测T

        ####SVD奇异值分解求旋转矩阵
        cov = torch.bmm(torch.bmm(x, diag_mat), y).contiguous().squeeze(0)

        u, _, v = torch.svd(cov)

        u = u.transpose(1, 0).contiguous()
        d = torch.det(torch.mm(v, u)).contiguous().view(1, 1, 1).contiguous()
        u = u.transpose(1, 0).contiguous().unsqueeze(0)

        ud = torch.cat((u[:, :, :-1], u[:, :, -1:] * d), dim=2)
        v = v.transpose(1, 0).contiguous().unsqueeze(0)

        pred_r = torch.bmm(ud, v).transpose(2, 1).contiguous()
        return pred_r, pred_t[:, 0, :].view(1, 3)

其中平移向量T的预测通过两组关键点质心坐标的变化表示，而旋转矩阵则需要根据SVD奇异值分解的方法求得，详见SVD计算旋转，平移矩阵