使用torchvision的目标检测模型微调示例

在这篇演示中，我们将在数据集Penn-Fudan Database for Pedestrian Detection and Segmentation上微调一个预训练模型Mask R-CNN。此数据集包含345个行人实例（instance）的170张图像，我们将使用它来阐述如何使用torchvision在自定义数据集上训练实例分割模型的一些新特性。

数据集说明

这一数据集包含报告在实验Object Detection Combining Recognition and Segmentation提出的行人检测的图像。每张图像来自于校园以及城市街道。我们的目标是图像上的行人。每一张图像上将会有至少一个行人。

在数据集中的标注了的行人的高度在[180, 390]像素。所有标注的行人都是直立状态。

总共有170张图像标注了345个标注行人，其中96张图像拍摄于宾夕法尼亚大学附近，另外74张图像拍摄于复旦大学附近。

定义数据集

在用于目标检测、实例分割以及人物关键点检测的参考脚本允许轻松地添加新的自定义数据集。数据集应该集成标准类torch.utils.data.Dataset，并且实现__len__和__getitem__方法。

唯一我们需要的要求就是__getitem__应该返回一下内容：

• boxes (FloatTensor[N, 4]：N个边界边框的坐标以[x0, y0, x1, y1]的格式给出，并且范围是从0到W和0到H
• labels (Int64Tensor[N]：每个边界边框的label，0总是代表背景类别
• image_id (Int64Tensor[1])：图像的标识符，它在图像数据集中应该是唯一的，并且用于评估阶段
• area (Tensor[N])：每个边界框的面积（area），这将在评估阶段计算COCO指标，并将其分为small、medium和large
• iscrowd (UInt8Tensor[N])：iscrowd=True的实例在评估阶段将被忽略
• （可选）masks (UInt8Tensor[N, H, W])：对于每一个目标的分割遮罩
• （可选）keypoints (FloatTensor[N, K, 3])：对于N个目标，它包含以[x, y, visibility]格式的K个关键点（keypoints）定义。visibility=0表示关键点不可见。注意：对于数据增强（data augmentation），翻转一个关键点（keypoint）的概念是独立于数据表示（data representation）的，并且你大概率要为你的新关键点表示（keypoint representation）调整references/detection/transforms.py文件

如果你的模型返回上面的方法，那么它将会在训练和评估阶段奏效，并且将使用pycocotools的评估脚本，这个工具可以使用pip install pycocotools安装。

注意：对于WIndows系统，请从gautamchitnis中安装pycocotools，使用如下命令：

pip install git+https://github.com/gautamchitnis/cocoapi.git@cocodataset-master#subdirectory=PythonAPI

对于labels需要注意，模型会将类别0认作位背景类别。如果你的数据集没有包含背景类，你不应该在你的labels中出现0。举个例子，假设你只有2个类别，猫和狗，你可以定义1（而不是0）代表猫以及2代表狗。所以，如果有个图像包含两个类别，你的labelstensor应该像这样[1,2]。

另外，如果你想在训练阶段使用aspect ration grouping（也就是每一个数据批量（data batch）的图片都有相似的长宽比），那么推荐做法就是实现get_height_and_width方法，这个方法返回图片的长度和宽度。如果这个方法没有被提供，我们将使用__getitem__查询数据集的所有元素，这将会加载图片到内存，并且会比提供了自定义方法的速度慢。

为PennFudan自定义数据集

下面开始为PennFudan数据编写一个Dataset类，在下载并提取出zip文件后，我们将有如下所示的文件结构：

PennFudanPed/
  PedMasks/
    FudanPed00001_mask.png
    FudanPed00002_mask.png
    FudanPed00003_mask.png
    FudanPed00004_mask.png
    ...
  PNGImages/
    FudanPed00001.png
    FudanPed00002.png
    FudanPed00003.png
    FudanPed00004.png

下面是一对图像和对应的分割遮罩的示例：

对于每一张图片都有对应的分割遮罩，其中每个颜色对应不同的实例（instance），让我们开始为此数据集编写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):
        self.root = root
        self.transforms = transforms
        # 加载所有图片文件，将它们排序
        # 确保图片与遮罩一一对应
        self.imgs = list(sorted(os.listdir(os.path.join(root, "PNGImages"))))
        self.masks = list(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")
        # 注意我们并没有把遮罩也转换为RGB
        # 因为每一个颜色都对应不同的实例
        # 0表示背景
        mask = Image.open(mask_path)
        # 将PIL图像转换成numpy的array
        mask = np.array(mask)
        # 实例以不同的颜色编码
        obj_ids = np.unique(mask)
        # 第一个id是背景，所以移除它
        obj_ids = obj_ids[1:]
        
        # 将编码的颜色遮罩分割为二元遮罩
        masks = mask == obj_ids[:, None, None]
        
        # 对于每个遮罩，获取他们的边界边框坐标
        num_objs = len(obj_ids)
        boxes = []
        for i in range(num_objs):
            # np.where 返回的为二维坐标
            pos = np.where(masks[i])
            xmin = np.min(pos[1])
            xmax = np.max(pos[1])
            ymin = np.min(pos[0])
            ymax = np.max(pos[0])
            boxes.append([xmin, ymin, xmax, ymax])
        # 将所有转换到torch.Tensor
        boxes = torch.as_tensor(boxes, dtype=torch.float32)
        # 仅只有1个类别：行人类
        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])
        # 假设所有的实例都不是人群（crowd）
        iscrowd = torch.zeros((num_objs,), dtype=torch.int64)
        
        target = {}
        target["boxes"] = boxes
        target["labels"] = labels
        target["masks"] = masks
        target["image_id"] = image_id
        target["area"] = area
        target["iscrowd"] = iscrowd
        
        if self.transforms is not None:
            img, target = self.transforms(img, target)

        return img, target
    
    def __len__(self):
        return len(self.imgs)

这就是数据集的所有内容。接下来让我们定义可以在这个数据集上进行预测的模型。

定义模型

在本文，我们将使用基于Faster R-CNN模型之上的Mask R-CNN。Faster R-CNN是可以同时预测在图像中潜在的目标的边界边框和类别分数的模型。

Mask R-CNN添加了额外的分支到Faster R-CNN上，此模型对于每一个实例也预测了分割遮罩。

有两种常见的情况，其中一种就是在torchvision乐园中使用想要修改可用的模型之一。首先，当我们从一个预训练模型开始，简单地微调一下最后一层即可。另一种方法便是当我们想使用不同的架构替换掉模型的骨干（例如，为了更快的预测）。

让我们在接下来的章节看一下我们将会如何做到以上两点。

1 一个预训练模型的微调

假设我们希望从一个在COCO上预训练的模型开始，并且对于我们特定的类别对其进行微调。

以下是一种可能的做法：

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

# 在COCO上加载预训练的模型
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)

# 替换分类器，由用户定义num_classes
num_classes = 2  # 类别 1（行人）+背景类别
# 获得分类器输入特征的数量
in_features = model.roi_heads.box_predictor.cls_score.in_features
# 替换预训练模型的头部
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

2 修改模型：添加新骨干

import torchvision
from torchvision.models.detection import FasterRCNN
from torchvision.models.detection.rpn import AnchorGenerator

# 为分类任务加载预训练模型，并且仅返回特征（features）
backbone = torchvision.models.mobilenet_v2(pretrained=True).features
# Faster RCNN需要知道骨干中输出通道的数量
# 对于mobilenet_v2是1280，所以我们需要为其添加
backbone.out_channels = 1280

# 让我们使用RPN在每一个空间位置生成5x3个锚点（anchor）
# 其中5种不同尺寸和3种不同长宽比
# 我们需要一个Tuple[Tuple[int]]，因为每一个特征图（feature map）潜在地会有不同的尺寸和长宽比
anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512),), aspect_ratios=((0.5, 1.0, 2.0),))

# 让我们定义什么是特征图（feature maps），我们将用其执行裁剪感兴趣区域ROI（region of interest），
# 以及在重新缩放之后裁剪尺寸。
# 如果你的架构返回一个Tensor，featmap_names被期待为[0]。通常来讲，骨干应该返回一个
# OrderedDict[Tensor]，并且在featmap_names中，你可以选择使用哪一个features maps。
roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=['0'], output_size=7, sampling_ratio=2)

# 将部件放进FasterRCNN模型中
model = FasterRCNN(backbone, num_classes=2,
                   rpn_anchor_generator=anchor_generator,
                   box_roi_pool=roi_pooler)

使用PennFudan数据集的实例分割模型

在我们的例子中，我们想要微调一个预训练模型，我们给定的数据集非常小，所以我们将使用以下方法1。

这里，我们也要计算实例分割遮罩，所以我们使用Mask R-CNN：

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

def get_model_instance_segmentation(num_classes):
    # 加载一个在COCO上的预训练的实例分割模型
    model = torchvision.models.detection.maskrcnn_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)
    
    # 获得对于遮罩分类器输入数量
    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

至此，这将使得model将在你自定义的训练集上准备开始训练和评估了。

整合

在references/detection中，我们有很多帮助函数来简化训练和评估检测模型。在这里，我们将使用references/detection/engin.py，references/detection/utils.py以及references/detection/transforms.py。直接将这些文件复制到你的文件目录下并使用它。

让我们编写用于数据增强（data augmentation）\transformation的帮助函数：

import transforms as T
from torchvision.transforms import ToTensor

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

测试forward()方法

在迭代整个数据集之前，让我们看一下在部分数据上模型期待的训练和测试时间。

import utils
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
dataset = PennFudanDataset('./PennFudanPed', get_transform(train=True))
data_loader = torch.utils.data.DataLoader(
    dataset, batch_size=2, shuffle=True, num_workers=4,
    collate_fn=utils.collate_fn)
# 训练
images, targets = next(iter(data_loader))
images = list(image for image in images)
targets = [{k: v for k, v in t.items()} for t in targets]
output = model(images, targets)  # 返回损失值（loss）和检测
# 测试
model.eval()
x = [torch.rand(3, 300, 400), torch.rand(3, 500, 400)]
predictions = model(x)

现在，让我们编写主函数，用来执行训练和测试：

from engine import train_one_epoch, evaluate
import utils


def main():
    # 若GPU可用便在GPU上训练
    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

    # 我们的数据集只有2个类别：背景和行人
    num_classes = 2
    # 使用我们的数据集和定义的transformations
    dataset = PennFudanDataset('PennFudanPed', get_transform(train=True))
    dataset_test = PennFudanDataset('PennFudanPed', get_transform(train=False))

    # 分割数据集为训练集和验证集
    indices = torch.randperm(len(dataset)).tolist()
    dataset = torch.utils.data.Subset(dataset, indices[:-50])
    dataset_test = torch.utils.data.Subset(dataset_test, indices[-50:])

    # 定义训练集和验证集的数据加载器
    data_loader = torch.utils.data.DataLoader(
        dataset, batch_size=2, shuffle=True, num_workers=4,
        collate_fn=utils.collate_fn)

    data_loader_test = torch.utils.data.DataLoader(
        dataset_test, batch_size=1, shuffle=False, num_workers=4,
        collate_fn=utils.collate_fn)

    # 使用帮助函数获得模型
    model = get_model_instance_segmentation(num_classes)

    # 将模型移动到指定设备上
    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)
    # 使用learning rate scheduler
    lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
                                                   step_size=3,
                                                   gamma=0.1)

    # 训练模型10个epochs
    num_epochs = 10

    for epoch in range(num_epochs):
        # train for one epoch, printing every 10 iterations
        train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq=10)
        # update the learning rate
        lr_scheduler.step()
        # evaluate on the test dataset
        evaluate(model, data_loader_test, device=device)

    print("That's it!")

最终的预测效果，可以使用数据集中的一张图片来验证

训练的模型预测了9个实例，让我们看下其中的几个实例：

在本文，对于实例分割模型，你已经学会了如何创建你自己在自定义数据集上的训练流程。为此，你编写了用户反悔图片和真实边界以及分割遮罩的torch.utils.data.Dataset类。为了在执行迁移学习和新数据集上，你还利用了预训练在COCO train2017上的Mask R-CNN模型。

对于更完整的示例，包含多机器、多显卡训练，查看torchvision仓库中的references/detection/train.py即可。