From 3984f0c91d0421df3ad3b3155bdd91c0f1e1fd15 Mon Sep 17 00:00:00 2001 From: liaoxingyu Date: Fri, 24 Apr 2020 12:16:18 +0800 Subject: [PATCH] refactor($modeling/meta): refactor heads output without intermediate variables generated by reid heads, make it more flexible --- fastreid/config/defaults.py | 5 ++++- fastreid/modeling/meta_arch/baseline.py | 7 ++++--- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/fastreid/config/defaults.py b/fastreid/config/defaults.py index 8ea3cf1..ae78c58 100644 --- a/fastreid/config/defaults.py +++ b/fastreid/config/defaults.py @@ -48,8 +48,11 @@ _C.MODEL.BACKBONE.PRETRAIN_PATH = '' _C.MODEL.HEADS = CN() _C.MODEL.HEADS.NAME = "BNneckHead" +# Number of identity _C.MODEL.HEADS.NUM_CLASSES = 751 -# Reduction dimension +# Input feature dimension +_C.MODEL.HEADS.IN_FEAT = 2048 +# Reduction dimension in head _C.MODEL.HEADS.REDUCTION_DIM = 512 # Pooling layer type _C.MODEL.HEADS.POOL_LAYER = 'avgpool' diff --git a/fastreid/modeling/meta_arch/baseline.py b/fastreid/modeling/meta_arch/baseline.py index 2e347c1..029ebfc 100644 --- a/fastreid/modeling/meta_arch/baseline.py +++ b/fastreid/modeling/meta_arch/baseline.py @@ -31,7 +31,9 @@ class Baseline(nn.Module): pool_layer = GeneralizedMeanPoolingP() else: pool_layer = nn.Identity() - self.heads = build_reid_heads(cfg, 2048, pool_layer) + + in_feat = cfg.MODEL.HEADS.IN_FEAT + self.heads = build_reid_heads(cfg, in_feat, pool_layer) def forward(self, inputs): images = inputs["images"] @@ -43,8 +45,7 @@ class Baseline(nn.Module): # training features = self.backbone(images) # (bs, 2048, 16, 8) - logits, global_feat = self.heads(features, targets) - return logits, global_feat, targets + return self.heads(features, targets) def inference(self, images): assert not self.training