PyTorch中冻结除指定层外的大部分层实现微调

问题描述 #

我有一个已在数据集上训练完成的模型，修改了最后两层（self.conv_6和self.sigmoid）后，希望进行微调操作：冻结模型中除这两层外的所有层，重新在该数据集上训练模型。模型代码如下：

class model(nn.Module):
def __init__(self, pretrained=False):
    super(model, self).__init__()
    
    self.conv_1 = nn.Conv3d(1024, 1024, kernel_size=(3,1,1), stride=(2,1,1), padding=(1,0,0))
    self.conv_2 = nn.Conv3d(1024, 1024, kernel_size=(3,1,1), stride=(2,1,1), padding=(1,0,0))
    self.conv_3 = nn.Conv3d(1024, 1024, kernel_size=(3,1,1), stride=(2,1,1), padding=(1,0,0))
    self.conv_4 = nn.Conv3d(1024, 1024, kernel_size=(3,1,1), stride=(2,1,1), padding=(1,0,0))
    self.conv_5 = nn.Conv3d(1024, 1024, kernel_size=(3,1,1), stride=(2,1,1), padding=(1,0,0))
    self.conv_6 = nn.Conv3d(1024, 1024, kernel_size=(3,1,1), stride=(2,1,1), padding=(1,0,0))

    
    self.sigmoid = nn.Sigmoid()
    
def forward(self, x):
    
    
    x = self.conv_1(x)
    x = self.conv_2(x)
    x = self.conv_3(x)
    x = self.conv_4(x)
    x = self.conv_5(x)
    
    x = self.conv_6(x)
    y = self.sigmoid(x) 
    return y

解决方案 #

1. 加载预训练模型并更新最后两层 #

先加载你已训练好的模型权重，再替换为你修改后的conv_6和sigmoid层。

2. 冻结目标外的所有层 #

遍历模型参数，将除conv_6和sigmoid之外的所有层参数的requires_grad设为False，这些层在训练时就不会被更新。

3. 配置优化器与训练流程 #

只把可训练的参数（即conv_6和sigmoid的参数）传入优化器，再执行常规训练循环即可。

完整代码示例：

import torch
import torch.nn as nn

# 复用你定义的模型类
class model(nn.Module):
    def __init__(self, pretrained=False):
        super(model, self).__init__()
        
        self.conv_1 = nn.Conv3d(1024, 1024, kernel_size=(3,1,1), stride=(2,1,1), padding=(1,0,0))
        self.conv_2 = nn.Conv3d(1024, 1024, kernel_size=(3,1,1), stride=(2,1,1), padding=(1,0,0))
        self.conv_3 = nn.Conv3d(1024, 1024, kernel_size=(3,1,1), stride=(2,1,1), padding=(1,0,0))
        self.conv_4 = nn.Conv3d(1024, 1024, kernel_size=(3,1,1), stride=(2,1,1), padding=(1,0,0))
        self.conv_5 = nn.Conv3d(1024, 1024, kernel_size=(3,1,1), stride=(2,1,1), padding=(1,0,0))
        self.conv_6 = nn.Conv3d(1024, 1024, kernel_size=(3,1,1), stride=(2,1,1), padding=(1,0,0))

        self.sigmoid = nn.Sigmoid()
        
    def forward(self, x):
        x = self.conv_1(x)
        x = self.conv_2(x)
        x = self.conv_3(x)
        x = self.conv_4(x)
        x = self.conv_5(x)
        
        x = self.conv_6(x)
        y = self.sigmoid(x) 
        return y

# 加载预训练权重（替换为你的权重文件路径）
my_model = model()
my_model.load_state_dict(torch.load("path/to/your/pretrained_model.pth"))

# 更新最后两层为你修改后的结构（示例中保持原结构，实际按你的修改调整）
my_model.conv_6 = nn.Conv3d(1024, 1024, kernel_size=(3,1,1), stride=(2,1,1), padding=(1,0,0))
my_model.sigmoid = nn.Sigmoid()

# 冻结非目标层
for name, param in my_model.named_parameters():
    if not (name.startswith("conv_6") or name.startswith("sigmoid")):
        param.requires_grad = False

# 仅传入可训练参数到优化器
optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, my_model.parameters()), lr=1e-4)

# 示例训练循环（根据你的数据集和任务调整）
criterion = nn.BCELoss()  # 二分类任务示例，替换为你的损失函数
epochs = 10
for epoch in range(epochs):
    my_model.train()
    total_loss = 0.0
    for batch_x, batch_y in your_dataloader:  # 替换为你的数据加载器
        optimizer.zero_grad()
        outputs = my_model(batch_x)
        loss = criterion(outputs, batch_y)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    print(f"Epoch {epoch+1}, Loss: {total_loss/len(your_dataloader):.4f}")

核心要点 #

• 冻结逻辑：requires_grad=False会阻止参数参与梯度计算，确保只有目标层被更新。
• 优化器效率：用filter筛选可训练参数，避免优化器处理不必要的参数，节省计算资源。
• 新层默认状态：你修改后的conv_6和sigmoid层参数默认requires_grad=True，会自动参与训练。

内容的提问来源于stack exchange，提问作者dtr43