如何在nnUNet中冻结权重仅微调最后一层（类别从3改1）

步骤1：加载预训练模型并冻结所有层 #

先加载你的预训练nnUNet模型，然后遍历所有参数关闭梯度更新，实现权重冻结：

import torch
from nnunet.network_architecture.generic_UNet import Generic_UNet

# 初始化原模型结构（参数需与预训练模型一致，比如输入通道数、初始特征数、网络深度等）
pretrained_model = Generic_UNet(
    input_channels=1,  # 替换为你的任务输入通道数
    base_num_features=32,
    num_classes=3,  # 原模型的类别数
    num_pool=4,
    conv_op=torch.nn.Conv2d,  # 3D任务替换为torch.nn.Conv3d
    norm_op=torch.nn.InstanceNorm2d,  # 对应原模型的归一化层
    dropout_op=torch.nn.Dropout2d,
)
# 加载预训练权重
pretrained_model.load_state_dict(torch.load("path/to/your/pretrained_weights.pth"))

# 冻结所有层的参数
for param in pretrained_model.parameters():
    param.requires_grad = False

步骤2：替换最后一层为输出1通道的卷积层 #

nnUNet的最终输出层是final_conv，我们直接替换这个层，让它输出1个类别，并开启该层的梯度更新：

# 获取原最后一层的输入通道数
in_channels = pretrained_model.final_conv.in_channels
# 替换最后一层（2D任务用Conv2d，3D用Conv3d）
pretrained_model.final_conv = torch.nn.Conv2d(in_channels, 1, kernel_size=1)

# 确保新层的参数可训练
for param in pretrained_model.final_conv.parameters():
    param.requires_grad = True

步骤3：调整损失函数与优化器 #

因为类别数改为1，需要使用二分类适配的损失函数，同时优化器只传入最后一层的参数：

# 方案1：使用BCEWithLogitsLoss（适合二分类）
criterion = torch.nn.BCEWithLogitsLoss()

# 方案2：使用nnUNet自带的Dice+CE损失（适配二分类场景）
# from nnunet.utilities.loss_functions.dice_loss import DC_and_CE_loss
# criterion = DC_and_CE_loss(
#     {'batch_dice': True, 'smooth': 1e-5, 'do_bg': False},  # do_bg设为False，因为只有1个前景类
#     {}
# )

# 仅优化最后一层的参数
optimizer = torch.optim.Adam(pretrained_model.final_conv.parameters(), lr=1e-4)

步骤4：验证冻结状态（可选） #

可以打印参数状态，确认只有最后一层参数允许梯度更新：

for name, param in pretrained_model.named_parameters():
    print(f"{name}: requires_grad={param.requires_grad}")

这样训练时，只有最后一层的权重会被更新，其余所有预训练层的权重保持不变，同时输出类别数改为1。

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