code: # Copyright (c) OpenMMLab. All rights reserved.
  import torch
  import torch.nn as nn
  from mmcv.runner import BaseModule
  
  from mmdet.models.builder import HEADS
  from ...core import bbox_cxcywh_to_xyxy
  
  
  @HEADS.register_module()
  class EmbeddingRPNHead(BaseModule):
  	"""RPNHead in the `Sparse R-CNN <https://arxiv.org/abs/2011.12450>`_ .
  
  	Unlike traditional RPNHead, this module does not need FPN input, but just
  	decode `init_proposal_bboxes` and expand the first dimension of
  	`init_proposal_bboxes` and `init_proposal_features` to the batch_size.
  
  	Args:
  		num_proposals (int): Number of init_proposals. Default 100.
  		proposal_feature_channel (int): Channel number of
  			init_proposal_feature. Defaults to 256.
  		init_cfg (dict or list[dict], optional): Initialization config dict.
  			Default: None
  	"""
  
  	def __init__(self,
  				 num_proposals=100,
  				 proposal_feature_channel=256,
  				 init_cfg=None,
  				 **kwargs):
  		assert init_cfg is None, 'To prevent abnormal initialization ' \
  								 'behavior, init_cfg is not allowed to be set'
  		super(EmbeddingRPNHead, self).__init__(init_cfg)
  		self.num_proposals = num_proposals
  		self.proposal_feature_channel = proposal_feature_channel
  		self._init_layers()
  
  	def _init_layers(self):
  		"""Initialize a sparse set of proposal boxes and proposal features."""
  		self.init_proposal_bboxes = nn.Embedding(self.num_proposals, 4)
  		self.init_proposal_features = nn.Embedding(self.num_proposals, self.proposal_feature_channel)
  
  	def init_weights(self):
  		"""Initialize the init_proposal_bboxes as normalized.
  
  		[c_x, c_y, w, h], and we initialize it to the size of  the entire
  		image.
  		"""
  		super(EmbeddingRPNHead, self).init_weights()
  		nn.init.constant_(self.init_proposal_bboxes.weight[:, :2], 0.5)
  		nn.init.constant_(self.init_proposal_bboxes.weight[:, 2:], 1)
  
  	def _decode_init_proposals(self, imgs, img_metas):
  		"""Decode init_proposal_bboxes according to the size of images and
  		expand dimension of init_proposal_features to batch_size.
  
  		Args:
  			imgs (list[Tensor]): List of FPN features.
  			img_metas (list[dict]): List of meta-information of
  				images. Need the img_shape to decode the init_proposals.
  
  		Returns:
  			Tuple(Tensor):
  
  				- proposals (Tensor): Decoded proposal bboxes,
  				  has shape (batch_size, num_proposals, 4).
  				- init_proposal_features (Tensor): Expanded proposal
  				  features, has shape
  				  (batch_size, num_proposals, proposal_feature_channel).
  				- imgs_whwh (Tensor): Tensor with shape
  				  (batch_size, 4), the dimension means
  				  [img_width, img_height, img_width, img_height].
  		"""
  		proposals = self.init_proposal_bboxes.weight.clone()
  		proposals = bbox_cxcywh_to_xyxy(proposals)
  		num_imgs = len(imgs[0])
  		imgs_whwh = []
  		for meta in img_metas:
  			h, w, _ = meta['img_shape']
  			imgs_whwh.append(imgs[0].new_tensor([[w, h, w, h]]))
  		imgs_whwh = torch.cat(imgs_whwh, dim=0)
  		imgs_whwh = imgs_whwh[:, None, :]
  
  		# imgs_whwh has shape (batch_size, 1, 4)
  		# The shape of proposals change from (num_proposals, 4)
  		# to (batch_size ,num_proposals, 4)
  		proposals = proposals * imgs_whwh
  
  		init_proposal_features = self.init_proposal_features.weight.clone()
  		init_proposal_features = init_proposal_features[None].expand(num_imgs, *init_proposal_features.size())
  		return proposals, init_proposal_features, imgs_whwh
  
  	def forward_dummy(self, img, img_metas):
  		"""Dummy forward function.
  
  		Used in flops calculation.
  		"""
  		return self._decode_init_proposals(img, img_metas)
  
  	def forward_train(self, img, img_metas):
  		"""Forward function in training stage."""
  		return self._decode_init_proposals(img, img_metas)
  
  	def simple_test_rpn(self, img, img_metas):
  		"""Forward function in testing stage."""
  		return self._decode_init_proposals(img, img_metas)
  
  	def simple_test(self, img, img_metas):
  		"""Forward function in testing stage."""
  		raise NotImplementedError
  
  	def aug_test_rpn(self, feats, img_metas):
  		raise NotImplementedError(
  			'EmbeddingRPNHead does not support test-time augmentation')
  
apis: torch.cat torch.nn.Embedding torch.nn.init.constant_