pytorch batchnormalization freeze

모델 내 requires_grad를 다 False로 뒀는데, 이상하게 앞단 모델이 엉망이 되었다.

 

초기모델과 가장 최근까지 학습했던 모델 내 파라미터를 하나씩 비교하던 중, 다음과 같이 backbone 안에 있던 batchnormalization 값은 다 변경되고 있었던 것을 발견했다.

 

backbone.layer1.0.bn1.running_mean

backbone.layer1.0.bn1.running_var

backbone.layer1.0.bn1.num_batches_tracked

 

이상 3가지 값들 위주로 변경되고 있었으며, 이 값들은 일반적인 layer 내 파라미터들처럼 optimizer가 loss로 update하는 값이 아니라고 한다. 해당 값들은 forward 과정에서 버퍼에 저장되어 변동되는 값들로, BN 내 값이 변경되지 않기 위해서는 각 BN layer를 .eval() 모드로 변경해주어야 한다.

 

아래 링크가 해당 내용에 대한 이슈를 정리해 놓은 포스팅이다.

How to freeze BN layers while training the rest of network (mean and var wont freeze)

 

결론적으로, BN 값이 변경되지 않기 위해서는 다음과 같이 수행하면 된다고 한다.

for module in model.modules():
    # print(module)
    if isinstance(module, nn.BatchNorm2d):
        if hasattr(module, 'weight'):
            module.weight.requires_grad_(False)
        if hasattr(module, 'bias'):
            module.bias.requires_grad_(False)
        module.eval()

 

+) 이랬는데도 변경되는것 같은데, module.affine, module.track_running_stats 을 False로 바꿔주었다

+) https://pytorch.org/docs/stable/generated/torch.nn.BatchNorm2d.html

BatchNorm2d — PyTorch 1.9.1 documentation