code: #!/usr/bin/env python3
  # -*- coding:utf-8 -*-
  # Copyright (c) Megvii Inc. All rights reserved.
  
  import math
  from loguru import logger
  
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  
  from yolox.utils import bboxes_iou, meshgrid
  
  from .losses import IOUloss
  from .network_blocks import BaseConv, DWConv
  
  
  class YOLOXHead(nn.Module):
      def __init__(
          self,
          num_classes,
          width=1.0,
          strides=[8, 16, 32],
          in_channels=[256, 512, 1024],
          act="silu",
          depthwise=False,
      ):
          """
          Args:
              act (str): activation type of conv. Defalut value: "silu".
              depthwise (bool): whether apply depthwise conv in conv branch. Defalut value: False.
          """
          super().__init__()
  
          self.n_anchors = 1
          self.num_classes = num_classes
          self.decode_in_inference = True  # for deploy, set to False
  
          self.cls_convs = nn.ModuleList()
          self.reg_convs = nn.ModuleList()
          self.cls_preds = nn.ModuleList()
          self.reg_preds = nn.ModuleList()
          self.obj_preds = nn.ModuleList()
          self.stems = nn.ModuleList()
          Conv = DWConv if depthwise else BaseConv
  
          for i in range(len(in_channels)):
              self.stems.append(
                  BaseConv(
                      in_channels=int(in_channels[i] * width),
                      out_channels=int(256 * width),
                      ksize=1,
                      stride=1,
                      act=act,
                  )
              )
              self.cls_convs.append(
                  nn.Sequential(
                      *[
                          Conv(
                              in_channels=int(256 * width),
                              out_channels=int(256 * width),
                              ksize=3,
                              stride=1,
                              act=act,
                          ),
                          Conv(
                              in_channels=int(256 * width),
                              out_channels=int(256 * width),
                              ksize=3,
                              stride=1,
                              act=act,
                          ),
                      ]
                  )
              )
              self.reg_convs.append(
                  nn.Sequential(
                      *[
                          Conv(
                              in_channels=int(256 * width),
                              out_channels=int(256 * width),
                              ksize=3,
                              stride=1,
                              act=act,
                          ),
                          Conv(
                              in_channels=int(256 * width),
                              out_channels=int(256 * width),
                              ksize=3,
                              stride=1,
                              act=act,
                          ),
                      ]
                  )
              )
              self.cls_preds.append(
                  nn.Conv2d(
                      in_channels=int(256 * width),
                      out_channels=self.n_anchors * self.num_classes,
                      kernel_size=1,
                      stride=1,
                      padding=0,
                  )
              )
              self.reg_preds.append(
                  nn.Conv2d(
                      in_channels=int(256 * width),
                      out_channels=4,
                      kernel_size=1,
                      stride=1,
                      padding=0,
                  )
              )
              self.obj_preds.append(
                  nn.Conv2d(
                      in_channels=int(256 * width),
                      out_channels=self.n_anchors * 1,
                      kernel_size=1,
                      stride=1,
                      padding=0,
                  )
              )
  
          self.use_l1 = False
          self.l1_loss = nn.L1Loss(reduction="none")
          self.bcewithlog_loss = nn.BCEWithLogitsLoss(reduction="none")
          self.iou_loss = IOUloss(reduction="none")
          self.strides = strides
          self.grids = [torch.zeros(1)] * len(in_channels)
  
      def initialize_biases(self, prior_prob):
          for conv in self.cls_preds:
              b = conv.bias.view(self.n_anchors, -1)
              b.data.fill_(-math.log((1 - prior_prob) / prior_prob))
              conv.bias = torch.nn.Parameter(b.view(-1), requires_grad=True)
  
          for conv in self.obj_preds:
              b = conv.bias.view(self.n_anchors, -1)
              b.data.fill_(-math.log((1 - prior_prob) / prior_prob))
              conv.bias = torch.nn.Parameter(b.view(-1), requires_grad=True)
  
      def forward(self, xin, labels=None, imgs=None):
          outputs = []
          origin_preds = []
          x_shifts = []
          y_shifts = []
          expanded_strides = []
  
          for k, (cls_conv, reg_conv, stride_this_level, x) in enumerate(
              zip(self.cls_convs, self.reg_convs, self.strides, xin)
          ):
              x = self.stems[k](x)
              cls_x = x
              reg_x = x
  
              cls_feat = cls_conv(cls_x)
              cls_output = self.cls_preds[k](cls_feat)
  
              reg_feat = reg_conv(reg_x)
              reg_output = self.reg_preds[k](reg_feat)
              obj_output = self.obj_preds[k](reg_feat)
  
              if self.training:
                  output = torch.cat([reg_output, obj_output, cls_output], 1)
                  output, grid = self.get_output_and_grid(
                      output, k, stride_this_level, xin[0].type()
                  )
                  x_shifts.append(grid[:, :, 0])
                  y_shifts.append(grid[:, :, 1])
                  expanded_strides.append(
                      torch.zeros(1, grid.shape[1])
                      .fill_(stride_this_level)
                      .type_as(xin[0])
                  )
                  if self.use_l1:
                      batch_size = reg_output.shape[0]
                      hsize, wsize = reg_output.shape[-2:]
                      reg_output = reg_output.view(
                          batch_size, self.n_anchors, 4, hsize, wsize
                      )
                      reg_output = reg_output.permute(0, 1, 3, 4, 2).reshape(
                          batch_size, -1, 4
                      )
                      origin_preds.append(reg_output.clone())
  
              else:
                  output = torch.cat(
                      [reg_output, obj_output.sigmoid(), cls_output.sigmoid()], 1
                  )
  
              outputs.append(output)
  
          if self.training:
              return self.get_losses(
                  imgs,
                  x_shifts,
                  y_shifts,
                  expanded_strides,
                  labels,
                  torch.cat(outputs, 1),
                  origin_preds,
                  dtype=xin[0].dtype,
              )
          else:
              self.hw = [x.shape[-2:] for x in outputs]
              # [batch, n_anchors_all, 85]
              outputs = torch.cat(
                  [x.flatten(start_dim=2) for x in outputs], dim=2
              ).permute(0, 2, 1)
              if self.decode_in_inference:
                  return self.decode_outputs(outputs, dtype=xin[0].type())
              else:
                  return outputs
  
      def get_output_and_grid(self, output, k, stride, dtype):
          grid = self.grids[k]
  
          batch_size = output.shape[0]
          n_ch = 5 + self.num_classes
          hsize, wsize = output.shape[-2:]
          if grid.shape[2:4] != output.shape[2:4]:
              yv, xv = meshgrid([torch.arange(hsize), torch.arange(wsize)])
              grid = torch.stack((xv, yv), 2).view(1, 1, hsize, wsize, 2).type(dtype)
              self.grids[k] = grid
  
          output = output.view(batch_size, self.n_anchors, n_ch, hsize, wsize)
          output = output.permute(0, 1, 3, 4, 2).reshape(
              batch_size, self.n_anchors * hsize * wsize, -1
          )
          grid = grid.view(1, -1, 2)
          output[..., :2] = (output[..., :2] + grid) * stride
          output[..., 2:4] = torch.exp(output[..., 2:4]) * stride
          return output, grid
  
      def decode_outputs(self, outputs, dtype):
          grids = []
          strides = []
          for (hsize, wsize), stride in zip(self.hw, self.strides):
              yv, xv = meshgrid([torch.arange(hsize), torch.arange(wsize)])
              grid = torch.stack((xv, yv), 2).view(1, -1, 2)
              grids.append(grid)
              shape = grid.shape[:2]
              strides.append(torch.full((*shape, 1), stride))
  
          grids = torch.cat(grids, dim=1).type(dtype)
          strides = torch.cat(strides, dim=1).type(dtype)
  
          outputs[..., :2] = (outputs[..., :2] + grids) * strides
          outputs[..., 2:4] = torch.exp(outputs[..., 2:4]) * strides
          return outputs
  
      def get_losses(
          self,
          imgs,
          x_shifts,
          y_shifts,
          expanded_strides,
          labels,
          outputs,
          origin_preds,
          dtype,
      ):
          bbox_preds = outputs[:, :, :4]  # [batch, n_anchors_all, 4]
          obj_preds = outputs[:, :, 4].unsqueeze(-1)  # [batch, n_anchors_all, 1]
          cls_preds = outputs[:, :, 5:]  # [batch, n_anchors_all, n_cls]
  
          # calculate targets
          nlabel = (labels.sum(dim=2) > 0).sum(dim=1)  # number of objects
  
          total_num_anchors = outputs.shape[1]
          x_shifts = torch.cat(x_shifts, 1)  # [1, n_anchors_all]
          y_shifts = torch.cat(y_shifts, 1)  # [1, n_anchors_all]
          expanded_strides = torch.cat(expanded_strides, 1)
          if self.use_l1:
              origin_preds = torch.cat(origin_preds, 1)
  
          cls_targets = []
          reg_targets = []
          l1_targets = []
          obj_targets = []
          fg_masks = []
  
          num_fg = 0.0
          num_gts = 0.0
  
          for batch_idx in range(outputs.shape[0]):
              num_gt = int(nlabel[batch_idx])
              num_gts += num_gt
              if num_gt == 0:
                  cls_target = outputs.new_zeros((0, self.num_classes))
                  reg_target = outputs.new_zeros((0, 4))
                  l1_target = outputs.new_zeros((0, 4))
                  obj_target = outputs.new_zeros((total_num_anchors, 1))
                  fg_mask = outputs.new_zeros(total_num_anchors).bool()
              else:
                  gt_bboxes_per_image = labels[batch_idx, :num_gt, 1:5]
                  gt_classes = labels[batch_idx, :num_gt, 0]
                  bboxes_preds_per_image = bbox_preds[batch_idx]
  
                  try:
                      (
                          gt_matched_classes,
                          fg_mask,
                          pred_ious_this_matching,
                          matched_gt_inds,
                          num_fg_img,
                      ) = self.get_assignments(  # noqa
                          batch_idx,
                          num_gt,
                          total_num_anchors,
                          gt_bboxes_per_image,
                          gt_classes,
                          bboxes_preds_per_image,
                          expanded_strides,
                          x_shifts,
                          y_shifts,
                          cls_preds,
                          bbox_preds,
                          obj_preds,
                          labels,
                          imgs,
                      )
                  except RuntimeError as e:
                      # TODO: the string might change, consider a better way
                      if "CUDA out of memory. " not in str(e):
                          raise  # RuntimeError might not caused by CUDA OOM
  
                      logger.error(
                          "OOM RuntimeError is raised due to the huge memory cost during label assignment. \
                             CPU mode is applied in this batch. If you want to avoid this issue, \
                             try to reduce the batch size or image size."
                      )
                      torch.cuda.empty_cache()
                      (
                          gt_matched_classes,
                          fg_mask,
                          pred_ious_this_matching,
                          matched_gt_inds,
                          num_fg_img,
                      ) = self.get_assignments(  # noqa
                          batch_idx,
                          num_gt,
                          total_num_anchors,
                          gt_bboxes_per_image,
                          gt_classes,
                          bboxes_preds_per_image,
                          expanded_strides,
                          x_shifts,
                          y_shifts,
                          cls_preds,
                          bbox_preds,
                          obj_preds,
                          labels,
                          imgs,
                          "cpu",
                      )
  
                  torch.cuda.empty_cache()
                  num_fg += num_fg_img
  
                  cls_target = F.one_hot(
                      gt_matched_classes.to(torch.int64), self.num_classes
                  ) * pred_ious_this_matching.unsqueeze(-1)
                  obj_target = fg_mask.unsqueeze(-1)
                  reg_target = gt_bboxes_per_image[matched_gt_inds]
                  if self.use_l1:
                      l1_target = self.get_l1_target(
                          outputs.new_zeros((num_fg_img, 4)),
                          gt_bboxes_per_image[matched_gt_inds],
                          expanded_strides[0][fg_mask],
                          x_shifts=x_shifts[0][fg_mask],
                          y_shifts=y_shifts[0][fg_mask],
                      )
  
              cls_targets.append(cls_target)
              reg_targets.append(reg_target)
              obj_targets.append(obj_target.to(dtype))
              fg_masks.append(fg_mask)
              if self.use_l1:
                  l1_targets.append(l1_target)
  
          cls_targets = torch.cat(cls_targets, 0)
          reg_targets = torch.cat(reg_targets, 0)
          obj_targets = torch.cat(obj_targets, 0)
          fg_masks = torch.cat(fg_masks, 0)
          if self.use_l1:
              l1_targets = torch.cat(l1_targets, 0)
  
          num_fg = max(num_fg, 1)
          loss_iou = (
              self.iou_loss(bbox_preds.view(-1, 4)[fg_masks], reg_targets)
          ).sum() / num_fg
          loss_obj = (
              self.bcewithlog_loss(obj_preds.view(-1, 1), obj_targets)
          ).sum() / num_fg
          loss_cls = (
              self.bcewithlog_loss(
                  cls_preds.view(-1, self.num_classes)[fg_masks], cls_targets
              )
          ).sum() / num_fg
          if self.use_l1:
              loss_l1 = (
                  self.l1_loss(origin_preds.view(-1, 4)[fg_masks], l1_targets)
              ).sum() / num_fg
          else:
              loss_l1 = 0.0
  
          reg_weight = 5.0
          loss = reg_weight * loss_iou + loss_obj + loss_cls + loss_l1
  
          return (
              loss,
              reg_weight * loss_iou,
              loss_obj,
              loss_cls,
              loss_l1,
              num_fg / max(num_gts, 1),
          )
  
      def get_l1_target(self, l1_target, gt, stride, x_shifts, y_shifts, eps=1e-8):
          l1_target[:, 0] = gt[:, 0] / stride - x_shifts
          l1_target[:, 1] = gt[:, 1] / stride - y_shifts
          l1_target[:, 2] = torch.log(gt[:, 2] / stride + eps)
          l1_target[:, 3] = torch.log(gt[:, 3] / stride + eps)
          return l1_target
  
      @torch.no_grad()
      def get_assignments(
          self,
          batch_idx,
          num_gt,
          total_num_anchors,
          gt_bboxes_per_image,
          gt_classes,
          bboxes_preds_per_image,
          expanded_strides,
          x_shifts,
          y_shifts,
          cls_preds,
          bbox_preds,
          obj_preds,
          labels,
          imgs,
          mode="gpu",
      ):
  
          if mode == "cpu":
              print("------------CPU Mode for This Batch-------------")
              gt_bboxes_per_image = gt_bboxes_per_image.cpu().float()
              bboxes_preds_per_image = bboxes_preds_per_image.cpu().float()
              gt_classes = gt_classes.cpu().float()
              expanded_strides = expanded_strides.cpu().float()
              x_shifts = x_shifts.cpu()
              y_shifts = y_shifts.cpu()
  
          fg_mask, is_in_boxes_and_center = self.get_in_boxes_info(
              gt_bboxes_per_image,
              expanded_strides,
              x_shifts,
              y_shifts,
              total_num_anchors,
              num_gt,
          )
  
          bboxes_preds_per_image = bboxes_preds_per_image[fg_mask]
          cls_preds_ = cls_preds[batch_idx][fg_mask]
          obj_preds_ = obj_preds[batch_idx][fg_mask]
          num_in_boxes_anchor = bboxes_preds_per_image.shape[0]
  
          if mode == "cpu":
              gt_bboxes_per_image = gt_bboxes_per_image.cpu()
              bboxes_preds_per_image = bboxes_preds_per_image.cpu()
  
          pair_wise_ious = bboxes_iou(gt_bboxes_per_image, bboxes_preds_per_image, False)
  
          gt_cls_per_image = (
              F.one_hot(gt_classes.to(torch.int64), self.num_classes)
              .float()
              .unsqueeze(1)
              .repeat(1, num_in_boxes_anchor, 1)
          )
          pair_wise_ious_loss = -torch.log(pair_wise_ious + 1e-8)
  
          if mode == "cpu":
              cls_preds_, obj_preds_ = cls_preds_.cpu(), obj_preds_.cpu()
  
          with torch.cuda.amp.autocast(enabled=False):
              cls_preds_ = (
                  cls_preds_.float().unsqueeze(0).repeat(num_gt, 1, 1).sigmoid_()
                  * obj_preds_.float().unsqueeze(0).repeat(num_gt, 1, 1).sigmoid_()
              )
              pair_wise_cls_loss = F.binary_cross_entropy(
                  cls_preds_.sqrt_(), gt_cls_per_image, reduction="none"
              ).sum(-1)
          del cls_preds_
  
          cost = (
              pair_wise_cls_loss
              + 3.0 * pair_wise_ious_loss
              + 100000.0 * (~is_in_boxes_and_center)
          )
  
          (
              num_fg,
              gt_matched_classes,
              pred_ious_this_matching,
              matched_gt_inds,
          ) = self.dynamic_k_matching(cost, pair_wise_ious, gt_classes, num_gt, fg_mask)
          del pair_wise_cls_loss, cost, pair_wise_ious, pair_wise_ious_loss
  
          if mode == "cpu":
              gt_matched_classes = gt_matched_classes.cpu()
              fg_mask = fg_mask.cpu()
              pred_ious_this_matching = pred_ious_this_matching.cpu()
              matched_gt_inds = matched_gt_inds.cpu()
  
          return (
              gt_matched_classes,
              fg_mask,
              pred_ious_this_matching,
              matched_gt_inds,
              num_fg,
          )
  
      def get_in_boxes_info(
          self,
          gt_bboxes_per_image,
          expanded_strides,
          x_shifts,
          y_shifts,
          total_num_anchors,
          num_gt,
      ):
          expanded_strides_per_image = expanded_strides[0]
          x_shifts_per_image = x_shifts[0] * expanded_strides_per_image
          y_shifts_per_image = y_shifts[0] * expanded_strides_per_image
          x_centers_per_image = (
              (x_shifts_per_image + 0.5 * expanded_strides_per_image)
              .unsqueeze(0)
              .repeat(num_gt, 1)
          )  # [n_anchor] -> [n_gt, n_anchor]
          y_centers_per_image = (
              (y_shifts_per_image + 0.5 * expanded_strides_per_image)
              .unsqueeze(0)
              .repeat(num_gt, 1)
          )
  
          gt_bboxes_per_image_l = (
              (gt_bboxes_per_image[:, 0] - 0.5 * gt_bboxes_per_image[:, 2])
              .unsqueeze(1)
              .repeat(1, total_num_anchors)
          )
          gt_bboxes_per_image_r = (
              (gt_bboxes_per_image[:, 0] + 0.5 * gt_bboxes_per_image[:, 2])
              .unsqueeze(1)
              .repeat(1, total_num_anchors)
          )
          gt_bboxes_per_image_t = (
              (gt_bboxes_per_image[:, 1] - 0.5 * gt_bboxes_per_image[:, 3])
              .unsqueeze(1)
              .repeat(1, total_num_anchors)
          )
          gt_bboxes_per_image_b = (
              (gt_bboxes_per_image[:, 1] + 0.5 * gt_bboxes_per_image[:, 3])
              .unsqueeze(1)
              .repeat(1, total_num_anchors)
          )
  
          b_l = x_centers_per_image - gt_bboxes_per_image_l
          b_r = gt_bboxes_per_image_r - x_centers_per_image
          b_t = y_centers_per_image - gt_bboxes_per_image_t
          b_b = gt_bboxes_per_image_b - y_centers_per_image
          bbox_deltas = torch.stack([b_l, b_t, b_r, b_b], 2)
  
          is_in_boxes = bbox_deltas.min(dim=-1).values > 0.0
          is_in_boxes_all = is_in_boxes.sum(dim=0) > 0
          # in fixed center
  
          center_radius = 2.5
  
          gt_bboxes_per_image_l = (gt_bboxes_per_image[:, 0]).unsqueeze(1).repeat(
              1, total_num_anchors
          ) - center_radius * expanded_strides_per_image.unsqueeze(0)
          gt_bboxes_per_image_r = (gt_bboxes_per_image[:, 0]).unsqueeze(1).repeat(
              1, total_num_anchors
          ) + center_radius * expanded_strides_per_image.unsqueeze(0)
          gt_bboxes_per_image_t = (gt_bboxes_per_image[:, 1]).unsqueeze(1).repeat(
              1, total_num_anchors
          ) - center_radius * expanded_strides_per_image.unsqueeze(0)
          gt_bboxes_per_image_b = (gt_bboxes_per_image[:, 1]).unsqueeze(1).repeat(
              1, total_num_anchors
          ) + center_radius * expanded_strides_per_image.unsqueeze(0)
  
          c_l = x_centers_per_image - gt_bboxes_per_image_l
          c_r = gt_bboxes_per_image_r - x_centers_per_image
          c_t = y_centers_per_image - gt_bboxes_per_image_t
          c_b = gt_bboxes_per_image_b - y_centers_per_image
          center_deltas = torch.stack([c_l, c_t, c_r, c_b], 2)
          is_in_centers = center_deltas.min(dim=-1).values > 0.0
          is_in_centers_all = is_in_centers.sum(dim=0) > 0
  
          # in boxes and in centers
          is_in_boxes_anchor = is_in_boxes_all | is_in_centers_all
  
          is_in_boxes_and_center = (
              is_in_boxes[:, is_in_boxes_anchor] & is_in_centers[:, is_in_boxes_anchor]
          )
          return is_in_boxes_anchor, is_in_boxes_and_center
  
      def dynamic_k_matching(self, cost, pair_wise_ious, gt_classes, num_gt, fg_mask):
          # Dynamic K
          # ---------------------------------------------------------------
          matching_matrix = torch.zeros_like(cost, dtype=torch.uint8)
  
          ious_in_boxes_matrix = pair_wise_ious
          n_candidate_k = min(10, ious_in_boxes_matrix.size(1))
          topk_ious, _ = torch.topk(ious_in_boxes_matrix, n_candidate_k, dim=1)
          dynamic_ks = torch.clamp(topk_ious.sum(1).int(), min=1)
          dynamic_ks = dynamic_ks.tolist()
          for gt_idx in range(num_gt):
              _, pos_idx = torch.topk(
                  cost[gt_idx], k=dynamic_ks[gt_idx], largest=False
              )
              matching_matrix[gt_idx][pos_idx] = 1
  
          del topk_ious, dynamic_ks, pos_idx
  
          anchor_matching_gt = matching_matrix.sum(0)
          if (anchor_matching_gt > 1).sum() > 0:
              _, cost_argmin = torch.min(cost[:, anchor_matching_gt > 1], dim=0)
              matching_matrix[:, anchor_matching_gt > 1] *= 0
              matching_matrix[cost_argmin, anchor_matching_gt > 1] = 1
          fg_mask_inboxes = matching_matrix.sum(0) > 0
          num_fg = fg_mask_inboxes.sum().item()
  
          fg_mask[fg_mask.clone()] = fg_mask_inboxes
  
          matched_gt_inds = matching_matrix[:, fg_mask_inboxes].argmax(0)
          gt_matched_classes = gt_classes[matched_gt_inds]
  
          pred_ious_this_matching = (matching_matrix * pair_wise_ious).sum(0)[
              fg_mask_inboxes
          ]
          return num_fg, gt_matched_classes, pred_ious_this_matching, matched_gt_inds
  
apis: torch.cuda.empty_cache torch.min torch.stack torch.nn.L1Loss torch.no_grad torch.zeros_like torch.full torch.exp torch.topk torch.log torch.nn.ModuleList torch.cuda.amp.autocast torch.nn.BCEWithLogitsLoss torch.arange torch.zeros torch.cat