[On-device Training] Yolov4 custom loss

[Marouan-st]

Hello,

I would like to implement a custom loss to be able to train on-device a yolov4-tiny model for object detection.

To compute the loss some post-processing must be performed on the output of the model, like computing bboxes iou and sum several losses (class loss + confidence loss + iou loss: cross entropy losses): see https://www.nature.com/articles/s41598-021-02225-y/figures/3

I don't see how to implement all these needed computations in the custom loss, especially how to provide the different losses with the post-processed input, since onnx loss functions takes as input String arguments (input name).

I'm using a yolov4-tiny model compiled from darknet and converted to onnx from a tensorflow implementation of the model.

The Torch implementation of this loss function (for the model i'm using) would look like this (inspired by this yolov4 loss tensorflow implementation):

def compute_loss(pred, conv, label, bboxes, STRIDES=[16, 32], NUM_CLASS=1, IOU_LOSS_THRESH=0.5, i=0):
      conv_shape  = conv.size()
      batch_size  = conv_shape[0]
      output_size = conv_shape[1]
      input_size  = STRIDES[i] * output_size
      conv = torch.reshape(conv, (batch_size, output_size, output_size, 3, 5 + NUM_CLASS))

      conv_raw_conf = conv[:, :, :, :, 4:5]
      conv_raw_prob = conv[:, :, :, :, 5:]

      pred_xywh     = pred[:, :, :, :, 0:4]
      pred_conf     = pred[:, :, :, :, 4:5]

      label_xywh    = label[:, :, :, :, 0:4]
      respond_bbox  = label[:, :, :, :, 4:5]
      label_prob    = label[:, :, :, :, 5:]

      giou = torch.unsqueeze(bbox_giou(pred_xywh, label_xywh), 0) # Here not sure...
      input_size = input_size.to(torch.float32)

      bbox_loss_scale = 2.0 - 1.0 * label_xywh[:, :, :, :, 2:3] * label_xywh[:, :, :, :, 3:4] / (input_size ** 2)
      giou_loss = respond_bbox * bbox_loss_scale * (1- giou)

      iou = bbox_iou(pred_xywh[:, :, :, :, np.newaxis, :], bboxes[:, np.newaxis, np.newaxis, np.newaxis, :, :])
      max_iou = torch.unsqueeze(torch.max(iou), 0)

      respond_bgd = (1.0 - respond_bbox) * (max_iou < IOU_LOSS_THRESH).to(torch.float32)

      conf_focal = torch.pow(respond_bbox - pred_conf, 2)

      conf_loss = conf_focal * (
              respond_bbox * torch.nn.functional.cross_entropy(input=conv_raw_conf, target=respond_bbox)
              +
              respond_bgd * torch.nn.functional.cross_entropy(input=conv_raw_conf, target=respond_bbox)
      )

      prob_loss = respond_bbox * torch.nn.functional.cross_entropy(input=conv_raw_prob, target=label_prob)
      giou_loss = torch.mean(torch.sum(giou_loss))
      conf_loss = torch.mean(torch.sum(conf_loss, axis=[1,2,3,4]))
      prob_loss = torch.mean(torch.sum(prob_loss, axis=[1,2,3,4]))

      return giou_loss + conf_loss + prob_loss

def bbox_iou(bboxes1, bboxes2):
      """
      @param bboxes1: (a, b, ..., 4)
      @param bboxes2: (A, B, ..., 4)
          x:X is 1:n or n:n or n:1
      @return (max(a,A), max(b,B), ...)
      ex) (4,):(3,4) -> (3,)
          (2,1,4):(2,3,4) -> (2,3)
      """
      bboxes1_area = bboxes1[..., 2] * bboxes1[..., 3]
      bboxes2_area = bboxes2[..., 2] * bboxes2[..., 3]

      bboxes1_coor = torch.concat(
          [
              bboxes1[..., :2] - bboxes1[..., 2:] * 0.5,
              bboxes1[..., :2] + bboxes1[..., 2:] * 0.5,
          ],
          axis=-1,
      )
      bboxes2_coor = torch.concat(
          [
              bboxes2[..., :2] - bboxes2[..., 2:] * 0.5,
              bboxes2[..., :2] + bboxes2[..., 2:] * 0.5,
          ],
          axis=-1,
      )

      left_up = torch.maximum(bboxes1_coor[..., :2], bboxes2_coor[..., :2])
      right_down = torch.minimum(bboxes1_coor[..., 2:], bboxes2_coor[..., 2:])

      inter_section = torch.maximum(right_down - left_up, 0.0)
      inter_area = inter_section[..., 0] * inter_section[..., 1]

      union_area = bboxes1_area + bboxes2_area - inter_area

      iou = torch.div(inter_area, union_area)

      return iou

def bbox_giou(bboxes1, bboxes2):
      """
      Generalized IoU
      @param bboxes1: (a, b, ..., 4)
      @param bboxes2: (A, B, ..., 4)
          x:X is 1:n or n:n or n:1
      @return (max(a,A), max(b,B), ...)
      ex) (4,):(3,4) -> (3,)
          (2,1,4):(2,3,4) -> (2,3)
      """
      bboxes1_area = bboxes1[..., 2] * bboxes1[..., 3]
      bboxes2_area = bboxes2[..., 2] * bboxes2[..., 3]
  
      bboxes1_coor = torch.concat(
          [
              bboxes1[..., :2] - bboxes1[..., 2:] * 0.5,
              bboxes1[..., :2] + bboxes1[..., 2:] * 0.5,
          ],
          axis=-1,
      )
      bboxes2_coor = torch.concat(
          [
              bboxes2[..., :2] - bboxes2[..., 2:] * 0.5,
              bboxes2[..., :2] + bboxes2[..., 2:] * 0.5,
          ],
          axis=-1,
      )
  
      left_up = torch.maximum(bboxes1_coor[..., :2], bboxes2_coor[..., :2])
      right_down = torch.minimum(bboxes1_coor[..., 2:], bboxes2_coor[..., 2:])
  
      inter_section = torch.maximum(right_down - left_up, 0.0)
      inter_area = inter_section[..., 0] * inter_section[..., 1]
  
      union_area = bboxes1_area + bboxes2_area - inter_area
  
      iou = torch.div(inter_area, union_area)
  
      enclose_left_up = torch.minimum(bboxes1_coor[..., :2], bboxes2_coor[..., :2])
      enclose_right_down = torch.maximum(
          bboxes1_coor[..., 2:], bboxes2_coor[..., 2:]
      )
  
      enclose_section = enclose_right_down - enclose_left_up
      enclose_area = enclose_section[..., 0] * enclose_section[..., 1]
  
      giou = iou - torch.div(enclose_area - union_area, enclose_area)
  
      return giou

Any suggestions?

Thank you

[baijumeswani]

Here is an idea for you to consider: Try to make the exported onnx model contain the loss and then when building the training artifacts, avoid passing in the loss argument:

class MyPTModelWithLoss:
    def __init__(self):
         ...

    def forward(self, ...):
        p, q, r = compute_logits()
        loss = loss1(p) + loss2(q) + loss3(r)
        return loss

pt_model = MyPTModelWithLoss(...)
torch.onnx.export(pt_model, ...)

onnx_model = onnx.load(<exported_onnx_model_path>)
artifacts.generate_artifacts(onnx_model, requires_grad=[...], frozen_params=[...], loss=None, optimizer=...)

Let me know how that goes.