update

Scheaven

2021-09-18 291deeb1fcf45dbf39a24aa72a213ff3fd6b3405

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
# encoding: utf-8 
""" 
@author:  liaoxingyu 
@contact: sherlockliao01@gmail.com 
""" 
 
import torch 
from torch import nn 
 
from layers import * 
from utils.weight_init import weights_init_kaiming, weights_init_classifier 
from .build import REID_HEADS_REGISTRY 
 
 
@REID_HEADS_REGISTRY.register() 
class AttrHead(nn.Module): 
    def __init__(self, cfg): 
        super().__init__() 
        # fmt: off 
        feat_dim      = cfg.MODEL.BACKBONE.FEAT_DIM 
        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 
        norm_type     = cfg.MODEL.HEADS.NORM 
 
        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 == 'amSoftmax':     self.classifier = AMSoftmax(cfg, feat_dim, num_classes) 
        else:                             raise KeyError(f"{cls_type} is not supported!") 
        # fmt: on 
 
        # bottleneck = [] 
        # if with_bnneck: 
        #     bottleneck.append(get_norm(norm_type, feat_dim, bias_freeze=True)) 
        bottleneck = [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): 
        """ 
        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 self.training: 
            return { 
                "cls_outputs": cls_outputs, 
            } 
        else: 
            cls_outputs = torch.sigmoid(cls_outputs) 
            return cls_outputs