class_docstr: Implements the PETR transformer.
  
  Args:
      as_two_stage (bool): Generate query from encoder features.
          Default: False.
      num_feature_levels (int): Number of feature maps from FPN:
          Default: 4.
      two_stage_num_proposals (int): Number of proposals when set
          `as_two_stage` as True. Default: 300.
class_docstr_tok: ["Implements","the","PETR","transformer",".","Args",":","as_two_stage","(","bool",")",":","Generate","query","from","encoder","features",".","Default",":","False",".","num_feature_levels","(","int",")",":","Number","of","feature","maps","from","FPN",":","Default",":","4.","two_stage_num_proposals","(","int",")",":","Number","of","proposals","when","set","`","as_two_stage","`","as","True",".","Default",":","300","."]
docstr: Forward function for `Transformer`.
  
  Args:
      mlvl_feats (list(Tensor)): Input queries from different level.
          Each element has shape [bs, embed_dims, h, w].
      mlvl_masks (list(Tensor)): The key_padding_mask from different
          level used for encoder and decoder, each element has shape
          [bs, h, w].
      query_embed (Tensor): The query embedding for decoder,
          with shape [num_query, c].
      mlvl_pos_embeds (list(Tensor)): The positional encoding
          of feats from different level, has the shape
           [bs, embed_dims, h, w].
      kpt_branches (obj:`nn.LayerList`): Keypoint Regression heads for
          feature maps from each decoder layer. Only would be passed when
          `with_box_refine` is Ture. Default to None.
      cls_branches (obj:`nn.LayerList`): Classification heads for
          feature maps from each decoder layer. Only would be passed when
          `as_two_stage` is Ture. Default to None.
  
  Returns:
      tuple[Tensor]: results of decoder containing the following tensor.
  
          - inter_states: Outputs from decoder. If
              `return_intermediate_dec` is True output has shape                     (num_dec_layers, bs, num_query, embed_dims), else has                     shape (1, bs, num_query, embed_dims).
          - init_reference_out: The initial value of reference                     points, has shape (bs, num_queries, 4).
          - inter_references_out: The internal value of reference                     points in decoder, has shape                     (num_dec_layers, bs,num_query, embed_dims)
          - enc_outputs_class: The classification score of proposals                     generated from encoder's feature maps, has shape                     (batch, h*w, num_classes).                     Only would be returned when `as_two_stage` is True,                     otherwise None.
          - enc_outputs_kpt_unact: The regression results generated from                     encoder's feature maps., has shape (batch, h*w, K*2).
              Only would be returned when `as_two_stage` is True,                     otherwise None.
docstr_tok: ["Forward","function","for","`","Transformer","`",".","Args",":","mlvl_feats","(","list","(","Tensor",")",")",":","Input","queries","from","different","level",".","Each","element","has","shape","[","bs",",","embed_dims",",","h",",","w","]",".","mlvl_masks","(","list","(","Tensor",")",")",":","The","key_padding_mask","from","different","level","used","for","encoder","and","decoder",",","each","element","has","shape","[","bs",",","h",",","w","]",".","query_embed","(","Tensor",")",":","The","query","embedding","for","decoder",",","with","shape","[","num_query",",","c","]",".","mlvl_pos_embeds","(","list","(","Tensor",")",")",":","The","positional","encoding","of","feats","from","different","level",",","has","the","shape","[","bs",",","embed_dims",",","h",",","w","]",".","kpt_branches","(","obj",":","`","nn.LayerList","`",")",":","Keypoint","Regression","heads","for","feature","maps","from","each","decoder","layer",".","Only","would","be","passed","when","`","with_box_refine","`","is","Ture",".","Default","to","None",".","cls_branches","(","obj",":","`","nn.LayerList","`",")",":","Classification","heads","for","feature","maps","from","each","decoder","layer",".","Only","would","be","passed","when","`","as_two_stage","`","is","Ture",".","Default","to","None",".","Returns",":","tuple","[","Tensor","]",":","results","of","decoder","containing","the","following","tensor",".","-","inter_states",":","Outputs","from","decoder",".","If","`","return_intermediate_dec","`","is","True","output","has","shape","(","num_dec_layers",",","bs",",","num_query",",","embed_dims",")",",","else","has","shape","(","1",",","bs",",","num_query",",","embed_dims",")",".","-","init_reference_out",":","The","initial","value","of","reference","points",",","has","shape","(","bs",",","num_queries",",","4",")",".","-","inter_references_out",":","The","internal","value","of","reference","points","in","decoder",",","has","shape","(","num_dec_layers",",","bs",",","num_query",",","embed_dims",")","-","enc_outputs_class",":","The","classification","score","of","proposals","generated","from","encoder","'s","feature","maps",",","has","shape","(","batch",",","h","*","w",",","num_classes",")",".","Only","would","be","returned","when","`","as_two_stage","`","is","True",",","otherwise","None",".","-","enc_outputs_kpt_unact",":","The","regression","results","generated","from","encoder","'s","feature","maps.",",","has","shape","(","batch",",","h","*","w",",","K","*","2",")",".","Only","would","be","returned","when","`","as_two_stage","`","is","True",",","otherwise","None","."]
code: def forward(
      self,
      mlvl_feats,
      mlvl_masks,
      query_embed,
      mlvl_pos_embeds,
      kpt_branches=None,
      cls_branches=None,
  ):
      """Forward function for `Transformer`.
  
          Args:
              mlvl_feats (list(Tensor)): Input queries from different level.
                  Each element has shape [bs, embed_dims, h, w].
              mlvl_masks (list(Tensor)): The key_padding_mask from different
                  level used for encoder and decoder, each element has shape
                  [bs, h, w].
              query_embed (Tensor): The query embedding for decoder,
                  with shape [num_query, c].
              mlvl_pos_embeds (list(Tensor)): The positional encoding
                  of feats from different level, has the shape
                   [bs, embed_dims, h, w].
              kpt_branches (obj:`nn.LayerList`): Keypoint Regression heads for
                  feature maps from each decoder layer. Only would be passed when
                  `with_box_refine` is Ture. Default to None.
              cls_branches (obj:`nn.LayerList`): Classification heads for
                  feature maps from each decoder layer. Only would be passed when
                  `as_two_stage` is Ture. Default to None.
  
          Returns:
              tuple[Tensor]: results of decoder containing the following tensor.
  
                  - inter_states: Outputs from decoder. If
                      `return_intermediate_dec` is True output has shape \
                      (num_dec_layers, bs, num_query, embed_dims), else has \
                      shape (1, bs, num_query, embed_dims).
                  - init_reference_out: The initial value of reference \
                      points, has shape (bs, num_queries, 4).
                  - inter_references_out: The internal value of reference \
                      points in decoder, has shape \
                      (num_dec_layers, bs,num_query, embed_dims)
                  - enc_outputs_class: The classification score of proposals \
                      generated from encoder's feature maps, has shape \
                      (batch, h*w, num_classes). \
                      Only would be returned when `as_two_stage` is True, \
                      otherwise None.
                  - enc_outputs_kpt_unact: The regression results generated from \
                      encoder's feature maps., has shape (batch, h*w, K*2).
                      Only would be returned when `as_two_stage` is True, \
                      otherwise None.
          """
      assert self.as_two_stage or query_embed is not None
  
      feat_flatten = []
      mask_flatten = []
      lvl_pos_embed_flatten = []
      spatial_shapes = []
      for lvl, (feat, mask, pos_embed) in enumerate(
          zip(mlvl_feats, mlvl_masks, mlvl_pos_embeds)
      ):
          bs, c, h, w = feat.shape
          spatial_shape = (h, w)
          spatial_shapes.append(spatial_shape)
          feat = feat.flatten(2).transpose((0, 2, 1))
          mask = mask.flatten(1)
          pos_embed = pos_embed.flatten(2).transpose((0, 2, 1))
          lvl_pos_embed = pos_embed + self.level_embeds[lvl].reshape(
              [1, 1, -1]
          )
          lvl_pos_embed_flatten.append(lvl_pos_embed)
          feat_flatten.append(feat)
          mask_flatten.append(mask)
      feat_flatten = paddle.concat(feat_flatten, 1)
      mask_flatten = paddle.concat(mask_flatten, 1)
      lvl_pos_embed_flatten = paddle.concat(lvl_pos_embed_flatten, 1)
      spatial_shapes_cumsum = paddle.to_tensor(
          np.array(spatial_shapes).prod(1).cumsum(0)
      )
      spatial_shapes = paddle.to_tensor(spatial_shapes, dtype="int64")
      level_start_index = paddle.concat(
          (
              paddle.zeros((1,), dtype=spatial_shapes.dtype),
              spatial_shapes_cumsum[:-1],
          )
      )
      valid_ratios = paddle.stack(
          [self.get_valid_ratio(m) for m in mlvl_masks], 1
      )
  
      reference_points = self.get_reference_points(
          spatial_shapes, valid_ratios
      )
  
      memory = self.encoder(
          src=feat_flatten,
          pos_embed=lvl_pos_embed_flatten,
          src_mask=mask_flatten,
          value_spatial_shapes=spatial_shapes,
          reference_points=reference_points,
          value_level_start_index=level_start_index,
          valid_ratios=valid_ratios,
      )
  
      bs, _, c = memory.shape
  
      hm_proto = None
      if self.training:
          hm_memory = paddle.slice(
              memory,
              starts=level_start_index[0],
              ends=level_start_index[1],
              axes=[1],
          )
          hm_pos_embed = paddle.slice(
              lvl_pos_embed_flatten,
              starts=level_start_index[0],
              ends=level_start_index[1],
              axes=[1],
          )
          hm_mask = paddle.slice(
              mask_flatten,
              starts=level_start_index[0],
              ends=level_start_index[1],
              axes=[1],
          )
          hm_reference_points = paddle.slice(
              reference_points,
              starts=level_start_index[0],
              ends=level_start_index[1],
              axes=[1],
          )[:, :, :1, :]
  
          # official code make a mistake of pos_embed to pose_embed, which disable pos_embed
          hm_memory = self.hm_encoder(
              src=hm_memory,
              pose_embed=hm_pos_embed,
              src_mask=hm_mask,
              value_spatial_shapes=spatial_shapes[[0]],
              reference_points=hm_reference_points,
              value_level_start_index=level_start_index[0],
              valid_ratios=valid_ratios[:, :1, :],
          )
          hm_memory = hm_memory.reshape(
              (bs, spatial_shapes[0, 0], spatial_shapes[0, 1], -1)
          )
          hm_proto = (hm_memory, mlvl_masks[0])
  
      if self.as_two_stage:
          (
              output_memory,
              output_proposals,
          ) = self.gen_encoder_output_proposals(
              memory, mask_flatten, spatial_shapes
          )
          enc_outputs_class = cls_branches[self.decoder.num_layers](
              output_memory
          )
          enc_outputs_kpt_unact = kpt_branches[self.decoder.num_layers](
              output_memory
          )
          enc_outputs_kpt_unact[..., 0::2] += output_proposals[..., 0:1]
          enc_outputs_kpt_unact[..., 1::2] += output_proposals[..., 1:2]
  
          topk = self.two_stage_num_proposals
          topk_proposals = paddle.topk(
              enc_outputs_class[..., 0], topk, axis=1
          )[1].unsqueeze(-1)
  
          # paddle.take_along_axis 对应torch.gather
          topk_kpts_unact = paddle.take_along_axis(
              enc_outputs_kpt_unact, topk_proposals, 1
          )
          topk_kpts_unact = topk_kpts_unact.detach()
  
          reference_points = F.sigmoid(topk_kpts_unact)
          init_reference_out = reference_points
          # learnable query and query_pos
          query_pos, query = paddle.split(
              query_embed, query_embed.shape[1] // c, axis=1
          )
          query_pos = query_pos.unsqueeze(0).expand((bs, -1, -1))
          query = query.unsqueeze(0).expand((bs, -1, -1))
      else:
          query_pos, query = paddle.split(
              query_embed, query_embed.shape[1] // c, axis=1
          )
          query_pos = query_pos.unsqueeze(0).expand((bs, -1, -1))
          query = query.unsqueeze(0).expand((bs, -1, -1))
          reference_points = F.sigmoid(self.reference_points(query_pos))
          init_reference_out = reference_points
  
      # decoder
      inter_states, inter_references = self.decoder(
          query=query,
          memory=memory,
          query_pos_embed=query_pos,
          memory_mask=mask_flatten,
          reference_points=reference_points,
          value_spatial_shapes=spatial_shapes,
          value_level_start_index=level_start_index,
          valid_ratios=valid_ratios,
          kpt_branches=kpt_branches,
      )
  
      inter_references_out = inter_references
      if self.as_two_stage:
          return (
              inter_states,
              init_reference_out,
              inter_references_out,
              enc_outputs_class,
              enc_outputs_kpt_unact,
              hm_proto,
              memory,
          )
      return (
          inter_states,
          init_reference_out,
          inter_references_out,
          None,
          None,
          None,
          None,
          None,
          hm_proto,
      )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import math
  import numpy
  import paddle
  import paddle.nn
  import paddle.nn.functional
  from paddle import ParamAttr
  from ppdet.core.workspace import register
  from ..layers import MultiHeadAttention, _convert_attention_mask
  from .utils import _get_clones
  from ..initializer import linear_init_, normal_, constant_, xavier_uniform_
label_desc: Use image node_id 8 for calling the PETRTransformer obj's underlying member method code with example usage: obj.forward(mlvl_feats, mlvl_masks, query_embed, mlvl_pos_embeds, kpt_branches, cls_branches) and returns: inter_states,  init_reference_out,  inter_references_out,  None,  None,  None,  None,  None,  hm_proto, inter_states,  init_reference_out,  inter_references_out,  enc_outputs_class,  enc_outputs_kpt_unact,  hm_proto,  memory