freeze bn layer by not require grad but still behave differently when model.train() vs. model.eval()
(model)
| 16 | # from pykalman import KalmanFilter |
| 17 | |
| 18 | def freeze_bn_layer(model): |
| 19 | ''' freeze bn layer by not require grad but still behave differently when model.train() vs. model.eval() ''' |
| 20 | print("Freezing BatchNorm Layers...") |
| 21 | for module in model.modules(): |
| 22 | if isinstance(module, nn.BatchNorm2d): |
| 23 | # print("this is a BN layer:", module) |
| 24 | if hasattr(module, 'weight'): |
| 25 | module.weight.requires_grad_(False) |
| 26 | if hasattr(module, 'bias'): |
| 27 | module.bias.requires_grad_(False) |
| 28 | return model |
| 29 | |
| 30 | def freeze_bn_layer_train(model): |
| 31 | ''' set batchnorm to eval() |
no outgoing calls
no test coverage detected