PyTorch模型微调实战指南

微调 torchvision 模型完整指南

一、概述

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

1. 目标检测与实例分割：使用 Mask R-CNN 在自定义数据集（Penn-Fudan 行人数据集）上微调，同时预测边界框和分割掩码。

2. 图像分类：使用 ResNet、AlexNet、VGG、SqueezeNet、DenseNet、Inception v3 等模型在蜜蜂/蚂蚁二分类数据集上微调或特征提取。

两种方法的共同步骤：

• 加载预训练模型

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

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

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

• 训练并评估

二、目标检测与实例分割微调（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])