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
# encoding: utf-8
"""
@author:  liaoxingyu
@contact: sherlockliao01@gmail.com
"""
 
from layers import *
from modeling.losses import *
from utils.weight_init import weights_init_kaiming, weights_init_classifier
from .build import REID_HEADS_REGISTRY
 
 
@REID_HEADS_REGISTRY.register()
class ReductionHead(nn.Module):
    def __init__(self, cfg, in_feat, num_classes, pool_layer):
        super().__init__()
        self._cfg = cfg
        reduction_dim = cfg.MODEL.HEADS.REDUCTION_DIM
        self.neck_feat = cfg.MODEL.HEADS.NECK_FEAT
 
        self.pool_layer = pool_layer
 
        self.bottleneck = nn.Sequential(
            nn.Conv2d(in_feat, reduction_dim, 1, 1, bias=False),
            get_norm(cfg.MODEL.HEADS.NORM, reduction_dim, cfg.MODEL.HEADS.NORM_SPLIT),
            nn.LeakyReLU(0.1, inplace=True),
        )
 
        self.bnneck = get_norm(cfg.MODEL.HEADS.NORM, reduction_dim, cfg.MODEL.HEADS.NORM_SPLIT, bias_freeze=True)
 
        self.bottleneck.apply(weights_init_kaiming)
        self.bnneck.apply(weights_init_kaiming)
 
        # identity classification layer
        cls_type = cfg.MODEL.HEADS.CLS_LAYER
        if cls_type == 'linear':    self.classifier = nn.Linear(reduction_dim, num_classes, bias=False)
        elif cls_type == 'arcface': self.classifier = Arcface(cfg, reduction_dim, num_classes)
        elif cls_type == 'circle':  self.classifier = Circle(cfg, reduction_dim, num_classes)
        else:
            raise KeyError(f"{cls_type} is invalid, please choose from "
                           f"'linear', 'arcface' and 'circle'.")
 
        self.classifier.apply(weights_init_classifier)
 
    def forward(self, features, targets=None):
        """
        See :class:`ReIDHeads.forward`.
        """
        features = self.pool_layer(features)
        global_feat = self.bottleneck(features)
        bn_feat = self.bnneck(global_feat)
        bn_feat = bn_feat[..., 0, 0]
 
        # Evaluation
        if not self.training: return bn_feat
 
        # Training
        try:
            cls_outputs = self.classifier(bn_feat)
            pred_class_logits = cls_outputs.detach()
        except TypeError:
            cls_outputs = self.classifier(bn_feat, targets)
            pred_class_logits = F.linear(F.normalize(bn_feat.detach()), F.normalize(self.classifier.weight.detach()))
        # Log prediction accuracy
        CrossEntropyLoss.log_accuracy(pred_class_logits, targets)
 
        if self.neck_feat == "before":  feat = global_feat[..., 0, 0]
        elif self.neck_feat == "after": feat = bn_feat
        else:
            raise KeyError("MODEL.HEADS.NECK_FEAT value is invalid, must choose from ('after' & 'before')")
 
        return cls_outputs, feat