TorchVision 目标检测模型微调教程

本教程在 Penn-Fudan 行人检测与分割数据集上微调预训练的 Mask R-CNN 模型。数据集包含 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 图像。PIL 的 Image.size 实际按 (W, H) 返回，编写代码时应区分这一顺序。
target：包含以下字段的字典。

字段	类型、形状与含义
`boxes`	`torchvision.tv_tensors.BoundingBoxes`，形状 `[N, 4]`。N 个边界框按 `[x0, y0, x1, y1]` 表示，坐标分别处于 0～W 和 0～H。
`labels`	整数 `torch.Tensor`，形状 `[N]`，每个框的类别；0 始终代表背景。
`image_id`	整数图像标识符，在整个数据集中唯一，评估时使用。
`area`	浮点 `torch.Tensor`，形状 `[N]`，边界框面积；COCO 评估依据它区分小、中、大目标的指标。
`iscrowd`	原文字段说明为 uint8 `torch.Tensor`，形状 `[N]`；`iscrowd=True` 的实例在评估中忽略。下方原始示例实际使用 `torch.int64`，两处表述存在差异。
`masks`（可选）	`torchvision.tv_tensors.Mask`，形状 `[N, H, W]`，每个目标的分割掩码。

符合以上要求的数据集，可以直接配合参考脚本的训练和评估代码。评估需要 pycocotools，可用 pip install pycocotools 安装。原教程针对 Windows 给出的安装方式是使用 gautamchitnis 的 cocoapi 分支：

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

模型把标签 0 视为背景。如果数据集没有背景类，labels 中就不应出现 0。例如只有猫和狗两类时，可用 1 表示猫、2 表示狗。一张图同时包含两类目标时，标签张量应类似 [1, 2]。

如果希望训练时按宽高比分组，使每个批次中的图像宽高比接近，建议再实现 get_height_and_width，返回图像高度和宽度。否则脚本会通过 __getitem__ 遍历数据集，把图像读入内存，速度较慢。

为 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.Image、BoundingBoxes 和 Mask，以便在目标检测与分割任务中使用 torchvision 内置的新变换 API。

torchvision.tv_tensors.TVTensor 是 torch.Tensor 的子类，因此包装后的对象仍是张量，继承普通张量 API。更多内容见 TVTensors 文档。

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)

数据集到此定义完成，接下来建立能够对其预测的模型。

定义模型

我们使用建立在 Faster R-CNN 基础上的 Mask R-CNN。Faster R-CNN 为图像中的候选目标预测边界框和类别得分。

Mask 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/ 中包含简化检测模型训练和评估的辅助函数。这里使用 engine.py 和 utils.py，同时需要该目录的其他依赖文件。将它们下载到工作目录后即可导入。在安装了 wget 的 Linux 环境中，原教程使用以下代码下载：

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

从 torchvision 0.15.0 开始，新变换 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])

原网页的 forward() 示例输出如下：

{'loss_classifier': tensor(0.1195, grad_fn=<NllLossBackward0>), 'loss_box_reg': tensor(0.0222, grad_fn=<DivBackward0>), 'loss_objectness': tensor(0.0159, grad_fn=<BinaryCrossEntropyWithLogitsBackward0>), 'loss_rpn_box_reg': tensor(0.0009, grad_fn=<DivBackward0>)}
{'boxes': tensor([], size=(0, 4), grad_fn=<StackBackward0>), 'labels': tensor([], dtype=torch.int64), 'scores': tensor([], grad_fn=<IndexBackward0>)}

训练和验证

我们希望使用 CUDA、MPS、MTIA、XPU 等加速器训练；没有可用加速器时使用 CPU。下面组合训练和验证过程。这里的 torch.accelerator 要求所安装的 PyTorch 提供该 API，不能仅凭 torchvision 版本下限判断兼容性。

from engine import train_one_epoch, evaluate

# train on the accelerator or on the CPU, if an accelerator is not available
device = torch.accelerator.current_accelerator() if torch.accelerator.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!")

原网页的训练输出如下，包含下载进度、弃用警告和两轮评估；这些是原站构建环境的记录，并非本译稿的运行结果：

Downloading: "https://download.pytorch.org/models/maskrcnn_resnet50_fpn_coco-bf2d0c1e.pth" to /var/lib/ci-user/.cache/torch/hub/checkpoints/maskrcnn_resnet50_fpn_coco-bf2d0c1e.pth

  0%|          | 0.00/170M [00:00<?, ?B/s]
 25%|██▍       | 42.1M/170M [00:00<00:00, 441MB/s]
 50%|█████     | 85.6M/170M [00:00<00:00, 449MB/s]
 76%|███████▌  | 129M/170M [00:00<00:00, 452MB/s]
100%|██████████| 170M/170M [00:00<00:00, 452MB/s]
/var/lib/workspace/intermediate_source/engine.py:30: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.
  with torch.cuda.amp.autocast(enabled=scaler is not None):
Epoch: [0]  [ 0/60]  eta: 0:00:51  lr: 0.000090  loss: 5.0142 (5.0142)  loss_classifier: 0.6467 (0.6467)  loss_box_reg: 0.3455 (0.3455)  loss_mask: 3.9954 (3.9954)  loss_objectness: 0.0207 (0.0207)  loss_rpn_box_reg: 0.0058 (0.0058)  time: 0.8572  data: 0.0140  max mem: 2144
Epoch: [0]  [10/60]  eta: 0:00:13  lr: 0.000936  loss: 1.7787 (2.8318)  loss_classifier: 0.4449 (0.4134)  loss_box_reg: 0.3037 (0.2730)  loss_mask: 0.9447 (2.1163)  loss_objectness: 0.0237 (0.0255)  loss_rpn_box_reg: 0.0033 (0.0037)  time: 0.2663  data: 0.0166  max mem: 2432
Epoch: [0]  [20/60]  eta: 0:00:09  lr: 0.001783  loss: 1.0360 (1.8413)  loss_classifier: 0.2012 (0.2987)  loss_box_reg: 0.2441 (0.2786)  loss_mask: 0.4199 (1.2388)  loss_objectness: 0.0153 (0.0203)  loss_rpn_box_reg: 0.0041 (0.0050)  time: 0.2078  data: 0.0169  max mem: 2551
Epoch: [0]  [30/60]  eta: 0:00:06  lr: 0.002629  loss: 0.5565 (1.4080)  loss_classifier: 0.1181 (0.2296)  loss_box_reg: 0.1978 (0.2499)  loss_mask: 0.2211 (0.9055)  loss_objectness: 0.0097 (0.0176)  loss_rpn_box_reg: 0.0054 (0.0054)  time: 0.2025  data: 0.0161  max mem: 2551
Epoch: [0]  [40/60]  eta: 0:00:04  lr: 0.003476  loss: 0.4310 (1.1655)  loss_classifier: 0.0596 (0.1856)  loss_box_reg: 0.1598 (0.2301)  loss_mask: 0.1767 (0.7295)  loss_objectness: 0.0071 (0.0152)  loss_rpn_box_reg: 0.0041 (0.0051)  time: 0.1940  data: 0.0149  max mem: 2551
Epoch: [0]  [50/60]  eta: 0:00:02  lr: 0.004323  loss: 0.3635 (1.0068)  loss_classifier: 0.0443 (0.1572)  loss_box_reg: 0.1468 (0.2100)  loss_mask: 0.1750 (0.6217)  loss_objectness: 0.0037 (0.0129)  loss_rpn_box_reg: 0.0036 (0.0051)  time: 0.1926  data: 0.0141  max mem: 2551
Epoch: [0]  [59/60]  eta: 0:00:00  lr: 0.005000  loss: 0.3605 (0.9150)  loss_classifier: 0.0436 (0.1409)  loss_box_reg: 0.1428 (0.1995)  loss_mask: 0.1683 (0.5578)  loss_objectness: 0.0030 (0.0115)  loss_rpn_box_reg: 0.0053 (0.0054)  time: 0.2062  data: 0.0149  max mem: 2982
Epoch: [0] Total time: 0:00:12 (0.2136 s / it)
creating index...
index created!
Test:  [ 0/50]  eta: 0:00:04  model_time: 0.0602 (0.0602)  evaluator_time: 0.0186 (0.0186)  time: 0.0964  data: 0.0169  max mem: 2982
Test:  [49/50]  eta: 0:00:00  model_time: 0.0437 (0.0460)  evaluator_time: 0.0061 (0.0071)  time: 0.0643  data: 0.0108  max mem: 2982
Test: Total time: 0:00:03 (0.0644 s / it)
Averaged stats: model_time: 0.0437 (0.0460)  evaluator_time: 0.0061 (0.0071)
Accumulating evaluation results...
DONE (t=0.01s).
Accumulating evaluation results...
DONE (t=0.01s).
IoU metric: bbox
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.712
 Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.992
 Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.882
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.625
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.414
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.721
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.295
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.759
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.759
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.650
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.700
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.764
IoU metric: segm
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.715
 Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.992
 Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.893
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.422
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.338
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.727
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.284
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.751
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.753
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.700
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.714
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.756
Epoch: [1]  [ 0/60]  eta: 0:00:13  lr: 0.005000  loss: 0.2362 (0.2362)  loss_classifier: 0.0231 (0.0231)  loss_box_reg: 0.0745 (0.0745)  loss_mask: 0.1294 (0.1294)  loss_objectness: 0.0056 (0.0056)  loss_rpn_box_reg: 0.0036 (0.0036)  time: 0.2181  data: 0.0172  max mem: 2982
Epoch: [1]  [10/60]  eta: 0:00:10  lr: 0.005000  loss: 0.3038 (0.2900)  loss_classifier: 0.0397 (0.0395)  loss_box_reg: 0.1018 (0.0966)  loss_mask: 0.1396 (0.1459)  loss_objectness: 0.0013 (0.0020)  loss_rpn_box_reg: 0.0050 (0.0062)  time: 0.2040  data: 0.0160  max mem: 2982
Epoch: [1]  [20/60]  eta: 0:00:08  lr: 0.005000  loss: 0.2842 (0.2925)  loss_classifier: 0.0330 (0.0383)  loss_box_reg: 0.0894 (0.0898)  loss_mask: 0.1473 (0.1569)  loss_objectness: 0.0010 (0.0022)  loss_rpn_box_reg: 0.0032 (0.0054)  time: 0.1993  data: 0.0157  max mem: 2982
Epoch: [1]  [30/60]  eta: 0:00:05  lr: 0.005000  loss: 0.2823 (0.2911)  loss_classifier: 0.0392 (0.0388)  loss_box_reg: 0.0759 (0.0899)  loss_mask: 0.1533 (0.1555)  loss_objectness: 0.0010 (0.0019)  loss_rpn_box_reg: 0.0040 (0.0050)  time: 0.1968  data: 0.0154  max mem: 2982
Epoch: [1]  [40/60]  eta: 0:00:03  lr: 0.005000  loss: 0.2657 (0.2806)  loss_classifier: 0.0281 (0.0364)  loss_box_reg: 0.0728 (0.0835)  loss_mask: 0.1364 (0.1542)  loss_objectness: 0.0008 (0.0016)  loss_rpn_box_reg: 0.0043 (0.0050)  time: 0.1997  data: 0.0147  max mem: 2982
Epoch: [1]  [50/60]  eta: 0:00:02  lr: 0.005000  loss: 0.2483 (0.2832)  loss_classifier: 0.0262 (0.0366)  loss_box_reg: 0.0648 (0.0850)  loss_mask: 0.1364 (0.1551)  loss_objectness: 0.0008 (0.0015)  loss_rpn_box_reg: 0.0050 (0.0051)  time: 0.2070  data: 0.0154  max mem: 2982
Epoch: [1]  [59/60]  eta: 0:00:00  lr: 0.005000  loss: 0.2748 (0.2792)  loss_classifier: 0.0334 (0.0362)  loss_box_reg: 0.0762 (0.0823)  loss_mask: 0.1364 (0.1540)  loss_objectness: 0.0007 (0.0014)  loss_rpn_box_reg: 0.0060 (0.0053)  time: 0.2121  data: 0.0166  max mem: 2982
Epoch: [1] Total time: 0:00:12 (0.2034 s / it)
creating index...
index created!
Test:  [ 0/50]  eta: 0:00:04  model_time: 0.0466 (0.0466)  evaluator_time: 0.0164 (0.0164)  time: 0.0805  data: 0.0170  max mem: 2982
Test:  [49/50]  eta: 0:00:00  model_time: 0.0399 (0.0408)  evaluator_time: 0.0031 (0.0039)  time: 0.0561  data: 0.0108  max mem: 2982
Test: Total time: 0:00:02 (0.0557 s / it)
Averaged stats: model_time: 0.0399 (0.0408)  evaluator_time: 0.0031 (0.0039)
Accumulating evaluation results...
DONE (t=0.01s).
Accumulating evaluation results...
DONE (t=0.01s).
IoU metric: bbox
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.781
 Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.993
 Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.939
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.384
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.578
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.790
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.319
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.818
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.818
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.650
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.786
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.822
IoU metric: segm
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.723
 Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.992
 Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.883
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.359
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.529
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.731
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.292
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.764
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.765
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.550
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.743
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.770
That's it!

原教程报告：训练一个 epoch 后，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))

原教程展示的结果符合预期。原网页记录脚本总耗时为 44.996 秒。

总结

本教程介绍了如何为自定义数据集建立目标检测训练流程：编写 torch.utils.data.Dataset 返回图像、真实边界框与分割掩码，再利用在 COCO train2017 上预训练的 Mask R-CNN 进行迁移学习。

如需包含多机、多 GPU 训练的更完整示例，请参阅 torchvision 仓库中的 references/detection/train.py。

原文：TorchVision Object Detection Finetuning Tutorial，PyTorch 官方教程。版权与代码许可归 PyTorch 教程、TorchVision 及其相应权利人；Penn-Fudan 数据集遵循其原始发布条件。

教程仓库使用 BSD 3-Clause License。保留原许可证如下：

BSD 3-Clause License

Copyright (c) 2017-2022, Pytorch contributors
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

* Redistributions of source code must retain the above copyright notice, this
  list of conditions and the following disclaimer.

* Redistributions in binary form must reproduce the above copyright notice,
  this list of conditions and the following disclaimer in the documentation
  and/or other materials provided with the distribution.

* Neither the name of the copyright holder nor the names of its
  contributors may be used to endorse or promote products derived from
  this software without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

文章版权归作者所有，未经允许请勿转载。

THE END