Pytoch freeze pre training weights (feature extraction and BN layer)

1. Read the pre training weight

pre_weights = torch.load(model_weights_path, map_location=device)

2. Read the weights of the same layer as the existing model parameter settings in the pre training weights , It can be used when the classification or the number of channels in some layers is modified

net = yourmodel()
pre_dict = {k: v for k, v in pre_weights.items() 
    if net.state_dict()[k].numel() == v.numel()}
# strict = False  Indicates that only the matching weights are read 
missing_keys, unexpected_keys = net.load_state_dict(pre_dict, strict = False)

3. Freeze the pre training weight of feature extraction layer

for params in net.features.parameters():
    params.requires_grad = False

4. because BN Layer parameters are calculated from the values of each channel , stay forward In the automatic implementation of , Instead of updating through gradient calculation and back propagation , Additional freezing required BN Layer weight

def freeze_bn(ly):
    classname = ly.__class__.__name__
    if classname.find('BatchNorm') != -1:
        ly.eval()
net.apply(freeze_bn)

5. Related links

1. pytorch Medium BN Layer Introduction _lpj822 The column -CSDN Blog

2. model.load_state_dict(state_dict, strict=False)_t20134297 The blog of -CSDN Blog