Hi, I am training on large images of size 1024*1024 for an image segmentation problem. Since I am able to use a batch size of 4 on 4gpu’s, I wanted to freeze the Batch normalization layer. When looking at the source code of BnFreeze callback I realized the function expects the gradients of the BN layer to be false. I could not find a function which does that in fastai. I remember using it in the previous version of fastai. So I wrote a function which does that. Wanted to check if I am doing it in a proper/recommended way.
Function to freeze BN layers.
def freeze_bn(m):
for l in children(m):
if isinstance(l, bn_types):
for o in l.parameters():
o.requires_grad = False
freeze_bn(l)
I pass model to freeze_bn function and then use BnFreeze callback.