mirror of
https://github.com/JDAI-CV/fast-reid.git
synced 2025-06-03 14:50:47 +08:00
fix(heads): fix bug in reduce head
add neck_feat from config, add inplace in leakyrelu for memory save
This commit is contained in:
parent
c21de64166
commit
94c86579a3
@ -9,3 +9,4 @@ from .build import REID_HEADS_REGISTRY, build_reid_heads
|
|||||||
# import all the meta_arch, so they will be registered
|
# import all the meta_arch, so they will be registered
|
||||||
from .linear_head import LinearHead
|
from .linear_head import LinearHead
|
||||||
from .bnneck_head import BNneckHead
|
from .bnneck_head import BNneckHead
|
||||||
|
from .reduction_head import ReductionHead
|
||||||
|
@ -13,15 +13,14 @@ from .build import REID_HEADS_REGISTRY
|
|||||||
class ReductionHead(nn.Module):
|
class ReductionHead(nn.Module):
|
||||||
def __init__(self, cfg, in_feat, num_classes, pool_layer=nn.AdaptiveAvgPool2d(1)):
|
def __init__(self, cfg, in_feat, num_classes, pool_layer=nn.AdaptiveAvgPool2d(1)):
|
||||||
super().__init__()
|
super().__init__()
|
||||||
|
|
||||||
reduction_dim = cfg.MODEL.HEADS.REDUCTION_DIM
|
reduction_dim = cfg.MODEL.HEADS.REDUCTION_DIM
|
||||||
|
self.neck_feat = cfg.MODEL.HEADS.NECK_FEAT
|
||||||
|
|
||||||
self.pool_layer = pool_layer
|
self.pool_layer = pool_layer
|
||||||
|
|
||||||
self.bottleneck = nn.Sequential(
|
self.bottleneck = nn.Sequential(
|
||||||
nn.Conv2d(in_feat, reduction_dim, 1, 1, bias=False),
|
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),
|
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),
|
nn.Dropout2d(0.5),
|
||||||
)
|
)
|
||||||
self.bnneck = get_norm(cfg.MODEL.HEADS.NORM, reduction_dim, cfg.MODEL.HEADS.NORM_SPLIT, bias_freeze=True)
|
self.bnneck = get_norm(cfg.MODEL.HEADS.NORM, reduction_dim, cfg.MODEL.HEADS.NORM_SPLIT, bias_freeze=True)
|
||||||
|
Loading…
x
Reference in New Issue
Block a user