text: import torch
  import torch.nn as nn
  import torch.nn.functional as F
  from mmcv.runner import auto_fp16, force_fp32
  from torch.nn.modules.utils import _pair
  
  from mmdet.core import build_bbox_coder, multi_apply, multiclass_nms
  from mmdet.models.builder import HEADS, build_loss
  from mmdet.models.losses import accuracy
  from mmdet.models.utils.CompactBilinearPooling import CompactBilinearPooling
  import pdb
  
  @HEADS.register_module()
  class RGBNBBoxHead(nn.Module):
      """Simplest RoI head, with only two fc layers for classification and
      regression respectively.
      在 RGBN 的原始定义中，来自 noise stream 的 roi 输出要和 来自rgb stream 的输出在 cls 分支要进行 bilinear pooling
      """
      def __init__(self,
                   with_avg_pool=False,
                   with_cls=True,
                   with_reg=True,
                   roi_feat_size=7,
                   in_channels=256,
                   num_classes=80,
                   pooling_size=8 * 2048,
                   bbox_coder=dict(
                       type='DeltaXYWHBBoxCoder',
                       target_means=[0., 0., 0., 0.],
                       target_stds=[0.1, 0.1, 0.2, 0.2]),
                   reg_class_agnostic=False,
                   reg_decoded_bbox=False,
                   loss_cls=dict(
                       type='CrossEntropyLoss',
                       use_sigmoid=False,
                       loss_weight=1.0),
                   loss_bbox=dict(
                       type='SmoothL1Loss', beta=1.0, loss_weight=1.0)):
          super(RGBNBBoxHead, self).__init__()
          assert with_cls or with_reg
          self.with_avg_pool = with_avg_pool
          self.with_cls = with_cls
          self.with_reg = with_reg
          self.roi_feat_size = _pair(roi_feat_size)
          self.roi_feat_area = self.roi_feat_size[0] * self.roi_feat_size[1]
          self.in_channels = in_channels
          self.num_classes = num_classes
          self.reg_class_agnostic = reg_class_agnostic
          self.reg_decoded_bbox = reg_decoded_bbox
          self.fp16_enabled = False
  
          self.bbox_coder = build_bbox_coder(bbox_coder)
          self.loss_cls = build_loss(loss_cls)
          self.loss_bbox = build_loss(loss_bbox)
  
          in_channels = self.in_channels
          in_channels_before = in_channels
          if self.with_avg_pool:
              self.avg_pool = nn.AvgPool2d(self.roi_feat_size)
          else:
              # 这里默认考虑了 flatten， 这里不不会立即扁平
              in_channels *= self.roi_feat_area
  
  
          if self.with_cls:
              # need to add background class
              # 这里增加了 bilinear pooling
              self.fc_cls = nn.Linear(pooling_size, num_classes + 1)
              # 使用 fuse 的方法
              # self.fc_cls = nn.Linear(in_channels, num_classes + 1)
  
          if self.with_reg:
              out_dim_reg = 4 if reg_class_agnostic else 4 * num_classes
              self.fc_reg = nn.Linear(in_channels, out_dim_reg)
          self.debug_imgs = None
          #TODO sum_pool 能否加上 avg pool
          self.compact_pooling = CompactBilinearPooling(input_dim1=in_channels_before,
                                                        input_dim2=in_channels_before,
                                                        output_dim=pooling_size, sum_pool=True)
          self.compact_pooling.train()
          # self.conv_fuse = nn.Conv2d(in_channels=in_channels_before * 2, out_channels=in_channels_before, kernel_size=1)
  
      def init_weights(self):
          # conv layers are already initialized by ConvModule
          if self.with_cls:
              # nn.init.normal_(self.fc_cls.weight, 0, 0.01)
              nn.init.xavier_uniform_(self.fc_cls.weight)
              nn.init.constant_(self.fc_cls.bias, 0)
              # fuse
              # nn.init.kaiming_uniform_(self.conv_fuse.weight, mode='fan_in', nonlinearity='relu')
              # nn.init.constant_(self.conv_fuse.bias, 0)
          if self.with_reg:
              nn.init.normal_(self.fc_reg.weight, 0, 0.001)
              nn.init.constant_(self.fc_reg.bias, 0)
  
      @auto_fp16()
      def forward(self, x, noise_stream):
          # shape : 【B, in_channels, roi_feat_size, roi_feat_size】
          # print(x.shape)
          # print(noise_stream.shape)
          if self.with_avg_pool:
              x = self.avg_pool(x)
          if self.with_cls:
              x_cls = self.compact_pooling(x.detach(), noise_stream)
              # cls_score = cls_score.view(-1, )  # 默认已经在 height 和 width 的维度上 reduce_sum
              # x = sign * \sqrt(x)
              # cls_score = torch.sign(cls_score) * torch.sqrt(cls_score) + 1e-12
              # x = L2(x)
              # print(cls_score.shape)
              # cls_score = torch.norm(cls_score, dim=1)
              #
              # print(cls_score.shape)
              # fuse
              # cls_score = self.conv_fuse(torch.cat((x, noise_stream), dim=1))
              # x_cls = torch.sign(x_cls) * torch.sqrt(x_cls) + 1e-12
              x_cls = torch.flatten(x_cls, start_dim=1)
              # cls_score = self.fc_cls(cls_score)
          else:
              raise ValueError('应该要有 cls')
  
          bbox_reg_branch = x.view(x.size(0), -1)
          cls_score = self.fc_cls(x_cls) if self.with_cls else None
          # 辅助分类的
          bbox_pred = self.fc_reg(bbox_reg_branch) if self.with_reg else None
          return cls_score, bbox_pred
  
      def _get_target_single(self, pos_bboxes, neg_bboxes, pos_gt_bboxes,
                             pos_gt_labels, cfg):
          num_pos = pos_bboxes.size(0)
          num_neg = neg_bboxes.size(0)
          num_samples = num_pos + num_neg
  
          # original implementation uses new_zeros since BG are set to be 0
          # now use empty & fill because BG cat_id = num_classes,
          # FG cat_id = [0, num_classes-1]
          labels = pos_bboxes.new_full((num_samples, ),
                                       self.num_classes,
                                       dtype=torch.long)
          label_weights = pos_bboxes.new_zeros(num_samples)
          bbox_targets = pos_bboxes.new_zeros(num_samples, 4)
          bbox_weights = pos_bboxes.new_zeros(num_samples, 4)
          if num_pos > 0:
              labels[:num_pos] = pos_gt_labels
              pos_weight = 1.0 if cfg.pos_weight <= 0 else cfg.pos_weight
              label_weights[:num_pos] = pos_weight
              if not self.reg_decoded_bbox:
                  pos_bbox_targets = self.bbox_coder.encode(
                      pos_bboxes, pos_gt_bboxes)
              else:
                  pos_bbox_targets = pos_gt_bboxes
              bbox_targets[:num_pos, :] = pos_bbox_targets
              bbox_weights[:num_pos, :] = 1
          if num_neg > 0:
              label_weights[-num_neg:] = 1.0
  
          return labels, label_weights, bbox_targets, bbox_weights
  
      def get_targets(self,
                      sampling_results,
                      gt_bboxes,
                      gt_labels,
                      rcnn_train_cfg,
                      concat=True):
          pos_bboxes_list = [res.pos_bboxes for res in sampling_results]
          neg_bboxes_list = [res.neg_bboxes for res in sampling_results]
          pos_gt_bboxes_list = [res.pos_gt_bboxes for res in sampling_results]
          pos_gt_labels_list = [res.pos_gt_labels for res in sampling_results]
          labels, label_weights, bbox_targets, bbox_weights = multi_apply(
              self._get_target_single,
              pos_bboxes_list,
              neg_bboxes_list,
              pos_gt_bboxes_list,
              pos_gt_labels_list,
              cfg=rcnn_train_cfg)
  
          if concat:
              labels = torch.cat(labels, 0)
              label_weights = torch.cat(label_weights, 0)
              bbox_targets = torch.cat(bbox_targets, 0)
              bbox_weights = torch.cat(bbox_weights, 0)
          return labels, label_weights, bbox_targets, bbox_weights
  
      @force_fp32(apply_to=('cls_score', 'bbox_pred'))
      def loss(self,
               cls_score,
               bbox_pred,
               rois,
               labels,
               label_weights,
               bbox_targets,
               bbox_weights,
               reduction_override=None):
          losses = dict()
          if cls_score is not None:
              avg_factor = max(torch.sum(label_weights > 0).float().item(), 1.)
              if cls_score.numel() > 0:
                  losses['loss_cls'] = self.loss_cls(
                      cls_score,
                      labels,
                      label_weights,
                      avg_factor=avg_factor,
                      reduction_override=reduction_override)
  
                  losses['acc'] = accuracy(cls_score, labels)
                  pdb.set_trace()
          if bbox_pred is not None:
              bg_class_ind = self.num_classes
              # 0~self.num_classes-1 are FG, self.num_classes is BG
              pos_inds = (labels >= 0) & (labels < bg_class_ind)
              # do not perform bounding box regression for BG anymore.
              if pos_inds.any():
                  if self.reg_decoded_bbox:
                      bbox_pred = self.bbox_coder.decode(rois[:, 1:], bbox_pred)
                  if self.reg_class_agnostic:
                      pos_bbox_pred = bbox_pred.view(
                          bbox_pred.size(0), 4)[pos_inds.type(torch.bool)]
                  else:
                      pos_bbox_pred = bbox_pred.view(
                          bbox_pred.size(0), -1,
                          4)[pos_inds.type(torch.bool),
                             labels[pos_inds.type(torch.bool)]]
                  losses['loss_bbox'] = self.loss_bbox(
                      pos_bbox_pred,
                      bbox_targets[pos_inds.type(torch.bool)],
                      bbox_weights[pos_inds.type(torch.bool)],
                      avg_factor=bbox_targets.size(0),
                      reduction_override=reduction_override)
              else:
                  losses['loss_bbox'] = bbox_pred[pos_inds].sum()
          return losses
  
      @force_fp32(apply_to=('cls_score', 'bbox_pred'))
      def get_bboxes(self,
                     rois,
                     cls_score,
                     bbox_pred,
                     img_shape,
                     scale_factor,
                     rescale=False,
                     cfg=None):
          if isinstance(cls_score, list):
              cls_score = sum(cls_score) / float(len(cls_score))
          scores = F.softmax(cls_score, dim=1) if cls_score is not None else None
  
          if bbox_pred is not None:
              bboxes = self.bbox_coder.decode(
                  rois[:, 1:], bbox_pred, max_shape=img_shape)
          else:
              bboxes = rois[:, 1:].clone()
              if img_shape is not None:
                  bboxes[:, [0, 2]].clamp_(min=0, max=img_shape[1])
                  bboxes[:, [1, 3]].clamp_(min=0, max=img_shape[0])
  
          if rescale and bboxes.size(0) > 0:
              if isinstance(scale_factor, float):
                  bboxes /= scale_factor
              else:
                  scale_factor = bboxes.new_tensor(scale_factor)
                  bboxes = (bboxes.view(bboxes.size(0), -1, 4) /
                            scale_factor).view(bboxes.size()[0], -1)
  
          if cfg is None:
              return bboxes, scores
          else:
              det_bboxes, det_labels = multiclass_nms(bboxes, scores,
                                                      cfg.score_thr, cfg.nms,
                                                      cfg.max_per_img)
  
              return det_bboxes, det_labels
  
      @force_fp32(apply_to=('bbox_preds', ))
      def refine_bboxes(self, rois, labels, bbox_preds, pos_is_gts, img_metas):
          """Refine bboxes during training.
  
          Args:
              rois (Tensor): Shape (n*bs, 5), where n is image number per GPU,
                  and bs is the sampled RoIs per image. The first column is
                  the image id and the next 4 columns are x1, y1, x2, y2.
              labels (Tensor): Shape (n*bs, ).
              bbox_preds (Tensor): Shape (n*bs, 4) or (n*bs, 4*#class).
              pos_is_gts (list[Tensor]): Flags indicating if each positive bbox
                  is a gt bbox.
              img_metas (list[dict]): Meta info of each image.
  
          Returns:
              list[Tensor]: Refined bboxes of each image in a mini-batch.
  
          Example:
              >>> # xdoctest: +REQUIRES(module:kwarray)
              >>> import kwarray
              >>> import numpy as np
              >>> from mmdet.core.bbox.demodata import random_boxes
              >>> self = BBoxHead(reg_class_agnostic=True)
              >>> n_roi = 2
              >>> n_img = 4
              >>> scale = 512
              >>> rng = np.random.RandomState(0)
              >>> img_metas = [{'img_shape': (scale, scale)}
              ...              for _ in range(n_img)]
              >>> # Create rois in the expected format
              >>> roi_boxes = random_boxes(n_roi, scale=scale, rng=rng)
              >>> img_ids = torch.randint(0, n_img, (n_roi,))
              >>> img_ids = img_ids.float()
              >>> rois = torch.cat([img_ids[:, None], roi_boxes], dim=1)
              >>> # Create other args
              >>> labels = torch.randint(0, 2, (n_roi,)).long()
              >>> bbox_preds = random_boxes(n_roi, scale=scale, rng=rng)
              >>> # For each image, pretend random positive boxes are gts
              >>> is_label_pos = (labels.numpy() > 0).astype(np.int)
              >>> lbl_per_img = kwarray.group_items(is_label_pos,
              ...                                   img_ids.numpy())
              >>> pos_per_img = [sum(lbl_per_img.get(gid, []))
              ...                for gid in range(n_img)]
              >>> pos_is_gts = [
              >>>     torch.randint(0, 2, (npos,)).byte().sort(
              >>>         descending=True)[0]
              >>>     for npos in pos_per_img
              >>> ]
              >>> bboxes_list = self.refine_bboxes(rois, labels, bbox_preds,
              >>>                    pos_is_gts, img_metas)
              >>> print(bboxes_list)
          """
          img_ids = rois[:, 0].long().unique(sorted=True)
          assert img_ids.numel() <= len(img_metas)
  
          bboxes_list = []
          for i in range(len(img_metas)):
              inds = torch.nonzero(
                  rois[:, 0] == i, as_tuple=False).squeeze(dim=1)
              num_rois = inds.numel()
  
              bboxes_ = rois[inds, 1:]
              label_ = labels[inds]
              bbox_pred_ = bbox_preds[inds]
              img_meta_ = img_metas[i]
              pos_is_gts_ = pos_is_gts[i]
  
              bboxes = self.regress_by_class(bboxes_, label_, bbox_pred_,
                                             img_meta_)
  
              # filter gt bboxes
              pos_keep = 1 - pos_is_gts_
              keep_inds = pos_is_gts_.new_ones(num_rois)
              keep_inds[:len(pos_is_gts_)] = pos_keep
  
              bboxes_list.append(bboxes[keep_inds.type(torch.bool)])
  
          return bboxes_list
  
      @force_fp32(apply_to=('bbox_pred', ))
      def regress_by_class(self, rois, label, bbox_pred, img_meta):
          """Regress the bbox for the predicted class. Used in Cascade R-CNN.
  
          Args:
              rois (Tensor): shape (n, 4) or (n, 5)
              label (Tensor): shape (n, )
              bbox_pred (Tensor): shape (n, 4*(#class)) or (n, 4)
              img_meta (dict): Image meta info.
  
          Returns:
              Tensor: Regressed bboxes, the same shape as input rois.
          """
          assert rois.size(1) == 4 or rois.size(1) == 5, repr(rois.shape)
  
          if not self.reg_class_agnostic:
              label = label * 4
              inds = torch.stack((label, label + 1, label + 2, label + 3), 1)
              bbox_pred = torch.gather(bbox_pred, 1, inds)
          assert bbox_pred.size(1) == 4
  
          if rois.size(1) == 4:
              new_rois = self.bbox_coder.decode(
                  rois, bbox_pred, max_shape=img_meta['img_shape'])
          else:
              bboxes = self.bbox_coder.decode(
                  rois[:, 1:], bbox_pred, max_shape=img_meta['img_shape'])
              new_rois = torch.cat((rois[:, [0]], bboxes), dim=1)
  
          return new_rois
  
  
  import torch.nn as nn
  from mmcv.cnn import ConvModule
  
  from mmdet.models.builder import HEADS
  from .bbox_head import BBoxHead
  
  
  @HEADS.register_module()
  class RGBNConvFCBBoxHead(BBoxHead):
      # 共享参数版本的， 需要加入
      r"""More general bbox head, with shared conv and fc layers and two optional
      separated branches.
  
      .. code-block:: none
  
                                      /-> cls convs -> cls fcs -> cls
          shared convs -> shared fcs
                                      \->  reg convs -> reg fcs -> reg
                                              
                                      
      """  # noqa: W605
  
      def __init__(self,
                   # 共享的 conv fc 数量
                   num_shared_convs=0,
                   num_shared_fcs=0,
                   # cls 的 conv 的 fcs
                   num_cls_convs=0,
                   num_cls_fcs=0,
                   num_reg_convs=0,
                   num_reg_fcs=0,
                   # Pooling 输出的 conv
                   conv_out_channels=256,
                   fc_out_channels=1024,
                   pooling_cls_out_dim=2048*8,
                   conv_cfg=None,
                   norm_cfg=None,
                   *args,
                   **kwargs):
          super(RGBNConvFCBBoxHead, self).__init__(*args, **kwargs)
          assert (num_shared_convs + num_shared_fcs + num_cls_convs +
                  num_cls_fcs + num_reg_convs + num_reg_fcs > 0)
          if num_cls_convs > 0 or num_reg_convs > 0:
              assert num_shared_fcs == 0
          if not self.with_cls:
              assert num_cls_convs == 0 and num_cls_fcs == 0
          if not self.with_reg:
              assert num_reg_convs == 0 and num_reg_fcs == 0
          self.num_shared_convs = num_shared_convs
          self.num_shared_fcs = num_shared_fcs
          self.num_cls_convs = num_cls_convs
          self.num_cls_fcs = num_cls_fcs
          self.num_reg_convs = num_reg_convs
          self.num_reg_fcs = num_reg_fcs
          self.conv_out_channels = conv_out_channels
          self.fc_out_channels = fc_out_channels
          self.conv_cfg = conv_cfg
          self.norm_cfg = norm_cfg
  
          # add shared convs and fcs
          self.shared_convs, self.shared_fcs, last_layer_dim = \
              self._add_conv_fc_branch(
                  self.num_shared_convs, self.num_shared_fcs, self.in_channels,
                  True)
          self.shared_out_channels = last_layer_dim
  
          # add cls specific branch
          self.cls_convs, self.cls_fcs, self.cls_last_dim = \
              self._add_conv_fc_branch(
                  self.num_cls_convs, self.num_cls_fcs, self.shared_out_channels)
          self.compact_pooling = CompactBilinearPooling(input_dim1=self.cls_last_dim, input_dim2=self.in_channels,
                                                        output_dim=pooling_cls_out_dim,
                                                        sum_pool=True)
  
          # add reg specific branch
          self.reg_convs, self.reg_fcs, self.reg_last_dim = \
              self._add_conv_fc_branch(
                  self.num_reg_convs, self.num_reg_fcs, self.shared_out_channels)
  
          if self.num_shared_fcs == 0 and not self.with_avg_pool:
              raise NotImplementedError()
              if self.num_cls_fcs == 0:
                  self.cls_last_dim *= self.roi_feat_area
              if self.num_reg_fcs == 0:
                  self.reg_last_dim *= self.roi_feat_area
  
          self.relu = nn.ReLU(inplace=True)
          # reconstruct fc_cls and fc_reg since input channels are changed
          if self.with_cls:
              # self.fc_cls = nn.Linear(self.cls_last_dim, self.num_classes + 1)
              self.fc_cls = nn.Linear(pooling_cls_out_dim, self.num_classes + 1)
          if self.with_reg:
              out_dim_reg = (4 if self.reg_class_agnostic else 4 *
                             self.num_classes)
              self.fc_reg = nn.Linear(self.reg_last_dim, out_dim_reg)
  
  
      def _add_conv_fc_branch(self,
                              num_branch_convs,
                              num_branch_fcs,
                              in_channels,
                              is_shared=False):
          """Add shared or separable branch.
  
          convs -> avg pool (optional) -> fcs
          """
          last_layer_dim = in_channels
          # add branch specific conv layers
          branch_convs = nn.ModuleList()
          if num_branch_convs > 0:
              for i in range(num_branch_convs):
                  conv_in_channels = (
                      last_layer_dim if i == 0 else self.conv_out_channels)
                  branch_convs.append(
                      ConvModule(
                          conv_in_channels,
                          self.conv_out_channels,
                          3,
                          padding=1,
                          conv_cfg=self.conv_cfg,
                          norm_cfg=self.norm_cfg))
              last_layer_dim = self.conv_out_channels
          # add branch specific fc layers
          branch_fcs = nn.ModuleList()
          if num_branch_fcs > 0:
              # for shared branch, only consider self.with_avg_pool
              # for separated branches, also consider self.num_shared_fcs
              if (is_shared
                      or self.num_shared_fcs == 0) and not self.with_avg_pool:
                  last_layer_dim *= self.roi_feat_area
              for i in range(num_branch_fcs):
                  fc_in_channels = (
                      last_layer_dim if i == 0 else self.fc_out_channels)
                  branch_fcs.append(
                      nn.Linear(fc_in_channels, self.fc_out_channels))
              last_layer_dim = self.fc_out_channels
          return branch_convs, branch_fcs, last_layer_dim
  
      def init_weights(self):
          super(RGBNConvFCBBoxHead, self).init_weights()
          # conv layers are already initialized by ConvModule
          for module_list in [self.shared_fcs, self.cls_fcs, self.reg_fcs]:
              for m in module_list.modules():
                  if isinstance(m, nn.Linear):
                      nn.init.xavier_uniform_(m.weight)
                      nn.init.constant_(m.bias, 0)
  
      def forward(self, x, noise_bilinear_pooling_out):
          print(x.shape)
          print(noise_bilinear_pooling_out.shape)
          # shared part
          if self.num_shared_convs > 0:
              for conv in self.shared_convs:
                  x = conv(x)
  
          if self.num_shared_fcs > 0:
              if self.with_avg_pool:
                  x = self.avg_pool(x)
  
              x = x.flatten(1)
  
              for fc in self.shared_fcs:
                  x = self.relu(fc(x))
          # separate branches
          x_cls = x
          x_reg = x
  
          for conv in self.cls_convs:
              x_cls = conv(x_cls)
  
          print(x_cls.shape)
          if x_cls.dim() > 2:
              if self.with_avg_pool:
                  x_cls = self.avg_pool(x_cls)
  
              # 合并 似乎还是需要 特征图的情况
              x_cls = self.compact_pooling(x_cls, noise_bilinear_pooling_out)
              # x = sign * \sqrt(x)
              # x_cls = torch.sign(x_cls) * torch.sqrt(x_cls) + 1e-12
              # # x = L2(x)
              # x_cls = torch.norm(x_cls)
              # 这里应该已经 在 height 和 width 维度上求和了
              # x_cls = x_cls.flatten(1)
          else:
              x_cls = x_cls.unsqueeze(-1).unsqueeze(-1)
              # 合并 似乎还是需要 特征图的情况
              x_cls = self.compact_pooling(x_cls, noise_bilinear_pooling_out)
              # x = sign * \sqrt(x)
              # x_cls = torch.sign(x_cls) * torch.sqrt(x_cls) + 1e-12
              # # x = L2(x)
              # x_cls = torch.norm(x_cls)
              # 这里应该已经 在 height 和 width 维度上求和了
              # x_cls = x_cls.flatten(1)
  
  
          for fc in self.cls_fcs:
              x_cls = self.relu(fc(x_cls))
  
          for conv in self.reg_convs:
              x_reg = conv(x_reg)
          if x_reg.dim() > 2:
              if self.with_avg_pool:
                  x_reg = self.avg_pool(x_reg)
              x_reg = x_reg.flatten(1)
          for fc in self.reg_fcs:
              x_reg = self.relu(fc(x_reg))
  
          cls_score = self.fc_cls(x_cls) if self.with_cls else None
          bbox_pred = self.fc_reg(x_reg) if self.with_reg else None
          return cls_score, bbox_pred
  
  
  @HEADS.register_module()
  class RGBNShared2FCBBoxHead(RGBNConvFCBBoxHead):
  
      def __init__(self, fc_out_channels=1024, *args, **kwargs):
          super(RGBNShared2FCBBoxHead, self).__init__(
              num_shared_convs=0,
              num_shared_fcs=2,
              num_cls_convs=0,
              num_cls_fcs=0,
              num_reg_convs=0,
              num_reg_fcs=0,
              fc_out_channels=fc_out_channels,
              *args,
              **kwargs)