PyTorch模型微调实战指南
微调 torchvision 模型完整指南一、概述本指南涵盖两种迁移学习场景:目标检测与实例分割:使用 Mask R-CNN 在自定义数据集(Penn-Fudan 行人数据集)上微调,同时预测边界框和分割掩码。图像分类:使用 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])