TorchVision 对象检测微调教程

Created On: Dec 14, 2023 | Last Updated: Jun 11, 2024 | Last Verified: Nov 05, 2024

在本教程中，我们将对预训练的 Mask R-CNN 模型在 Penn-Fudan 行人检测和分割数据库 上进行微调。该数据库包含170张图片和345个行人实例，我们将使用它来演示如何利用 torchvision 的新功能来训练针对自定义数据集的对象检测和实例分割模型。

备注

本教程仅适用于 torchvision 版本 >=0.16 或 nightly。如果您使用的是 torchvision<=0.15，请改为按照 这篇教程。

定义数据集

对象检测、实例分割和人物关键点检测的参考脚本允许轻松添加新的自定义数据集。数据集应继承标准类 torch.utils.data.Dataset，并实现 __len__ 和 __getitem__。

我们要求数据集 __getitem__ 返回一个元组:

• image: 形状为 [3, H, W] 的 torchvision.tv_tensors.Image`（纯张量）或大小为 ``(H, W)` 的 PIL 图像

• target: 包含以下字段的字典

  • boxes, 形状为 [N, 4] 的 torchvision.tv_tensors.BoundingBoxes：以 [x0, y0, x1, y1] 格式定义的 N 个边界框的坐标，范围从 0 到 W 和 0 到 H

  • labels, 形状为 [N] 的整数 torch.Tensor：每个边界框的标签。标签 0 始终代表背景类。

  • image_id, int: 图像标识符。在数据集中所有图像中应该是唯一的，用于评估期间

  • area, 形状为 [N] 的 float torch.Tensor: 边界框的面积。在评估 COCO 指标时，用于将指标分数划分为小型、中型和大型框。

  • iscrowd, 形状为 [N] 的 uint8 torch.Tensor: 实例其中 iscrowd=True 将在评估期间被忽略。

  • (可选) masks, 形状为 [N, H, W] 的 torchvision.tv_tensors.Mask: 每个对象的分割掩码

如果您的数据集符合上述要求，则可以支持参考脚本中的训练和评估代码。评估代码将使用 pycocotools 库中的脚本，该库可通过 pip install pycocotools 安装。

备注

对于 Windows，请从 gautamchitnis 安装 pycocotools ，命令如下

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

关于 labels 的一点注意事项。模型将类 0 视为背景。如果您的数据集中不包含背景类，则您的 labels 中不应有 0。例如，假设您只有两个类，猫 和 狗，您可以定义 1``（而不是 ``0）来表示 猫，定义 2 来表示 狗。因此，假如有一个图片同时包含这两个类，那么您的 labels 张量应为 [1, 2]。

此外，如果您希望在训练期间使用纵横比分组（以使每个批次仅包含具有相似纵横比的图像），建议实现一个 get_height_and_width 方法以返回图像的高度和宽度。如果未提供此方法，我们通过 __getitem__ 查询数据集的所有元素，这会加载图像到内存，速度比提供自定义方法慢。

为 PennFudan 编写自定义数据集

让我们为 PennFudan 数据集编写一个数据集。首先下载数据集并提取 ZIP 文件:

wget https://www.cis.upenn.edu/~jshi/ped_html/PennFudanPed.zip -P data
cd data && unzip PennFudanPed.zip

我们有以下文件夹结构:

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

以下是一个图片和分割掩码的示例配对

import matplotlib.pyplot as plt
from torchvision.io import read_image


image = read_image("data/PennFudanPed/PNGImages/FudanPed00046.png")
mask = read_image("data/PennFudanPed/PedMasks/FudanPed00046_mask.png")

plt.figure(figsize=(16, 8))
plt.subplot(121)
plt.title("Image")
plt.imshow(image.permute(1, 2, 0))
plt.subplot(122)
plt.title("Mask")
plt.imshow(mask.permute(1, 2, 0))

因此，每个图片都有一个对应的分割掩码，其中每种颜色对应一个实例。让我们为这个数据集编写一个 torch.utils.data.Dataset 类。在以下代码中，我们将图片、边界框和掩码包装到 torchvision.tv_tensors.TVTensor 类中，以便能为给定的对象检测和分割任务应用 torchvision 内置变换（新变换 API）。具体来说，图片张量将由 torchvision.tv_tensors.Image 包装，边界框由 torchvision.tv_tensors.BoundingBoxes 包装，掩码由 torchvision.tv_tensors.Mask 包装。由于 torchvision.tv_tensors.TVTensor 是 torch.Tensor 的子类，包装的对象也是张量并继承了普通 torch.Tensor 的 API。关于 torchvision tv_tensors 的更多信息，请参阅 这份文档。

import os
import torch

from torchvision.io import read_image
from torchvision.ops.boxes import masks_to_boxes
from torchvision import tv_tensors
from torchvision.transforms.v2 import functional as F


class PennFudanDataset(torch.utils.data.Dataset):
    def __init__(self, root, transforms):
        self.root = root
        self.transforms = transforms
        # load all image files, sorting them to
        # ensure that they are aligned
        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):
        # load images and masks
        img_path = os.path.join(self.root, "PNGImages", self.imgs[idx])
        mask_path = os.path.join(self.root, "PedMasks", self.masks[idx])
        img = read_image(img_path)
        mask = read_image(mask_path)
        # instances are encoded as different colors
        obj_ids = torch.unique(mask)
        # first id is the background, so remove it
        obj_ids = obj_ids[1:]
        num_objs = len(obj_ids)

        # split the color-encoded mask into a set
        # of binary masks
        masks = (mask == obj_ids[:, None, None]).to(dtype=torch.uint8)

        # get bounding box coordinates for each mask
        boxes = masks_to_boxes(masks)

        # there is only one class
        labels = torch.ones((num_objs,), dtype=torch.int64)

        image_id = idx
        area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])
        # suppose all instances are not crowd
        iscrowd = torch.zeros((num_objs,), dtype=torch.int64)

        # Wrap sample and targets into torchvision tv_tensors:
        img = tv_tensors.Image(img)

        target = {}
        target["boxes"] = tv_tensors.BoundingBoxes(boxes, format="XYXY", canvas_size=F.get_size(img))
        target["masks"] = tv_tensors.Mask(masks)
        target["labels"] = labels
        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)

数据集就此定义完成。现在让我们定义一个能够在这个数据集上进行预测的模型。

定义您的模型

在本教程中，我们将使用 Mask R-CNN，它基于 Faster R-CNN。 Faster R-CNN 是一种模型，可预测图片中潜在对象的边界框和类别分数。

Mask R-CNN 在 Faster R-CNN 的基础上增加了一个分支，还预测每个实例的分割掩码。

TorchVision 模型库中有两种常见的情况，一个是我们希望从一个预训练模型开始，只微调最后的一层。另一个是我们希望替换模型的主干，使用不同的主干（例如，以获得更快的预测）。

接下来我们将演示如何实现上述两种情况。

1 - 从预训练模型进行微调

假设您希望从一个在 COCO 上预训练的模型开始，并只对其进行微调以适配您的特定类别。可以按以下方式实现：

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

# load a model pre-trained on COCO
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(weights="DEFAULT")

# replace the classifier with a new one, that has
# num_classes which is user-defined
num_classes = 2  # 1 class (person) + background
# get number of input features for the classifier
in_features = model.roi_heads.box_predictor.cls_score.in_features
# replace the pre-trained head with a new one
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

# load a pre-trained model for classification and return
# only the features
backbone = torchvision.models.mobilenet_v2(weights="DEFAULT").features
# ``FasterRCNN`` needs to know the number of
# output channels in a backbone. For mobilenet_v2, it's 1280
# so we need to add it here
backbone.out_channels = 1280

# let's make the RPN generate 5 x 3 anchors per spatial
# location, with 5 different sizes and 3 different aspect
# ratios. We have a Tuple[Tuple[int]] because each feature
# map could potentially have different sizes and
# aspect ratios
anchor_generator = AnchorGenerator(
    sizes=((32, 64, 128, 256, 512),),
    aspect_ratios=((0.5, 1.0, 2.0),)
)

# let's define what are the feature maps that we will
# use to perform the region of interest cropping, as well as
# the size of the crop after rescaling.
# if your backbone returns a Tensor, featmap_names is expected to
# be [0]. More generally, the backbone should return an
# ``OrderedDict[Tensor]``, and in ``featmap_names`` you can choose which
# feature maps to use.
roi_pooler = torchvision.ops.MultiScaleRoIAlign(
    featmap_names=['0'],
    output_size=7,
    sampling_ratio=2
)

# put the pieces together inside a Faster-RCNN model
model = FasterRCNN(
    backbone,
    num_classes=2,
    rpn_anchor_generator=anchor_generator,
    box_roi_pool=roi_pooler
)

为 PennFudan 数据集编写对象检测和实例分割模型

在我们的案例中，由于数据集非常小，我们希望从预训练模型开始微调，因此我们将采用第一种方法。

在这里，我们希望还能够计算实例分割掩码，因此将使用 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):
    # load an instance segmentation model pre-trained on COCO
    model = torchvision.models.detection.maskrcnn_resnet50_fpn(weights="DEFAULT")

    # get number of input features for the classifier
    in_features = model.roi_heads.box_predictor.cls_score.in_features
    # replace the pre-trained head with a new one
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

    # now get the number of input features for the mask classifier
    in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
    hidden_layer = 256
    # and replace the mask predictor with a new one
    model.roi_heads.mask_predictor = MaskRCNNPredictor(
        in_features_mask,
        hidden_layer,
        num_classes
    )

    return model

完成后，model 就可以准备在您的自定义数据集上进行训练和评估了。

将所有部分整合到一起

在 references/detection/ 文件夹中，我们有一些辅助函数可以简化训练和评估检测模型的过程。在这里我们将使用 references/detection/engine.py 和 references/detection/utils.py。您只需将 references/detection 文件夹中的所有内容都下载到您的文件夹中并在此使用。在 Linux 上如果有 wget，可以使用以下命令下载它们：

os.system("wget https://raw.githubusercontent.com/pytorch/vision/main/references/detection/engine.py")
os.system("wget https://raw.githubusercontent.com/pytorch/vision/main/references/detection/utils.py")
os.system("wget https://raw.githubusercontent.com/pytorch/vision/main/references/detection/coco_utils.py")
os.system("wget https://raw.githubusercontent.com/pytorch/vision/main/references/detection/coco_eval.py")
os.system("wget https://raw.githubusercontent.com/pytorch/vision/main/references/detection/transforms.py")

自版本 v0.15.0 起，torchvision 提供了 新变换 API ，可以轻松编写用于对象检测和分割任务的数据增强管道。

让我们为数据增强/变换编写一些辅助函数：

from torchvision.transforms import v2 as T


def get_transform(train):
    transforms = []
    if train:
        transforms.append(T.RandomHorizontalFlip(0.5))
    transforms.append(T.ToDtype(torch.float, scale=True))
    transforms.append(T.ToPureTensor())
    return T.Compose(transforms)

测试 forward() 方法（可选）

在遍历数据集之前，最好查看模型在训练和推断时对样本数据的预期。

import utils

model = torchvision.models.detection.fasterrcnn_resnet50_fpn(weights="DEFAULT")
dataset = PennFudanDataset('data/PennFudanPed', get_transform(train=True))
data_loader = torch.utils.data.DataLoader(
    dataset,
    batch_size=2,
    shuffle=True,
    collate_fn=utils.collate_fn
)

# For Training
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)  # Returns losses and detections
print(output)

# For inference
model.eval()
x = [torch.rand(3, 300, 400), torch.rand(3, 500, 400)]
predictions = model(x)  # Returns predictions
print(predictions[0])

现在让我们编写一个主函数来执行训练和验证：

from engine import train_one_epoch, evaluate

# train on the GPU or on the CPU, if a GPU is not available
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

# our dataset has two classes only - background and person
num_classes = 2
# use our dataset and defined transformations
dataset = PennFudanDataset('data/PennFudanPed', get_transform(train=True))
dataset_test = PennFudanDataset('data/PennFudanPed', get_transform(train=False))

# split the dataset in train and test set
indices = torch.randperm(len(dataset)).tolist()
dataset = torch.utils.data.Subset(dataset, indices[:-50])
dataset_test = torch.utils.data.Subset(dataset_test, indices[-50:])

# define training and validation data loaders
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
)

# get the model using our helper function
model = get_model_instance_segmentation(num_classes)

# move model to the right device
model.to(device)

# construct an optimizer
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
)

# and a learning rate scheduler
lr_scheduler = torch.optim.lr_scheduler.StepLR(
    optimizer,
    step_size=3,
    gamma=0.1
)

# let's train it just for 2 epochs
num_epochs = 2

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!")

经过一个训练周期后，我们获得了 COCO 风格的 mAP > 50 和分割掩码 mAP 为 65。

但预测结果是什么样的呢？让我们选一个数据集中的图片并查看一下

import matplotlib.pyplot as plt

from torchvision.utils import draw_bounding_boxes, draw_segmentation_masks


image = read_image("data/PennFudanPed/PNGImages/FudanPed00046.png")
eval_transform = get_transform(train=False)

model.eval()
with torch.no_grad():
    x = eval_transform(image)
    # convert RGBA -> RGB and move to device
    x = x[:3, ...].to(device)
    predictions = model([x, ])
    pred = predictions[0]


image = (255.0 * (image - image.min()) / (image.max() - image.min())).to(torch.uint8)
image = image[:3, ...]
pred_labels = [f"pedestrian: {score:.3f}" for label, score in zip(pred["labels"], pred["scores"])]
pred_boxes = pred["boxes"].long()
output_image = draw_bounding_boxes(image, pred_boxes, pred_labels, colors="red")

masks = (pred["masks"] > 0.7).squeeze(1)
output_image = draw_segmentation_masks(output_image, masks, alpha=0.5, colors="blue")


plt.figure(figsize=(12, 12))
plt.imshow(output_image.permute(1, 2, 0))

结果看起来不错！

总结

在本教程中，您已经学习如何为自定义数据集上的对象检测模型创建自己的训练管道。为此，您编写了一个 torch.utils.data.Dataset 类来返回图片以及真实边界框和分割掩码。您还利用了一个在 COCO train2017 上预训练的 Mask R-CNN 模型，以便在这个新数据集上进行迁移学习。

若需更完整的示例，包括多机/多GPU训练，请查看 torchvision 仓库中提供的 references/detection/train.py。