Pytorch冻结网络层

方法一：

# 冻结网络层
def freeze(model):
    for param in model.parameters():
        param.requires_grad = False

1.如果是冻结某一层（传需要冻结的那一层）

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)

        self.conv2 = nn.Conv2d(6, 16, 5)
        freeze(self.conv2) # 冻结self.conv2层
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        # print(self)
        self.fc3 = nn.Linear(84, 10)

2.如果是冻结某一层上面所有层（传self）

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)

        self.conv2 = nn.Conv2d(6, 16, 5)
        freeze(self) # 冻结self.conv1和self.conv2层
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        # print(self)
        self.fc3 = nn.Linear(84, 10)

方法二：

可以直接在前向传播中：

class Net(nn.Module):
    def __init__():
        self.layer1 = xx
        self.layer2 = xx
        self.fc = xx

    def forward(self.x):
        # 只训练self.fc层，将layer1和layer2都冻结
        with torch.no_grad():
            x = self.layer1(x)
            x = self.layer2(x)
        x = self.fc(x)
        return x

注意：如果用方法二，不可以只冻结某一层，因为设置了no_grad()，在经过需要冻结的那一层后，梯度就不会再向前传递了，结果是前面的都被冻结。