From 329764bb608c14242c8905f7533d6aa066f70261 Mon Sep 17 00:00:00 2001 From: liaoxingyu Date: Wed, 29 Apr 2020 21:29:48 +0800 Subject: [PATCH] refactor(heads): move num_classes out from heads set parameter num_classes in meta_arch to easily modify different heads fc layer --- fastreid/modeling/heads/bnneck_head.py | 7 +++---- fastreid/modeling/heads/build.py | 4 ++-- fastreid/modeling/heads/linear_head.py | 8 +++----- fastreid/modeling/heads/reduction_head.py | 7 +++---- fastreid/modeling/meta_arch/baseline.py | 3 ++- 5 files changed, 13 insertions(+), 16 deletions(-) diff --git a/fastreid/modeling/heads/bnneck_head.py b/fastreid/modeling/heads/bnneck_head.py index 7137f14..ca83794 100644 --- a/fastreid/modeling/heads/bnneck_head.py +++ b/fastreid/modeling/heads/bnneck_head.py @@ -11,9 +11,8 @@ from ...layers import * @REID_HEADS_REGISTRY.register() class BNneckHead(nn.Module): - def __init__(self, cfg, in_feat, pool_layer=nn.AdaptiveAvgPool2d(1)): + def __init__(self, cfg, in_feat, num_classes, pool_layer=nn.AdaptiveAvgPool2d(1)): super().__init__() - self._num_classes = cfg.MODEL.HEADS.NUM_CLASSES self.pool_layer = nn.Sequential( pool_layer, @@ -24,13 +23,13 @@ class BNneckHead(nn.Module): # identity classification layer if cfg.MODEL.HEADS.CLS_LAYER == 'linear': - self.classifier = nn.Linear(in_feat, self._num_classes, bias=False) + self.classifier = nn.Linear(in_feat, num_classes, bias=False) elif cfg.MODEL.HEADS.CLS_LAYER == 'arcface': self.classifier = Arcface(cfg, in_feat) elif cfg.MODEL.HEADS.CLS_LAYER == 'circle': self.classifier = Circle(cfg, in_feat) else: - self.classifier = nn.Linear(in_feat, self._num_classes, bias=False) + self.classifier = nn.Linear(in_feat, num_classes, bias=False) def forward(self, features, targets=None): """ diff --git a/fastreid/modeling/heads/build.py b/fastreid/modeling/heads/build.py index ddbb90b..d093c0a 100644 --- a/fastreid/modeling/heads/build.py +++ b/fastreid/modeling/heads/build.py @@ -16,9 +16,9 @@ The call is expected to return an :class:`ROIHeads`. """ -def build_reid_heads(cfg, in_feat, pool_layer): +def build_reid_heads(cfg, in_feat, num_classes, pool_layer): """ Build REIDHeads defined by `cfg.MODEL.REID_HEADS.NAME`. """ head = cfg.MODEL.HEADS.NAME - return REID_HEADS_REGISTRY.get(head)(cfg, in_feat, pool_layer) + return REID_HEADS_REGISTRY.get(head)(cfg, in_feat, num_classes, pool_layer) diff --git a/fastreid/modeling/heads/linear_head.py b/fastreid/modeling/heads/linear_head.py index e1e272e..8b1b148 100644 --- a/fastreid/modeling/heads/linear_head.py +++ b/fastreid/modeling/heads/linear_head.py @@ -10,10 +10,8 @@ from ...layers import * @REID_HEADS_REGISTRY.register() class LinearHead(nn.Module): - - def __init__(self, cfg, in_feat, pool_layer=nn.AdaptiveAvgPool2d(1)): + def __init__(self, cfg, in_feat, num_classes, pool_layer=nn.AdaptiveAvgPool2d(1)): super().__init__() - self._num_classes = cfg.MODEL.HEADS.NUM_CLASSES self.pool_layer = nn.Sequential( pool_layer, @@ -22,13 +20,13 @@ class LinearHead(nn.Module): # identity classification layer if cfg.MODEL.HEADS.CLS_LAYER == 'linear': - self.classifier = nn.Linear(in_feat, self._num_classes, bias=False) + self.classifier = nn.Linear(in_feat, num_classes, bias=False) elif cfg.MODEL.HEADS.CLS_LAYER == 'arcface': self.classifier = Arcface(cfg, in_feat) elif cfg.MODEL.HEADS.CLS_LAYER == 'circle': self.classifier = Circle(cfg, in_feat) else: - self.classifier = nn.Linear(in_feat, self._num_classes, bias=False) + self.classifier = nn.Linear(in_feat, num_classes, bias=False) def forward(self, features, targets=None): """ diff --git a/fastreid/modeling/heads/reduction_head.py b/fastreid/modeling/heads/reduction_head.py index 0208f14..3081c95 100644 --- a/fastreid/modeling/heads/reduction_head.py +++ b/fastreid/modeling/heads/reduction_head.py @@ -11,9 +11,8 @@ from ...layers import * @REID_HEADS_REGISTRY.register() class ReductionHead(nn.Module): - def __init__(self, cfg, in_feat, pool_layer=nn.AdaptiveAvgPool2d(1)): + def __init__(self, cfg, in_feat, num_classes, pool_layer=nn.AdaptiveAvgPool2d(1)): super().__init__() - self._num_classes = cfg.MODEL.HEADS.NUM_CLASSES reduction_dim = cfg.MODEL.HEADS.REDUCTION_DIM self.pool_layer = nn.Sequential( @@ -34,13 +33,13 @@ class ReductionHead(nn.Module): # identity classification layer if cfg.MODEL.HEADS.CLS_LAYER == 'linear': - self.classifier = nn.Linear(reduction_dim, self._num_classes, bias=False) + self.classifier = nn.Linear(reduction_dim, num_classes, bias=False) elif cfg.MODEL.HEADS.CLS_LAYER == 'arcface': self.classifier = Arcface(cfg, reduction_dim) elif cfg.MODEL.HEADS.CLS_LAYER == 'circle': self.classifier = Circle(cfg, reduction_dim) else: - self.classifier = nn.Linear(reduction_dim, self._num_classes, bias=False) + self.classifier = nn.Linear(reduction_dim, num_classes, bias=False) def forward(self, features, targets=None): """ diff --git a/fastreid/modeling/meta_arch/baseline.py b/fastreid/modeling/meta_arch/baseline.py index 07b3f77..d11470e 100644 --- a/fastreid/modeling/meta_arch/baseline.py +++ b/fastreid/modeling/meta_arch/baseline.py @@ -32,7 +32,8 @@ class Baseline(nn.Module): pool_layer = nn.Identity() in_feat = cfg.MODEL.HEADS.IN_FEAT - self.heads = build_reid_heads(cfg, in_feat, pool_layer) + num_classes = cfg.MODEL.HEADS.NUM_CLASSES + self.heads = build_reid_heads(cfg, in_feat, num_classes, pool_layer) def forward(self, inputs): images = inputs["images"]