注意
跳转至页面底部下载完整示例代码。
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，或者纯张量（tensor），或者尺寸为(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]的浮点型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]。
此外，如果您想在训练过程中使用长宽比分组（即每个 batch 只包含长宽比相似的图像），建议同时实现一个 让我们为 PennFudan 数据集编写一个数据集类。首先，下载数据集并解压 zip 文件。 我们有如下文件夹结构： 这是一个图像和分割掩码对的示例： 每张图像都有一个对应的分割掩码，其中每种颜色对应一个不同的实例。让我们为该数据集编写一个 数据集部分就到这里。现在让我们定义一个可以在此数据集上进行预测的模型。get_height_and_width 方法，返回图像的高度和宽度。如果未提供此方法，我们将通过 __getitem__ 查询数据集的所有元素，这将导致图像被加载到内存中，比提供自定义方法要慢。
为 PennFudan 编写自定义数据集#
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 0x7f8f999ce6e0>
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 模型库中的可用模型时，有两种常见情况。第一种是我们想从预训练模型开始，仅微调最后一层；第二种是我们想用不同的主干网络（backbone）替换现有的（例如为了更快的预测速度）。
让我们看看在接下来的章节中如何实现这些。
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]
24%|██▍ | 38.5M/160M [00:00<00:00, 403MB/s]
48%|████▊ | 77.0M/160M [00:00<00:00, 390MB/s]
72%|███████▏ | 116M/160M [00:00<00:00, 395MB/s]
97%|█████████▋| 156M/160M [00:00<00:00, 404MB/s]
100%|██████████| 160M/160M [00:00<00:00, 401MB/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, 409MB/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.1935, grad_fn=<NllLossBackward0>), 'loss_box_reg': tensor(0.0858, grad_fn=<DivBackward0>), 'loss_objectness': tensor(0.0259, grad_fn=<BinaryCrossEntropyWithLogitsBackward0>), 'loss_rpn_box_reg': tensor(0.0060, 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]
24%|██▍ | 40.9M/170M [00:00<00:00, 428MB/s]
49%|████▉ | 83.5M/170M [00:00<00:00, 439MB/s]
75%|███████▍ | 127M/170M [00:00<00:00, 446MB/s]
100%|██████████| 170M/170M [00:00<00:00, 447MB/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:45 lr: 0.000090 loss: 4.4292 (4.4292) loss_classifier: 0.6091 (0.6091) loss_box_reg: 0.2716 (0.2716) loss_mask: 3.5082 (3.5082) loss_objectness: 0.0389 (0.0389) loss_rpn_box_reg: 0.0014 (0.0014) time: 0.7543 data: 0.0218 max mem: 1705
Epoch: [0] [10/60] eta: 0:00:13 lr: 0.000936 loss: 1.7559 (2.3560) loss_classifier: 0.4656 (0.4174) loss_box_reg: 0.2727 (0.3042) loss_mask: 1.0085 (1.6076) loss_objectness: 0.0205 (0.0215) loss_rpn_box_reg: 0.0041 (0.0053) time: 0.2766 data: 0.0177 max mem: 2512
Epoch: [0] [20/60] eta: 0:00:09 lr: 0.001783 loss: 0.9175 (1.5340) loss_classifier: 0.1842 (0.2892) loss_box_reg: 0.2396 (0.2591) loss_mask: 0.3673 (0.9591) loss_objectness: 0.0177 (0.0210) loss_rpn_box_reg: 0.0037 (0.0056) time: 0.2202 data: 0.0162 max mem: 2542
Epoch: [0] [30/60] eta: 0:00:07 lr: 0.002629 loss: 0.4917 (1.1908) loss_classifier: 0.1004 (0.2203) loss_box_reg: 0.1724 (0.2359) loss_mask: 0.1966 (0.7093) loss_objectness: 0.0115 (0.0184) loss_rpn_box_reg: 0.0045 (0.0069) time: 0.2129 data: 0.0149 max mem: 2542
Epoch: [0] [40/60] eta: 0:00:04 lr: 0.003476 loss: 0.4547 (1.0100) loss_classifier: 0.0631 (0.1820) loss_box_reg: 0.1763 (0.2262) loss_mask: 0.1659 (0.5799) loss_objectness: 0.0075 (0.0149) loss_rpn_box_reg: 0.0067 (0.0069) time: 0.2159 data: 0.0155 max mem: 2561
Epoch: [0] [50/60] eta: 0:00:02 lr: 0.004323 loss: 0.4203 (0.9047) loss_classifier: 0.0506 (0.1600) loss_box_reg: 0.1966 (0.2222) loss_mask: 0.1716 (0.5033) loss_objectness: 0.0013 (0.0123) loss_rpn_box_reg: 0.0044 (0.0069) time: 0.2204 data: 0.0164 max mem: 2824
Epoch: [0] [59/60] eta: 0:00:00 lr: 0.005000 loss: 0.4211 (0.8255) loss_classifier: 0.0553 (0.1439) loss_box_reg: 0.1866 (0.2116) loss_mask: 0.1587 (0.4523) loss_objectness: 0.0010 (0.0107) loss_rpn_box_reg: 0.0054 (0.0070) time: 0.2260 data: 0.0167 max mem: 2824
Epoch: [0] Total time: 0:00:13 (0.2298 s / it)
creating index...
index created!
Test: [ 0/50] eta: 0:00:05 model_time: 0.0884 (0.0884) evaluator_time: 0.0075 (0.0075) time: 0.1092 data: 0.0127 max mem: 2824
Test: [49/50] eta: 0:00:00 model_time: 0.0399 (0.0554) evaluator_time: 0.0030 (0.0050) time: 0.0648 data: 0.0086 max mem: 2824
Test: Total time: 0:00:03 (0.0699 s / it)
Averaged stats: model_time: 0.0399 (0.0554) evaluator_time: 0.0030 (0.0050)
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.675
Average Precision (AP) @[ IoU=0.50 | area= all | maxDets=100 ] = 0.991
Average Precision (AP) @[ IoU=0.75 | area= all | maxDets=100 ] = 0.850
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.682
Average Precision (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.677
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 1 ] = 0.324
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 10 ] = 0.737
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.737
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.758
Average Recall (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.735
IoU metric: segm
Average Precision (AP) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.717
Average Precision (AP) @[ IoU=0.50 | area= all | maxDets=100 ] = 0.994
Average Precision (AP) @[ IoU=0.75 | area= all | maxDets=100 ] = 0.871
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.545
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.332
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 10 ] = 0.757
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.758
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.733
Average Recall (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.761
Epoch: [1] [ 0/60] eta: 0:00:12 lr: 0.005000 loss: 0.2042 (0.2042) loss_classifier: 0.0236 (0.0236) loss_box_reg: 0.0754 (0.0754) loss_mask: 0.1004 (0.1004) loss_objectness: 0.0036 (0.0036) loss_rpn_box_reg: 0.0012 (0.0012) time: 0.2142 data: 0.0163 max mem: 2824
Epoch: [1] [10/60] eta: 0:00:11 lr: 0.005000 loss: 0.3474 (0.3296) loss_classifier: 0.0510 (0.0481) loss_box_reg: 0.1155 (0.1250) loss_mask: 0.1556 (0.1471) loss_objectness: 0.0010 (0.0026) loss_rpn_box_reg: 0.0063 (0.0068) time: 0.2201 data: 0.0165 max mem: 2824
Epoch: [1] [20/60] eta: 0:00:08 lr: 0.005000 loss: 0.2951 (0.3029) loss_classifier: 0.0454 (0.0444) loss_box_reg: 0.0938 (0.1007) loss_mask: 0.1488 (0.1494) loss_objectness: 0.0012 (0.0021) loss_rpn_box_reg: 0.0056 (0.0063) time: 0.2146 data: 0.0161 max mem: 2824
Epoch: [1] [30/60] eta: 0:00:06 lr: 0.005000 loss: 0.2519 (0.2987) loss_classifier: 0.0326 (0.0421) loss_box_reg: 0.0623 (0.0951) loss_mask: 0.1488 (0.1523) loss_objectness: 0.0014 (0.0024) loss_rpn_box_reg: 0.0056 (0.0068) time: 0.2094 data: 0.0150 max mem: 2824
Epoch: [1] [40/60] eta: 0:00:04 lr: 0.005000 loss: 0.2519 (0.2927) loss_classifier: 0.0326 (0.0406) loss_box_reg: 0.0706 (0.0921) loss_mask: 0.1334 (0.1516) loss_objectness: 0.0014 (0.0023) loss_rpn_box_reg: 0.0040 (0.0061) time: 0.2166 data: 0.0153 max mem: 2890
Epoch: [1] [50/60] eta: 0:00:02 lr: 0.005000 loss: 0.2467 (0.2817) loss_classifier: 0.0372 (0.0401) loss_box_reg: 0.0694 (0.0860) loss_mask: 0.1305 (0.1477) loss_objectness: 0.0012 (0.0023) loss_rpn_box_reg: 0.0031 (0.0057) time: 0.2168 data: 0.0163 max mem: 2890
Epoch: [1] [59/60] eta: 0:00:00 lr: 0.005000 loss: 0.2582 (0.2838) loss_classifier: 0.0387 (0.0405) loss_box_reg: 0.0725 (0.0862) loss_mask: 0.1367 (0.1489) loss_objectness: 0.0012 (0.0022) loss_rpn_box_reg: 0.0041 (0.0059) time: 0.2152 data: 0.0169 max mem: 2890
Epoch: [1] Total time: 0:00:12 (0.2149 s / it)
creating index...
index created!
Test: [ 0/50] eta: 0:00:03 model_time: 0.0445 (0.0445) evaluator_time: 0.0062 (0.0062) time: 0.0638 data: 0.0127 max mem: 2890
Test: [49/50] eta: 0:00:00 model_time: 0.0389 (0.0399) evaluator_time: 0.0026 (0.0033) time: 0.0523 data: 0.0086 max mem: 2890
Test: Total time: 0:00:02 (0.0528 s / it)
Averaged stats: model_time: 0.0389 (0.0399) evaluator_time: 0.0026 (0.0033)
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.778
Average Precision (AP) @[ IoU=0.50 | area= all | maxDets=100 ] = 0.996
Average Precision (AP) @[ IoU=0.75 | area= all | maxDets=100 ] = 0.929
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.767
Average Precision (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.782
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 1 ] = 0.356
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 10 ] = 0.815
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.815
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.800
Average Recall (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.817
IoU metric: segm
Average Precision (AP) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.741
Average Precision (AP) @[ IoU=0.50 | area= all | maxDets=100 ] = 0.996
Average Precision (AP) @[ IoU=0.75 | area= all | maxDets=100 ] = 0.917
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.588
Average Precision (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.750
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 1 ] = 0.345
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 10 ] = 0.776
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.776
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.725
Average Recall (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.782
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))
<matplotlib.image.AxesImage object at 0x7f8fd6bd4250>
结果看起来不错！
总结#
在本教程中，您学习了如何为自定义数据集上的目标检测模型创建训练流水线。为此，您编写了一个返回图像、真值框和分割掩码的 torch.utils.data.Dataset 类。您还利用了在 COCO train2017 上预训练的 Mask R-CNN 模型，在该新数据集上执行了迁移学习。
有关包含多机/多 GPU 训练的更完整示例，请查看 torchvision 存储库中提供的 references/detection/train.py。
脚本总运行时间: (0 分 48.060 秒)