Batch norm not freezed in train?

I have a problem with training. Model gives good results on train data and absolutely trash on test data, even in case train and test data is equal. I try different learning rates and backbones, and have strange behavior.

I suppose that batchnorms not freezed in train\eval time with freeze_bn flags, because when we do model.train or model.eval it changes batchnorm behaviour. I think we can easy imporve that with

...
        self.bn_freezed = freeze_bn
        if freeze_bn:
            self.freeze_bn()

    def forward(self, input):
        if self.bn_freezed:
            self.freeze_bn()
        x, low_level_feat = self.backbone(input)
        x = self.aspp(x)
        x = self.decoder(x, low_level_feat)
        x = F.interpolate(x, size=input.size()[2:], mode='bilinear', align_corners=True)

        return x
...

May be you know more beatuiful way to freeze bn and not freeze other layers, it will be cool if you share it.

[vazyavazya]

@jaemin93 yes, it how batchnorm works, but when you use flag freeze_bn == True, it is not freezed in this original code in this repo, it become unfreezed if you call model.train() and statistic become shifted. When I use code my abowe, it improve quality on evaluation, because I have small batch size, and I want to train with absolutely freezed batchnorms.

[jaemin93]

@jfzhang95 Thank you!