Centerfusion训练过程中的损失函数初始化以及计算过程

损失函数的定义在src\lib\trainer.py23行左右,首先介绍损失函数的初始化部分
注：损失函数的调用在src\lib\trainer.py 134行

  def _get_losses(self, opt):
    loss_order = ['hm', 'wh', 'reg', 'ltrb', 'hps', 'hm_hp', \
      'hp_offset', 'dep', 'dep_sec', 'dim', 'rot', 'rot_sec',
      'amodel_offset', 'ltrb_amodal', 'tracking', 'nuscenes_att', 'velocity']
    #首先定义最后模型需要回归的头，之后根据这里的头来调用相应的损失函数进行回归
    loss_states = ['tot'] + [k for k in loss_order if k in opt.heads]
    loss = GenericLoss(opt) #这里就是损失函数的调用，都封装在了GenericLoss中
    return loss_states, loss

接下来详细介绍下GenericLoss这个类

class GenericLoss(torch.nn.Module):
  def __init__(self, opt):
    super(GenericLoss, self).__init__()
    self.crit = FastFocalLoss(opt=opt)        #FocalLoss用于之后的分类损失

这里只放了fastfocalloss类的初始化部分，之后在介绍函数具体运算的时候会放上所有的forward函数以及对应需要调用的函数，在该部分对于之后要用到的类也是主要介绍init函数。

class FastFocalLoss(nn.Module):
  '''
  Reimplemented focal loss, exactly the same as the CornerNet version.
  Faster and costs much less memory.
  '''
  def __init__(self, opt=None):
    super(FastFocalLoss, self).__init__()
    self.only_neg_loss = _only_neg_loss   #这里调用了负样本的损失函数 

论文中对于负样本的损失为$(1-Y_{xyz}{)}^{\beta }({\hat{Y}}_{xyz}{)}^{\alpha }\log (1-{\hat{Y}}_{xyz})$，其中$ Y_{xyz}$为groundtruth，$\widehat Y_{xyz}$为预测的标签。下面的函数就是根据这个公式进行编写的。

def _only_neg_loss(pred, gt):
  gt = torch.pow(1 - gt, 4)
  neg_loss = torch.log(1 - pred) * torch.pow(pred, 2) * gt
  return neg_loss.sum()

代码继续
下面几个损失大家可以在debug的时候详细看下，具体都是python包中的函数调用，就不一一介绍了

    self.crit_reg = RegWeightedL1Loss()         #用于计算hm的回归损失
     if 'rot' in opt.heads:
      self.crit_rot = BinRotLoss()              #用于计算alpha观测角的损失
    if 'nuscenes_att' in opt.heads:
      self.crit_nuscenes_att = WeightedBCELoss()#用于计算物体类别的损失
    self.opt = opt
    self.crit_dep = DepthLoss()                 #用于估计深度的损失

到这里损失函数的定义就完成了，接下来介绍损失函数的forward函数。
forward函数在src\lib\trainer.py46行左右

class GenericLoss(torch.nn.Module):
  def __init__(self, opt):
    pass

  def _sigmoid_output(self, output):
    if 'hm' in output:
      output['hm'] = _sigmoid(output['hm'])
    if 'hm_hp' in output:
      output['hm_hp'] = _sigmoid(output['hm_hp'])
    if 'dep' in output:
      output['dep'] = 1. / (output['dep'].sigmoid() + 1e-6) - 1.
    if 'dep_sec' in output and self.opt.sigmoid_dep_sec:
      output['dep_sec'] = 1. / (output['dep_sec'].sigmoid() + 1e-6) - 1.
    return output

  def forward(self, outputs, batch):
    opt = self.opt                           
    losses = {head: 0 for head in opt.heads}

    for s in range(opt.num_stacks):                
      output = outputs[s]
      output = self._sigmoid_output(output)  #将输出归一化到0-1之间

      if 'hm' in output:                     #利用focal对hm回归的分别计算损失
        losses['hm'] += self.crit(
          output['hm'], batch['hm'], batch['ind'], 
          batch['mask'], batch['cat']) / opt.num_stacks

接下介绍focalloss的forward函数

class FastFocalLoss(nn.Module):
  '''
  Reimplemented focal loss, exactly the same as the CornerNet version.
  Faster and costs much less memory.
  '''
  def __init__(self, opt=None):
    pass

  def forward(self, out, target, ind, mask, cat):
    '''
    Arguments:
      out, target: B x C x H x W   1*10*112*200
      ind, mask: B x M              
      cat (category id for peaks): B x M
    '''
    neg_loss = self.only_neg_loss(out, target)                       #计算负样本的损失
    pos_pred_pix = _tranpose_and_gather_feat(out, ind) # B x M x C

	def _tranpose_and_gather_feat(feat, ind):
	  feat = feat.permute(0, 2, 3, 1).contiguous()     #首先将feat从B*C*H*W转换成 BHWC
	  feat = feat.view(feat.size(0), -1, feat.size(3)) #之后在变成B*HW*C
	  '''
	  这里进行维度转换的原因是由于当时在dataset中设置中心点的时候，也是首先将图像展开成向量，之后保存中心点在向量中的位置
	  '''
	  feat = _gather_feat(feat, ind)
	  '''
	  _gather_feat函数如下边所示，ind中保存的就是原来2d框中心点在图像中的位置，
	  这里由于要使用gather函数获取到预测特征对应位置的值，所以将ind由原来的1*128大小扩展成1*128*10
	  加上原来1*128中的第一个值为7458，那么1*128*10中的第一行向量所有的值都为7458，这里是为了取出10类类别中所有的位置
	  用于之后的类别二次删选。
	  最后得到的特则为1*128*10
	  '''
	  return feat
	
	def _gather_feat(feat, ind):
	  dim = feat.size(2)
	  ind = ind.unsqueeze(2).expand(ind.size(0), ind.size(1), dim)
	  feat = feat.gather(1, ind)
	  return feat

    pos_pred = pos_pred_pix.gather(2, cat.unsqueeze(2)) # B x M    
    '''
    根据每个框的值进行二次删选,加入第一个框也就是上述7458位置对应的类别为第一类，
    那么只是取出了feat中的第一个样本的第一个元素做为预测的概率
    '''
    num_pos = mask.sum()
    pos_loss = torch.log(pos_pred) * torch.pow(1 - pos_pred, 2) * \
               mask.unsqueeze(2)
    #根据focalloss的正样本的公式进行计算损失
    pos_loss = pos_loss.sum() #改图像中对应的所有的bbox的个数
    if num_pos == 0:
      return - neg_loss
    return - (pos_loss + neg_loss) / num_pos #计算平均损失

代码继续

      if 'dep' in output:
        losses['dep'] += self.crit_dep(
          output['dep'], batch['dep'], batch['ind'], 
          batch['dep_mask'], batch['cat']) / opt.num_stacks

class DepthLoss(nn.Module):
  def __init__(self, opt=None):
    super(DepthLoss, self).__init__()

  def forward(self, output, target, ind, mask, cat):
    '''
    Arguments:
      out, target: B x C x H x W
      ind, mask: B x M
      cat (category id for peaks): B x M
    '''
    pred = _tranpose_and_gather_feat(output, ind) # B x M x (C) #获取框中心点对应的预测的深度值
    if pred.shape[2] > 1:
      pred = pred.gather(2, cat.unsqueeze(2)) # B x M
    loss = F.l1_loss(pred * mask, target * mask, reduction='sum') #利用l1损失计算depth的损失值
    loss = loss / (mask.sum() + 1e-4)
    return loss

      regression_heads = [
        'reg', 'wh', 'tracking', 'ltrb', 'ltrb_amodal', 'hps', 
        'dim', 'amodel_offset', 'velocity']

      for head in regression_heads:         #利用回归损失计算其他回归头的损失
        if head in output:
          losses[head] += self.crit_reg(
            output[head], batch[head + '_mask'],
            batch['ind'], batch[head]) / opt.num_stacks
      
      if 'hm_hp' in output:                              #没有执行
        losses['hm_hp'] += self.crit(
          output['hm_hp'], batch['hm_hp'], batch['hp_ind'], 
          batch['hm_hp_mask'], batch['joint']) / opt.num_stacks
        if 'hp_offset' in output:                       #没有执行
          losses['hp_offset'] += self.crit_reg(
            output['hp_offset'], batch['hp_offset_mask'],
            batch['hp_ind'], batch['hp_offset']) / opt.num_stacks
        
      if 'rot' in output:
        losses['rot'] += self.crit_rot(                                    #用于回归alpha观测角
          output['rot'], batch['rot_mask'], batch['ind'], batch['rotbin'],
          batch['rotres']) / opt.num_stacks
         

class BinRotLoss(nn.Module):
  def __init__(self):
    super(BinRotLoss, self).__init__()
  
  def forward(self, output, mask, ind, rotbin, rotres):
    pred = _tranpose_and_gather_feat(output, ind)          #获取对应位置的alpha观测角相关预测值
    loss = compute_rot_loss(pred, rotbin, rotres, mask)
    return loss

def compute_rot_loss(output, target_bin, target_res, mask):
    # output: (B, 128, 8) [bin1_cls[0], bin1_cls[1], bin1_sin, bin1_cos, 
    #                 bin2_cls[0], bin2_cls[1], bin2_sin, bin2_cos]
    # target_bin: (B, 128, 2) [bin1_cls, bin2_cls]
    # target_res: (B, 128, 2) [bin1_res, bin2_res]
    # mask: (B, 128, 1)
    '''
    这里用于计算alpha观测角的损失，作者在计算该部分的时候将分成区域+对应区域的残差角进行预测
    首先见360°分成两个相交的空间(bin)，之后计算预测角度在该区域内的残差角
    如果残差角在相交的区域，那么bin1，bin2对应的值都为1，如果只在一个区域内，则对应位置的值为1，另一个为0
    在预测残差角的时候用的是sin 和cos两个进行预测，
    '''
    output = output.view(-1, 8)
    target_bin = target_bin.view(-1, 2)
    target_res = target_res.view(-1, 2)
    mask = mask.view(-1, 1)
    #前四维的表示bin1的相关预测值，前两个表示在那个bin中，后两个是对应的sin cos值，后四维表示bin2
    #这里是计算bin对应的损失，由于是分类损失用的是交叉熵损失
    loss_bin1 = compute_bin_loss(output[:, 0:2], target_bin[:, 0], mask)        
    loss_bin2 = compute_bin_loss(output[:, 4:6], target_bin[:, 1], mask)
    
    #这里计算残差角的损失
    #如果groundtruth对应bin位置有值的话，则计算对应的sin，cos损失
    loss_res = torch.zeros_like(loss_bin1)
    if target_bin[:, 0].nonzero().shape[0] > 0:
        idx1 = target_bin[:, 0].nonzero()[:, 0]
        valid_output1 = torch.index_select(output, 0, idx1.long())
        valid_target_res1 = torch.index_select(target_res, 0, idx1.long())
        loss_sin1 = compute_res_loss(
          valid_output1[:, 2], torch.sin(valid_target_res1[:, 0]))
        loss_cos1 = compute_res_loss(
          valid_output1[:, 3], torch.cos(valid_target_res1[:, 0]))
        loss_res += loss_sin1 + loss_cos1
    if target_bin[:, 1].nonzero().shape[0] > 0:
        idx2 = target_bin[:, 1].nonzero()[:, 0]
        valid_output2 = torch.index_select(output, 0, idx2.long())
        valid_target_res2 = torch.index_select(target_res, 0, idx2.long())
        loss_sin2 = compute_res_loss(
          valid_output2[:, 6], torch.sin(valid_target_res2[:, 1]))
        loss_cos2 = compute_res_loss(
          valid_output2[:, 7], torch.cos(valid_target_res2[:, 1]))
        loss_res += loss_sin2 + loss_cos2
    return loss_bin1 + loss_bin2 + loss_res

      if 'nuscenes_att' in output:                            #计算状态损失
        losses['nuscenes_att'] += self.crit_nuscenes_att(
          output['nuscenes_att'], batch['nuscenes_att_mask'],
          batch['ind'], batch['nuscenes_att']) / opt.num_stacks
      
      if 'dep_sec' in output:                                  #计算第二个回归头的深度损失
        losses['dep_sec'] += self.crit_dep(
          output['dep_sec'], batch['dep'], batch['ind'], 
          batch['dep_mask'], batch['cat']) / opt.num_stacks
      
      if 'rot_sec' in output:                                 #计算第二个回归头的角度损失
        losses['rot_sec'] += self.crit_rot(
          output['rot_sec'], batch['rot_mask'], batch['ind'], batch['rotbin'],
          batch['rotres']) / opt.num_stacks

    losses['tot'] = 0
    for head in opt.heads:                                  #计算总的损失
      losses['tot'] += opt.weights[head] * losses[head] 

    return losses['tot'], losses