From e4f09cecf50a54540f7633190387ea060b1bc49e Mon Sep 17 00:00:00 2001 From: xiaohangzhan Date: Wed, 8 Jul 2020 23:18:54 +0800 Subject: [PATCH] support bias in init_weights in ClsHead --- openselfsup/models/heads/cls_head.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/openselfsup/models/heads/cls_head.py b/openselfsup/models/heads/cls_head.py index e8b9fd5f..f9213be8 100644 --- a/openselfsup/models/heads/cls_head.py +++ b/openselfsup/models/heads/cls_head.py @@ -25,13 +25,13 @@ class ClsHead(nn.Module): self.avg_pool = nn.AdaptiveAvgPool2d((1, 1)) self.fc_cls = nn.Linear(in_channels, num_classes) - def init_weights(self, init_linear='normal', std=0.01): + def init_weights(self, init_linear='normal', std=0.01, bias=0.): assert init_linear in ['normal', 'kaiming'], \ "Undefined init_linear: {}".format(init_linear) for m in self.modules(): if isinstance(m, nn.Linear): if init_linear == 'normal': - normal_init(m, std=std) + normal_init(m, std=std, bias=bias) else: kaiming_init(m, mode='fan_in', nonlinearity='relu') elif isinstance(m,