pytorch中神经网络中的冻结层

import torch
import torch.nn as nn

# 假设我们有一个简单的神经网络
class SimpleNet(nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3)
        self.fc1 = nn.Linear(64 * 6 * 6, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = x.view(-1, 64 * 6 * 6)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# 创建模型实例
model = SimpleNet()

# 冻结前两层
for name, param in model.named_parameters():
    if 'conv1' in name or 'conv2' in name:
        param.requires_grad = False

# 打印需要更新的参数
for name, param in model.named_parameters():
    print(name, param.requires_grad)