Inheritance diagram for fastreid.modeling.heads.attr_head.AttrHead:
Public Member Functions | |
| __init__ (self, cfg) | |
| forward (self, features, targets=None) | |
Public Attributes | |
| bottleneck | |
Detailed Description
Definition at line 16 of file attr_head.py.
Constructor & Destructor Documentation
◆ __init__()
| fastreid.modeling.heads.attr_head.AttrHead.__init__ | ( | self, | |
| cfg ) |
Definition at line 17 of file attr_head.py.
17 def __init__(self, cfg):
18 super().__init__()
19 # fmt: off
20 feat_dim = cfg.MODEL.BACKBONE.FEAT_DIM
21 num_classes = cfg.MODEL.HEADS.NUM_CLASSES
22 pool_type = cfg.MODEL.HEADS.POOL_LAYER
23 cls_type = cfg.MODEL.HEADS.CLS_LAYER
24 with_bnneck = cfg.MODEL.HEADS.WITH_BNNECK
25 norm_type = cfg.MODEL.HEADS.NORM
26
27 if pool_type == 'fastavgpool': self.pool_layer = FastGlobalAvgPool2d()
28 elif pool_type == 'avgpool': self.pool_layer = nn.AdaptiveAvgPool2d(1)
29 elif pool_type == 'maxpool': self.pool_layer = nn.AdaptiveMaxPool2d(1)
30 elif pool_type == 'gempoolP': self.pool_layer = GeneralizedMeanPoolingP()
31 elif pool_type == 'gempool': self.pool_layer = GeneralizedMeanPooling()
32 elif pool_type == "avgmaxpool": self.pool_layer = AdaptiveAvgMaxPool2d()
33 elif pool_type == 'clipavgpool': self.pool_layer = ClipGlobalAvgPool2d()
34 elif pool_type == "identity": self.pool_layer = nn.Identity()
35 elif pool_type == "flatten": self.pool_layer = Flatten()
36 else: raise KeyError(f"{pool_type} is not supported!")
37
38 # Classification layer
39 if cls_type == 'linear': self.classifier = nn.Linear(feat_dim, num_classes, bias=False)
40 elif cls_type == 'arcSoftmax': self.classifier = ArcSoftmax(cfg, feat_dim, num_classes)
41 elif cls_type == 'circleSoftmax': self.classifier = CircleSoftmax(cfg, feat_dim, num_classes)
42 elif cls_type == 'amSoftmax': self.classifier = AMSoftmax(cfg, feat_dim, num_classes)
43 else: raise KeyError(f"{cls_type} is not supported!")
44 # fmt: on
45
46 # bottleneck = []
47 # if with_bnneck:
48 # bottleneck.append(get_norm(norm_type, feat_dim, bias_freeze=True))
49 bottleneck = [nn.BatchNorm1d(num_classes)]
50
51 self.bottleneck = nn.Sequential(*bottleneck)
52
53 self.bottleneck.apply(weights_init_kaiming)
54 self.classifier.apply(weights_init_classifier)
55
Member Function Documentation
◆ forward()
| fastreid.modeling.heads.attr_head.AttrHead.forward | ( | self, | |
| features, | |||
| targets = None ) |
See :class:`ReIDHeads.forward`.
Definition at line 56 of file attr_head.py.
56 def forward(self, features, targets=None):
57 """
58 See :class:`ReIDHeads.forward`.
59 """
60 global_feat = self.pool_layer(features)
61 global_feat = global_feat[..., 0, 0]
62
63 classifier_name = self.classifier.__class__.__name__
64 # fmt: off
65 if classifier_name == 'Linear': cls_outputs = self.classifier(global_feat)
66 else: cls_outputs = self.classifier(global_feat, targets)
67 # fmt: on
68
69 cls_outputs = self.bottleneck(cls_outputs)
70
71 if self.training:
72 return {
73 "cls_outputs": cls_outputs,
74 }
75 else:
76 cls_outputs = torch.sigmoid(cls_outputs)
77 return cls_outputs
Member Data Documentation
◆ bottleneck
| fastreid.modeling.heads.attr_head.AttrHead.bottleneck |
Definition at line 51 of file attr_head.py.
The documentation for this class was generated from the following file:
- smreid/fastreid/modeling/heads/attr_head.py
