pytorch模型冻结及BN层存在的坑

一、pytorch模型冻结

冻结模型常用的方法有两种，一种是将需要冻结的参数的梯度设置为False，另一种是只将不动结的参数传给优化器。
方法一：将需要冻结的参数的梯度设置为False。

model = ...
# 冻结模型的某些层
for name,param in model.named_parameters():
    if 'conv' in name:  # 假设我们冻结所有卷积层
        param.requires_grad = False

方法二：只将不动结模型的参数传递给优化器

冻结模型中除了fc1层之外的所有层
optimizer = optim.SGD(model.fc1.parameters(), lr=0.01)

下面我们来举一个简单的例子来说明。

import torch
import torch.nn as nn
import torch.optim as optim

class Model(nn.Module):
    def __init__(self):
        super(Model,self).__init__()
        self.fc1 = nn.Linear(5,3)
        self.fc2 = nn.Linear(3,2)

    def forward(self,x):
        return self.fc2(self.fc1(x))

model = Model()

print(model.fc1.weight)

optim = optim.SGD(model.parameters(),lr = 0.01)
loss_fn = nn.CrossEntropyLoss()

x = torch.randn((1,5))
label = torch.randn((1,2))

for epoch in range(10):
    output = model(x)
    loss = loss_fn(output,label)
    optim.zero_grad()
    loss.backward()
    optim.step()

print(model.fc1.weight)

首先，我们定义一个简单的模型，然后对其进行训练，然后我们分别打印fc1层在训练前和训练后的权重。可以看到，在没有冻结模型的情况下，模型权重在训练过程中会改变。

接着我们使用方法一去冻结fc1的网络参数，不冻结fc2的网络参数，去看一下最终的网络参数变化情况。

import torch
import torch.nn as nn
import torch.optim as optim

class Model(nn.Module):
    def __init__(self):
        super(Model,self).__init__()
        self.fc1 = nn.Linear(5,3)
        self.fc2 = nn.Linear(3,2)

    def forward(self,x):
        return self.fc2(self.fc1(x))

model = Model()


for name,param in model.named_parameters():
    if 'fc1' in name:
        param.requires_grad = False
print(model.fc1.weight,model.fc2.weight)

optim = optim.SGD(model.parameters(),lr = 0.01)
loss_fn = nn.CrossEntropyLoss()

x = torch.randn((1,5))
label = torch.randn((1,2))

for epoch in range(10):
    output = model(x)
    loss = loss_fn(output,label)
    optim.zero_grad()
    loss.backward()
    optim.step()

print(model.fc1.weight,model.fc2.weight)

可以看到，此时fc1的requires_grad被设置为False后，模型的参数在训练前后没有改变；而fc2层的requires_grad为True,所以其训练前后的网络参数发生了改变。

注意：使用方法一时，一定要注意不要拼写错误，requires_grad拼写错误的话也不会报错，但是如果拼写错误的话就不能正常把梯度设置为False了，所以有的时候明明把requires_grad设置为False了却还是没有冻结住模型，去看看是不是单词拼写错了，我有好几次都是拼写错误导致模型没有冻结住但是没有及时发现。

二、pytorch中BN层存在的坑

在训练过程中，有时候在导入预训练模型，在只进行前向传播的情况下网络每次的输出不同，按照道理说只进行前向传播没有进行参数更新，网络参数应该不变，每次的输出应该相同，为什么会出现这个问题呢？
这是因为BN层搞的鬼，BN层的参数（均值和方差）是在前向传播过程中更新的。在训练模式下，BN层会计算输入数据的批次均值和方差，并使用这些统计量来归一化输入数据。同时，BN层还会更新其运行均值和运行方差的估计值，这些估计值在模型评估时使用。在PyTorch中，这个过程是自动进行的，只要模型处于训练模式。在评估模式下，BN层会使用之前训练过程中更新的运行均值和方差来归一化数据，而不是使用当前批次的统计量。这是BN层在训练和评估模式下的一个关键区别，也是为什么在使用BN层时需要正确地在训练和评估模式之间切换的原因。因此，模型在做推理的过程中，如果模型中有BN层，要把模型设置成评估模式，即model.eval()，否则在推理的过程中BN层的参数会变化。

如果模型在训练过程中使用了BN层，通常建议保持BN层的统计信息更新，以获得最佳的模型性能。只有在某些特殊情况下，比如模型微调或者特定层的性能问题，才可能需要冻结BN层的参数。那么如何冻结住BN层呢？同样的道理，只需要把BN层的requries_grad设置为False就可以了。