mirrors/fast-reid
1
0
Fork 0
mirror of https://github.com/JDAI-CV/fast-reid.git synced 2025-06-03 14:50:47 +08:00
Code Issues Projects Releases Wiki Activity

fix(heads): fix bug in reduce head

Browse Source 
add neck_feat from config, add inplace in leakyrelu for memory save
This commit is contained in:
liaoxingyu 2020-05-23 10:41:13 +08:00
parent c21de64166
commit 94c86579a3
2 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
fastreid/modeling/heads/__init__.py
View File

@ -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

5
fastreid/modeling/heads/reduction_head.py
View File

@ -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)