From 94c86579a37d411d0afffba1447e5d53f8177cdc Mon Sep 17 00:00:00 2001 From: liaoxingyu Date: Sat, 23 May 2020 10:41:13 +0800 Subject: [PATCH] fix(heads): fix bug in reduce head add neck_feat from config, add inplace in leakyrelu for memory save --- fastreid/modeling/heads/__init__.py | 1 + fastreid/modeling/heads/reduction_head.py | 5 ++--- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/fastreid/modeling/heads/__init__.py b/fastreid/modeling/heads/__init__.py index ab2bcd8..24a5462 100644 --- a/fastreid/modeling/heads/__init__.py +++ b/fastreid/modeling/heads/__init__.py @@ -9,3 +9,4 @@ from .build import REID_HEADS_REGISTRY, build_reid_heads # import all the meta_arch, so they will be registered from .linear_head import LinearHead from .bnneck_head import BNneckHead +from .reduction_head import ReductionHead diff --git a/fastreid/modeling/heads/reduction_head.py b/fastreid/modeling/heads/reduction_head.py index f587c70..747add2 100644 --- a/fastreid/modeling/heads/reduction_head.py +++ b/fastreid/modeling/heads/reduction_head.py @@ -13,15 +13,14 @@ from .build import REID_HEADS_REGISTRY class ReductionHead(nn.Module): def __init__(self, cfg, in_feat, num_classes, pool_layer=nn.AdaptiveAvgPool2d(1)): super().__init__() - 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, bias_freeze=True), - nn.LeakyReLU(0.1), + nn.LeakyReLU(0.1, inplace=True), nn.Dropout2d(0.5), ) self.bnneck = get_norm(cfg.MODEL.HEADS.NORM, reduction_dim, cfg.MODEL.HEADS.NORM_SPLIT, bias_freeze=True)