Should I use model.eval() when I freeze BatchNorm layers to finetune?

vision
5

It depends if they were set to .eval() before, but the default mode is train() after loading the model.
If you want to set the complete model to eval mode, just use model.eval().
Alternatively, if you just want to apply it on all batch norm layers, you could use:

def set_bn_eval(module):
    if isinstance(module, torch.nn.modules.batchnorm._BatchNorm):
        module.eval()
        
model.apply(set_bn_eval)
9 Likes