Hi,
I’m currently working on finetuning a large CNN for semantic segmentation and due to GPU memory limitations I can only use a batch size of one.
The CNN I’m using has a bunch of batch normalization layers, which I wan’t to fix during training (since batch normalization with batch size 1 does not make sense). The idea is to set the mode of the batchnorm layers to eval during training. I use the following code to do this
net.train()
for module in net.modules():
        if isinstance(module, torch.nn.modules.BatchNorm1d):
            module.eval()
        if isinstance(module, torch.nn.modules.BatchNorm2d):
            module.eval()
        if isinstance(module, torch.nn.modules.BatchNorm3d):
            module.eval()                        
            

But I can’t seem to get the training to work.
I’ve debugging by doing the same on training code that works fine with a large batch size and not freezing the batch normalization layers. Fixing the batch norm layers will then make the training diverge.
Am I missing something here, does the batchnorm layers backprop function work in eval mode? Is there anything else I need to be doing?
Thanks!