text: # encoding: utf-8
  """
  @author:  liaoxingyu
  @contact: sherlockliao01@gmail.com
  """
  
  import torch.nn.functional as F
  from torch import nn
  
  from fastreid.layers import *
  from fastreid.utils.weight_init import weights_init_kaiming, weights_init_classifier
  from .build import REID_HEADS_REGISTRY
  
  
  @REID_HEADS_REGISTRY.register()
  
  
  class ADD_AttrHead(nn.Module):
      def __init__(self, cfg):
          super().__init__()
          # fmt: off
          feat_dim      = cfg.MODEL.HEADS.IN_FEAT
          num_classes   = cfg.MODEL.HEADS.NUM_CLASSES
          pool_type     = cfg.MODEL.HEADS.POOL_LAYER
          cls_type      = cfg.MODEL.HEADS.CLS_LAYER
          with_bnneck   = cfg.MODEL.HEADS.WITH_BNNECK
  
          if pool_type == 'fastavgpool':   self.pool_layer = FastGlobalAvgPool2d()
          elif pool_type == 'avgpool':     self.pool_layer = nn.AdaptiveAvgPool2d(1)
          elif pool_type == 'maxpool':     self.pool_layer = nn.AdaptiveMaxPool2d(1)
          elif pool_type == 'gempoolP':    self.pool_layer = GeneralizedMeanPoolingP()
          elif pool_type == 'gempool':     self.pool_layer = GeneralizedMeanPooling()
          elif pool_type == "avgmaxpool":  self.pool_layer = AdaptiveAvgMaxPool2d()
          elif pool_type == 'clipavgpool': self.pool_layer = ClipGlobalAvgPool2d()
          elif pool_type == "identity":    self.pool_layer = nn.Identity()
          elif pool_type == "flatten":     self.pool_layer = Flatten()
          else:                            raise KeyError(f"{pool_type} is not supported!")
  
          # Classification layer
          if cls_type == 'linear':          self.classifier = nn.Linear(feat_dim, num_classes, bias=False)
          elif cls_type == 'arcSoftmax':    self.classifier = ArcSoftmax(cfg, feat_dim, num_classes)
          elif cls_type == 'circleSoftmax': self.classifier = CircleSoftmax(cfg, feat_dim, num_classes)
          elif cls_type == 'cosSoftmax':     self.classifier = CosSoftmax(cfg, feat_dim, num_classes)
          else:                             raise KeyError(f"{cls_type} is not supported!")
          # fmt: on
  
          bottleneck = []
          if with_bnneck:
              bottleneck.append(nn.BatchNorm1d(num_classes))
  
          self.bottleneck = nn.Sequential(*bottleneck)
  
          self.bottleneck.apply(weights_init_kaiming)
          self.classifier.apply(weights_init_classifier)
  
      def forward(self, features, targets=None ,training=None):
          """
          See :class:`ReIDHeads.forward`.
          """
          global_feat = self.pool_layer(features)
          global_feat = global_feat[..., 0, 0]
  
          classifier_name = self.classifier.__class__.__name__
          # fmt: off
          if classifier_name == 'Linear': cls_outputs = self.classifier(global_feat)
          else:                           cls_outputs = self.classifier(global_feat, targets)
          # fmt: on
  
          cls_outputs = self.bottleneck(cls_outputs)
  
          if not self.training:
              #cls_outputs = torch.sigmoid(cls_outputs)
              return {
                  "cls_outputs": cls_outputs,
              }
          else:
              return {
                  "cls_outputs": cls_outputs,
              }
  
  