PyTorch模型微调实战指南

微调 torchvision 模型完整指南

一、概述

本指南涵盖两种迁移学习场景：

两种方法的共同步骤：

• 加载预训练模型

• 根据自定义任务的类别数修改最后一层（或必要层）

• 选择更新全部参数（微调）或仅更新新层（特征提取）

• 定义损失函数、优化器、数据加载器

• 训练并评估

二、目标检测与实例分割微调（Mask R-CNN）

2.1 数据集准备

官方源（推荐）：Penn-Fudan Database for Pedestrian Detection and Segmentation。在页面中找到 "Download zipped file here" 的链接即可下载

Penn-Fudan 数据集结构：

PennFudanPed/
PNGImages/ # 原始图像
PedMasks/ # 分割掩码（不同实例用不同颜色编码）

每张图像对应一个掩码图，掩码中像素值 0 表示背景，非零值表示不同实例。

2.2 自定义数据集类 PennFudanDataset

关键要求：继承 torch.utils.data.Dataset，实现 __len__ 和 __getitem__。 返回字典 target 必须包含：

• boxes (FloatTensor[N,4])：[x0, y0, x1, y1] 格式

• labels (Int64Tensor[N])

• image_id (Int64Tensor[1])

• area (Tensor[N])

• iscrowd (UInt8Tensor[N])

• masks (UInt8Tensor[N,H,W]) – 可选，但 Mask R-CNN 需要

代码解释：

class PennFudanDataset(object):
def __getitem__(self, idx):
# 加载图像和 mask
img = Image.open(img_path).convert("RGB")
mask = Image.open(mask_path)
mask = np.array(mask)

# 获取所有实例 ID（非零值）
obj_ids = np.unique(mask)[1:]
# 将 mask 拆分为多个二值 mask
masks = mask == obj_ids[:, None, None]

# 计算每个实例的边界框
boxes = []
for i in range(len(obj_ids)):
pos = np.where(masks[i])
xmin, xmax = np.min(pos[1]), np.max(pos[1])
ymin, ymax = np.min(pos[0]), np.max(pos[0])
boxes.append([xmin, ymin, xmax, ymax])

# 转换为 tensor
boxes = torch.as_tensor(boxes, dtype=torch.float32)
labels = torch.ones((num_objs,), dtype=torch.int64) # 只有一类：人
masks = torch.as_tensor(masks, dtype=torch.uint8)
image_id = torch.tensor([idx])
area = (boxes[:, 3] – boxes[:, 1]) * (boxes[:, 2] – boxes[:, 0])
iscrowd = torch.zeros((num_objs,), dtype=torch.int64)

target = {"boxes": boxes, "labels": labels, "masks": masks,
"image_id": image_id, "area": area, "iscrowd": iscrowd}
return img, target

2.3 模型构建函数 get_model_instance_segmentation

关键点：替换分类头和 mask 预测头，适配自定义类别数（此处为 2：背景 + 人）。

def get_model_instance_segmentation(num_classes):
# 加载 COCO 预训练的 Mask R-CNN
model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True)

# 替换 box 分类器
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

# 替换 mask 预测器
in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
hidden_layer = 256
model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask,
hidden_layer,
num_classes)
return model

2.4 训练流程

辅助模块：需要从 torchvision/references/detection/ 复制 engine.py, utils.py, transforms.py。

• engine.py: pytorch/vision/blob/main/references/detection/engine.py

• utils.py: pytorch/vision/blob/main/references/detection/utils.py

• transforms.py: pytorch/vision/blob/main/references/detection/transforms.py

数据增强：

def get_transform(train):
transforms = [T.ToTensor()]
if train:
transforms.append(T.RandomHorizontalFlip(0.5))
return T.Compose(transforms)

主训练代码结构：

def main():
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
num_classes = 2

# 创建数据集（训练集+验证集，按 170:50 划分）
dataset = PennFudanDataset('PennFudanPed', get_transform(train=True))
dataset_test = PennFudanDataset('PennFudanPed', get_transform(train=False))
indices = torch.randperm(len(dataset)).tolist()
dataset = Subset(dataset, indices[:-50])
dataset_test = Subset(dataset_test, indices[-50:])

# DataLoader 中需使用 utils.collate_fn
data_loader = DataLoader(dataset, batch_size=2, shuffle=True,
collate_fn=utils.collate_fn)

model = get_model_instance_segmentation(num_classes).to(device)
optimizer = torch.optim.SGD([p for p in model.parameters() if p.requires_grad],
lr=0.005, momentum=0.9, weight_decay=0.0005)
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)

for epoch in range(10):
train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq=10)
lr_scheduler.step()
evaluate(model, data_loader_test, device=device)

2.5 结果

• 训练 1 个 epoch 后：bbox AP 60.6%，mask AP 70.4%

• 训练 10 个 epoch 后：bbox AP 79.9%，mask AP 76.1%，模型能准确检测图像中的行人并给出分割掩码。

三、分类模型微调（ResNet, AlexNet, VGG, SqueezeNet, DenseNet, Inception v3）

3.1 两种策略

3.2 辅助函数

设置参数冻结：

def set_parameter_requires_grad(model, feature_extracting):
if feature_extracting:
for param in model.parameters():
param.requires_grad = False

训练与验证函数 train_model：

• 支持 Inception v3 的双输出（主输出 + 辅助输出）损失计算

• 每个 epoch 后计算验证准确率，保存最佳模型

• 返回最佳模型和验证准确率历史

3.3 各模型最后一层修改方法

统一初始化函数 initialize_model 根据 model_name 返回修改后的模型和输入尺寸（Inception v3 需 299×299，其余 224×224）。

3.4 数据准备与优化器

数据增强（以 ResNet 为例）：

from torchvision import transforms

data_transforms = {
'train': transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
]),
'val': transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) # 补全参数
]),
}

优化器：只更新 requires_grad=True 的参数（特征提取时只有新层可训练）。

params_to_update = model.parameters()
if feature_extract:
params_to_update = [p for p in model.parameters() if p.requires_grad]
optimizer = optim.SGD(params_to_update, lr=0.001, momentum=0.9)

3.5 训练与对比

• 预训练模型微调（SqueezeNet, 15 epochs）：验证准确率 92.8%

• 从头训练：验证准确率仅 45.8%（几乎无效），充分说明迁移学习的巨大优势。

可视化对比：绘制验证准确率随 epoch 变化曲线，预训练模型收敛快且精度高。

四、关键注意事项与实操建议

五、总结

通过这两个完整教程，您应该掌握了：

• 如何为自定义目标检测/分割数据集编写 Dataset 类，并利用预训练 Mask R-CNN 微调。

• 如何对不同架构的分类模型进行迁移学习（微调或特征提取），包括修改最后一层、冻结参数、训练评估。

• 迁移学习能够在小数据集上显著提升性能，避免从头训练的失败。

所有代码均可直接运行（需下载对应数据集），是实践计算机视觉迁移学习的极佳起点。您可以根据自己的任务调整类别数、骨干网络和超参数。