注意
跳转至页面底部下载完整示例代码。
TorchVision 目标检测微调教程#
创建日期：2023年12月14日 | 最后更新：2025年09月05日 | 最后验证：2024年11月05日
在本教程中，我们将基于 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 图像。target: 一个包含以下字段的字典
boxes，形状为[N, 4]的torchvision.tv_tensors.BoundingBoxes：以[x0, y0, x1, y1]格式表示的N个边界框坐标，范围从0到W以及0到H。labels，形状为[N]的整数torch.Tensor：每个边界框的标签。0始终代表背景类。image_id，整数：图像标识符。它在数据集的所有图像中必须唯一，并在评估期间使用。area，形状为[N]的浮点型torch.Tensor：边界框的面积。这在进行 COCO 指标评估时用于区分小、中、大尺寸框的指标分数。iscrowd，形状为[N]的 uint8torch.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))
<matplotlib.image.AxesImage object at 0x7f2a4bfdc8e0>
因此，每张图像都有一个对应的分割掩码，其中每种颜色对应一个不同的实例。让我们为该数据集编写一个 torch.utils.data.Dataset 类。在下面的代码中，我们将图像、边界框和掩码包装到 torchvision.tv_tensors.TVTensor 类中，以便我们可以应用 torchvision 内置的转换（新 Transforms 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)
Downloading: "https://download.pytorch.org/models/fasterrcnn_resnet50_fpn_coco-258fb6c6.pth" to /var/lib/ci-user/.cache/torch/hub/checkpoints/fasterrcnn_resnet50_fpn_coco-258fb6c6.pth
0%| | 0.00/160M [00:00<?, ?B/s]
15%|█▍ | 23.6M/160M [00:00<00:00, 246MB/s]
30%|██▉ | 47.1M/160M [00:00<00:01, 78.7MB/s]
51%|█████ | 81.4M/160M [00:00<00:00, 138MB/s]
75%|███████▍ | 120M/160M [00:00<00:00, 199MB/s]
100%|██████████| 160M/160M [00:00<00:00, 195MB/s]
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
)
Downloading: "https://download.pytorch.org/models/mobilenet_v2-7ebf99e0.pth" to /var/lib/ci-user/.cache/torch/hub/checkpoints/mobilenet_v2-7ebf99e0.pth
0%| | 0.00/13.6M [00:00<?, ?B/s]
100%|██████████| 13.6M/13.6M [00:00<00:00, 402MB/s]
用于 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):
# 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")
0
自 v0.15.0 版本起，torchvision 提供了新的 Transforms 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])
{'loss_classifier': tensor(0.1177, grad_fn=<NllLossBackward0>), 'loss_box_reg': tensor(0.0478, grad_fn=<DivBackward0>), 'loss_objectness': tensor(0.0177, grad_fn=<BinaryCrossEntropyWithLogitsBackward0>), 'loss_rpn_box_reg': tensor(0.0034, 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 等加速器上训练模型。现在让我们编写执行训练和验证的主要函数。
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]
19%|█▉ | 31.9M/170M [00:00<00:00, 334MB/s]
38%|███▊ | 63.8M/170M [00:00<00:00, 227MB/s]
57%|█████▋ | 97.0M/170M [00:00<00:00, 271MB/s]
74%|███████▍ | 126M/170M [00:00<00:00, 283MB/s]
94%|█████████▍| 160M/170M [00:00<00:00, 306MB/s]
100%|██████████| 170M/170M [00:00<00:00, 293MB/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:43 lr: 0.000090 loss: 4.2264 (4.2264) loss_classifier: 0.7407 (0.7407) loss_box_reg: 0.1434 (0.1434) loss_mask: 3.3278 (3.3278) loss_objectness: 0.0123 (0.0123) loss_rpn_box_reg: 0.0023 (0.0023) time: 0.7270 data: 0.0130 max mem: 1735
Epoch: [0] [10/60] eta: 0:00:12 lr: 0.000936 loss: 1.7446 (2.2678) loss_classifier: 0.4658 (0.4625) loss_box_reg: 0.2624 (0.2520) loss_mask: 0.9557 (1.5286) loss_objectness: 0.0158 (0.0205) loss_rpn_box_reg: 0.0029 (0.0042) time: 0.2598 data: 0.0148 max mem: 2431
Epoch: [0] [20/60] eta: 0:00:09 lr: 0.001783 loss: 0.8862 (1.5441) loss_classifier: 0.1974 (0.3176) loss_box_reg: 0.2163 (0.2526) loss_mask: 0.3872 (0.9474) loss_objectness: 0.0158 (0.0209) loss_rpn_box_reg: 0.0038 (0.0055) time: 0.2111 data: 0.0151 max mem: 2431
Epoch: [0] [30/60] eta: 0:00:06 lr: 0.002629 loss: 0.5413 (1.2130) loss_classifier: 0.1047 (0.2414) loss_box_reg: 0.1981 (0.2414) loss_mask: 0.2393 (0.7077) loss_objectness: 0.0089 (0.0171) loss_rpn_box_reg: 0.0049 (0.0054) time: 0.2089 data: 0.0148 max mem: 2468
Epoch: [0] [40/60] eta: 0:00:04 lr: 0.003476 loss: 0.4665 (1.0342) loss_classifier: 0.0595 (0.1971) loss_box_reg: 0.1941 (0.2320) loss_mask: 0.1919 (0.5850) loss_objectness: 0.0051 (0.0143) loss_rpn_box_reg: 0.0029 (0.0058) time: 0.2059 data: 0.0142 max mem: 2629
Epoch: [0] [50/60] eta: 0:00:02 lr: 0.004323 loss: 0.3742 (0.9122) loss_classifier: 0.0432 (0.1681) loss_box_reg: 0.1327 (0.2200) loss_mask: 0.1830 (0.5059) loss_objectness: 0.0032 (0.0120) loss_rpn_box_reg: 0.0044 (0.0062) time: 0.2101 data: 0.0145 max mem: 2629
Epoch: [0] [59/60] eta: 0:00:00 lr: 0.005000 loss: 0.3293 (0.8370) loss_classifier: 0.0414 (0.1492) loss_box_reg: 0.1250 (0.2068) loss_mask: 0.1881 (0.4638) loss_objectness: 0.0026 (0.0106) loss_rpn_box_reg: 0.0065 (0.0065) time: 0.2154 data: 0.0154 max mem: 2761
Epoch: [0] Total time: 0:00:13 (0.2188 s / it)
creating index...
index created!
Test: [ 0/50] eta: 0:00:04 model_time: 0.0857 (0.0857) evaluator_time: 0.0037 (0.0037) time: 0.0982 data: 0.0083 max mem: 2761
Test: [49/50] eta: 0:00:00 model_time: 0.0548 (0.0705) evaluator_time: 0.0087 (0.0108) time: 0.0807 data: 0.0105 max mem: 2761
Test: Total time: 0:00:04 (0.0935 s / it)
Averaged stats: model_time: 0.0548 (0.0705) evaluator_time: 0.0087 (0.0108)
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.543
Average Precision (AP) @[ IoU=0.50 | area= all | maxDets=100 ] = 0.988
Average Precision (AP) @[ IoU=0.75 | area= all | maxDets=100 ] = 0.480
Average Precision (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = -1.000
Average Precision (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.504
Average Precision (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.547
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 1 ] = 0.201
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 10 ] = 0.640
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.640
Average Recall (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = -1.000
Average Recall (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.708
Average Recall (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.633
IoU metric: segm
Average Precision (AP) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.675
Average Precision (AP) @[ IoU=0.50 | area= all | maxDets=100 ] = 0.988
Average Precision (AP) @[ IoU=0.75 | area= all | maxDets=100 ] = 0.899
Average Precision (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = -1.000
Average Precision (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.470
Average Precision (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.686
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 1 ] = 0.246
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 10 ] = 0.726
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.732
Average Recall (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = -1.000
Average Recall (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.715
Average Recall (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.733
Epoch: [1] [ 0/60] eta: 0:00:13 lr: 0.005000 loss: 0.3812 (0.3812) loss_classifier: 0.0428 (0.0428) loss_box_reg: 0.1530 (0.1530) loss_mask: 0.1784 (0.1784) loss_objectness: 0.0003 (0.0003) loss_rpn_box_reg: 0.0067 (0.0067) time: 0.2214 data: 0.0207 max mem: 2761
Epoch: [1] [10/60] eta: 0:00:10 lr: 0.005000 loss: 0.3633 (0.3207) loss_classifier: 0.0399 (0.0389) loss_box_reg: 0.1299 (0.1189) loss_mask: 0.1423 (0.1550) loss_objectness: 0.0011 (0.0022) loss_rpn_box_reg: 0.0055 (0.0057) time: 0.2122 data: 0.0167 max mem: 2761
Epoch: [1] [20/60] eta: 0:00:08 lr: 0.005000 loss: 0.2990 (0.3257) loss_classifier: 0.0368 (0.0406) loss_box_reg: 0.1124 (0.1161) loss_mask: 0.1544 (0.1611) loss_objectness: 0.0011 (0.0019) loss_rpn_box_reg: 0.0046 (0.0059) time: 0.2097 data: 0.0157 max mem: 2761
Epoch: [1] [30/60] eta: 0:00:06 lr: 0.005000 loss: 0.2805 (0.3061) loss_classifier: 0.0397 (0.0400) loss_box_reg: 0.0868 (0.1063) loss_mask: 0.1467 (0.1525) loss_objectness: 0.0011 (0.0020) loss_rpn_box_reg: 0.0041 (0.0053) time: 0.2080 data: 0.0146 max mem: 2761
Epoch: [1] [40/60] eta: 0:00:04 lr: 0.005000 loss: 0.2565 (0.2929) loss_classifier: 0.0341 (0.0371) loss_box_reg: 0.0703 (0.0953) loss_mask: 0.1218 (0.1540) loss_objectness: 0.0005 (0.0017) loss_rpn_box_reg: 0.0034 (0.0048) time: 0.2057 data: 0.0139 max mem: 2761
Epoch: [1] [50/60] eta: 0:00:02 lr: 0.005000 loss: 0.2341 (0.2850) loss_classifier: 0.0295 (0.0365) loss_box_reg: 0.0598 (0.0886) loss_mask: 0.1382 (0.1538) loss_objectness: 0.0003 (0.0015) loss_rpn_box_reg: 0.0034 (0.0045) time: 0.1970 data: 0.0140 max mem: 2761
Epoch: [1] [59/60] eta: 0:00:00 lr: 0.005000 loss: 0.2341 (0.2799) loss_classifier: 0.0334 (0.0361) loss_box_reg: 0.0598 (0.0857) loss_mask: 0.1382 (0.1523) loss_objectness: 0.0005 (0.0014) loss_rpn_box_reg: 0.0037 (0.0044) time: 0.1989 data: 0.0142 max mem: 2761
Epoch: [1] Total time: 0:00:12 (0.2050 s / it)
creating index...
index created!
Test: [ 0/50] eta: 0:00:02 model_time: 0.0407 (0.0407) evaluator_time: 0.0019 (0.0019) time: 0.0512 data: 0.0083 max mem: 2761
Test: [49/50] eta: 0:00:00 model_time: 0.0411 (0.0433) evaluator_time: 0.0041 (0.0055) time: 0.0579 data: 0.0104 max mem: 2761
Test: Total time: 0:00:03 (0.0602 s / it)
Averaged stats: model_time: 0.0411 (0.0433) evaluator_time: 0.0041 (0.0055)
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.716
Average Precision (AP) @[ IoU=0.50 | area= all | maxDets=100 ] = 0.989
Average Precision (AP) @[ IoU=0.75 | area= all | maxDets=100 ] = 0.867
Average Precision (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = -1.000
Average Precision (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.510
Average Precision (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.726
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 1 ] = 0.262
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 10 ] = 0.767
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.767
Average Recall (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = -1.000
Average Recall (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.731
Average Recall (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.770
IoU metric: segm
Average Precision (AP) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.704
Average Precision (AP) @[ IoU=0.50 | area= all | maxDets=100 ] = 0.976
Average Precision (AP) @[ IoU=0.75 | area= all | maxDets=100 ] = 0.898
Average Precision (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = -1.000
Average Precision (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.379
Average Precision (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.711
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 1 ] = 0.260
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 10 ] = 0.746
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.746
Average Recall (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = -1.000
Average Recall (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.762
Average Recall (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.745
That's it!
经过一个 epoch 的训练，我们获得了 COCO 风格的 mAP > 50，以及 65 的 mask mAP。
那么预测结果看起来如何呢？让我们取数据集中的一张图像来验证一下。
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))
<matplotlib.image.AxesImage object at 0x7f2a8d75af50>
结果看起来不错！
总结#
在本教程中，您学习了如何为自定义数据集上的目标检测模型创建自己的训练流水线。为此，您编写了一个返回图像、真实边界框和分割掩码的 torch.utils.data.Dataset 类。您还利用了在 COCO train2017 上预训练的 Mask R-CNN 模型，在该新数据集上执行了迁移学习。
如需查看更完整的示例（包含多机器/多 GPU 训练），请查看 torchvision 仓库中的 references/detection/train.py。
脚本总运行时间： (0 分 48.336 秒)