PyTorch实战：TorchVision目标检测模型微调完整教程

一、什么是微调（Finetuning）？

微调（Finetuning）是指在已经预训练好的模型基础上，使用自己的数据对模型进行进一步训练，使之快速适应特定任务。

在深度学习领域，从头训练一个模型往往需要大量计算资源与数据。微调的出现极大降低了训练成本和门槛，帮助开发者在有限数据集上快速实现高性能模型。

在本教程中，我们将以目标检测（Object Detection）为例，带领大家用PyTorch和TorchVision完成微调过程。

二、PyTorch与TorchVision基础介绍

PyTorch是什么？

PyTorch是Facebook开发的开源深度学习框架，因其灵活易用的动态图机制，迅速成为了学术界与工业界的主流工具。

TorchVision是什么？

TorchVision是PyTorch官方提供的计算机视觉工具包，包含：

• o 常用的图像数据集（如COCO、CIFAR）
• o 预训练模型（如ResNet、VGG、Faster R-CNN）
• o 数据转换工具（如Transforms）

我们本次教程的核心便是利用TorchVision提供的预训练目标检测模型进行微调。

三、数据集介绍与准备

本教程使用的Penn-Fudan Pedestrian Dataset，包含170张图像，用于行人检测任务。每张图像都附有行人的语义掩膜(mask)标注。

数据下载链接：
Penn-Fudan Pedestrian Dataset下载地址

数据集结构说明

下载解压后，文件夹结构如下：

PennFudanPed/
├── Annotation        # 标注信息（XML格式，此处未用到）
├── PedMasks          # 行人掩膜（mask）
└── PNGImages         # 原始图片

数据示例图：

数据集示例图（PyTorch官方教程）

（图片来源：PyTorch官方教程）

数据标注格式

掩膜(mask)图像以整数形式标注，每个像素表示所属对象的类别编号，背景为0，每个对象（如行人）依次编号（1,2,…）。

四、自定义Dataset类创建（含详细中文注释）

我们定义一个继承torch.utils.data.Dataset的数据加载类：

import os
import numpy as np
import torch
from PIL import Image

class PennFudanDataset(torch.utils.data.Dataset):
    def __init__(self, root, transforms=None):
        self.root = root
        self.transforms = transforms
        self.imgs = sorted(os.listdir(os.path.join(root, "PNGImages")))
        self.masks = sorted(os.listdir(os.path.join(root, "PedMasks")))

    def __getitem__(self, idx):
        # 读取图像和掩膜
        img_path = os.path.join(self.root, "PNGImages", self.imgs[idx])
        mask_path = os.path.join(self.root, "PedMasks", self.masks[idx])
        img = Image.open(img_path).convert("RGB")
        mask = Image.open(mask_path)

        mask = np.array(mask)
        obj_ids = np.unique(mask)[1:]  # 去除背景编号0

        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])

        # 转换为torch.Tensor类型
        boxes = torch.as_tensor(boxes, dtype=torch.float32)
        labels = torch.ones((len(obj_ids),), dtype=torch.int64)  # 行人标签为1

        image_id = torch.tensor([idx])
        area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])
        iscrowd = torch.zeros((len(obj_ids),), dtype=torch.int64)

        target = {}
        target["boxes"] = boxes
        target["labels"] = labels
        target["image_id"] = image_id
        target["area"] = area
        target["iscrowd"] = iscrowd

        if self.transforms:
            img, target = self.transforms(img, target)

        return img, target

    def __len__(self):
        return len(self.imgs)

五、模型选择与修改（预训练模型微调）

本教程使用TorchVision内置的Faster R-CNN模型：

Faster R-CNN网络结构：

Faster R-CNN结构图（PyTorch官方教程）

（图片来源：PyTorch官方教程）

使用预训练模型代码（含中文注释）

import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor

def get_model(num_classes):
    # 加载在COCO上预训练的Faster R-CNN
    model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)

    # 获取分类器输入特征数
    in_features = model.roi_heads.box_predictor.cls_score.in_features

    # 替换分类头，适应我们的数据集（只有背景和行人两个类别）
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

    return model

六、模型训练流程详解（代码详解）

1. 数据加载与数据增强

import utils
from engine import train_one_epoch, evaluate
import transforms as T

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

2. 完整训练脚本

device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

dataset = PennFudanDataset('PennFudanPed', get_transform(train=True))
dataset_test = PennFudanDataset('PennFudanPed', get_transform(train=False))

data_loader = torch.utils.data.DataLoader(dataset, batch_size=2, shuffle=True, collate_fn=utils.collate_fn)
data_loader_test = torch.utils.data.DataLoader(dataset_test, batch_size=1, shuffle=False, collate_fn=utils.collate_fn)

model = get_model(num_classes=2)
model.to(device)

params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params, lr=0.005, momentum=0.9, weight_decay=0.0005)

num_epochs = 10
for epoch in range(num_epochs):
    # 训练一个epoch
    train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq=10)
    # 在测试集上评估
    evaluate(model, data_loader_test, device=device)

七、模型评估与测试（代码与指标分析）

模型训练后，使用以下代码进行预测与可视化：

import matplotlib.pyplot as plt
from torchvision.utils import draw_bounding_boxes

model.eval()
img, _ = dataset_test[0]
with torch.no_grad():
    prediction = model([img.to(device)])

boxes = prediction[0]['boxes'].cpu()
scores = prediction[0]['scores'].cpu()

img = (img * 255).type(torch.uint8)
img_with_boxes = draw_bounding_boxes(img, boxes, labels=[f'{s:.2f}' for s in scores])

plt.figure(figsize=(12, 8))
plt.imshow(img_with_boxes.permute(1, 2, 0))
plt.axis('off')
plt.show()

预测效果图示例：

模型预测结果示例图（PyTorch官方教程）

（图片来源：PyTorch官方教程）

八、常见问题及排错技巧

1. 1. 模型不收敛怎么办？
   调整学习率，检查数据标注正确性。
2. 2. GPU内存不足怎么办？
   减小batch_size或裁剪图像尺寸。
3. 3. 为什么模型预测框位置偏差较大？
   检查训练数据标注框准确性，增加数据增强方法。

九、总结与延伸阅读

本教程详细介绍了使用PyTorch与TorchVision完成目标检测微调的全过程，帮助初中级开发者掌握核心技巧，推荐进一步阅读PyTorch官方教程深入学习。

原文链接：PyTorch TorchVision目标检测微调官方教程