torchvision对象检测微调教程

原文：https://pytorch.org/tutorials/intermediate/torchvision_tutorial.html

小费

为了充分利用本教程，我们建议使用此 Colab 版本。 这将使您可以尝试以下信息。

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

定义数据集

用于训练对象检测，实例细分和人员关键点检测的参考脚本可轻松支持添加新的自定义数据集。 数据集应继承自标准torch.utils.data.Dataset类，并实现__len__和__getitem__。

我们唯一需要的特异性是数据集__getitem__应该返回：

  • 图像：大小为(H, W)的 PIL 图像
  • 目标：包含以下字段的字典
    • boxes (FloatTensor[N, 4])：[x0, y0, x1, y1]格式的N边界框的坐标，范围从0至W，从0至H
    • labels (Int64Tensor[N])：每个边界框的标签。 0始终代表背景类。
    • image_id (Int64Tensor[1])：图像标识符。 它在数据集中的所有图像之间应该是唯一的，并在评估过程中使用
    • area (Tensor[N])：边界框的区域。 在使用 COCO 度量进行评估时，可使用此值来区分小盒子，中盒子和大盒子之间的度量得分。
    • iscrowd (UInt8Tensor[N])：iscrowd = True的实例在评估期间将被忽略。
    • （可选）masks (UInt8Tensor[N, H, W])：每个对象的分割蒙版
    • （可选）keypoints (FloatTensor[N, K, 3])：对于 N 个对象中的每一个，它包含[x, y, visibility]格式的 K 个关键点，以定义对象。 可见性为 0 表示关键点不可见。 请注意，对于数据扩充，翻转关键点的概念取决于数据表示形式，您可能应该将references/detection/transforms.py修改为新的关键点表示形式

如果您的模型返回上述方法，则它们将使其适用于训练和评估，并将使用pycocotools中的评估脚本。

注意

对于 Windows，请使用命令从gautamchitnis安装pycocotools

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

关于labels的注解。 该模型将0类作为背景。 如果您的数据集不包含背景类，则labels中不应包含0。 例如，假设您只有猫和狗两类，则可以定义1来表示猫和0代表狗。 因此，例如，如果其中一个图像同时具有两个类，则您的labels张量应类似于[1,2]。

此外，如果要在训练过程中使用宽高比分组（以便每个批量仅包含具有相似长宽比的图像），则建议您还实现get_height_and_width方法，该方法返回图像的高度和宽度。 如果未提供此方法，我们将通过__getitem__查询数据集的所有元素，这会将图像加载到内存中，并且比提供自定义方法慢。

为 PennFudan 编写自定义数据集

让我们为 PennFudan 数据集编写一个数据集。 在下载并解压缩 zip 文件之后，我们具有以下文件夹结构：

  1. PennFudanPed/
  2.   PedMasks/
  3.     FudanPed00001_mask.png
  4.     FudanPed00002_mask.png
  5.     FudanPed00003_mask.png
  6.     FudanPed00004_mask.png
  7.     ...
  8.   PNGimg/
  9.     FudanPed00001.png
  10.     FudanPed00002.png
  11.     FudanPed00003.png
  12.     FudanPed00004.png

这是一对图像和分割蒙版的一个示例

intermediate/../../_static/img/tv_tutorial/tv_image01.png intermediate/../../_static/img/tv_tutorial/tv_image02.png

因此，每个图像都有一个对应的分割蒙版，其中每个颜色对应一个不同的实例。 让我们为此数据集编写一个torch.utils.data.Dataset类。

  1. import os
  2. import numpy as np
  3. import torch
  4. from PIL import Image
  6. class PennFudanDataset(object):
  7.     def __init__(self, root, transforms):
  8.         self.root = root
  9.         self.transforms = transforms
  10.         # load all image files, sorting them to
  11.         # ensure that they are aligned
  12.         self.imgs = list(sorted(os.listdir(os.path.join(root, "PNGImages"))))
  13.         self.masks = list(sorted(os.listdir(os.path.join(root, "PedMasks"))))
  15.     def __getitem__(self, idx):
  16.         # load images ad masks
  17.         img_path = os.path.join(self.root, "PNGImages", self.imgs[idx])
  18.         mask_path = os.path.join(self.root, "PedMasks", self.masks[idx])
  19.         img = Image.open(img_path).convert("RGB")
  20.         # note that we haven't converted the mask to RGB,
  21.         # because each color corresponds to a different instance
  22.         # with 0 being background
  23.         mask = Image.open(mask_path)
  24.         # convert the PIL Image into a numpy array
  25.         mask = np.array(mask)
  26.         # instances are encoded as different colors
  27.         obj_ids = np.unique(mask)
  28.         # first id is the background, so remove it
  29.         obj_ids = obj_ids[1:]
  31.         # split the color-encoded mask into a set
  32.         # of binary masks
  33.         masks = mask == obj_ids[:, None, None]
  35.         # get bounding box coordinates for each mask
  36.         num_objs = len(obj_ids)
  37.         boxes = []
  38.         for i in range(num_objs):
  39.             pos = np.where(masks[i])
  40.             xmin = np.min(pos[1])
  41.             xmax = np.max(pos[1])
  42.             ymin = np.min(pos[0])
  43.             ymax = np.max(pos[0])
  44.             boxes.append([xmin, ymin, xmax, ymax])
  46.         # convert everything into a torch.Tensor
  47.         boxes = torch.as_tensor(boxes, dtype=torch.float32)
  48.         # there is only one class
  49.         labels = torch.ones((num_objs,), dtype=torch.int64)
  50.         masks = torch.as_tensor(masks, dtype=torch.uint8)
  52.         image_id = torch.tensor([idx])
  53.         area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])
  54.         # suppose all instances are not crowd
  55.         iscrowd = torch.zeros((num_objs,), dtype=torch.int64)
  57.         target = {}
  58.         target["boxes"] = boxes
  59.         target["labels"] = labels
  60.         target["masks"] = masks
  61.         target["image_id"] = image_id
  62.         target["area"] = area
  63.         target["iscrowd"] = iscrowd
  65.         if self.transforms is not None:
  66.             img, target = self.transforms(img, target)
  68.         return img, target
  70.     def __len__(self):
  71.         return len(self.imgs)

这就是数据集的全部内容。 现在，我们定义一个可以对该数据集执行预测的模型。

定义模型

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

intermediate/../../_static/img/tv_tutorial/tv_image03.png

Mask R-CNN 在 Faster R-CNN 中增加了一个分支，该分支还可以预测每个实例的分割掩码。

intermediate/../../_static/img/tv_tutorial/tv_image04.png

在两种常见情况下，可能要修改torchvision模型动物园中的可用模型之一。 首先是当我们想从预先训练的模型开始，然后微调最后一层时。 另一个是当我们想要用另一个模型替换主干时（例如，为了更快的预测）。

在以下各节中，让我们看看如何做一个或另一个。

1-将预训练模型用于微调

假设您要从在 COCO 上经过预训练的模型开始，并希望针对您的特定类对其进行微调。 这是一种可行的方法：

  1. import torchvision
  2. from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
  4. # load a model pre-trained pre-trained on COCO
  5. model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
  7. # replace the classifier with a new one, that has
  8. # num_classes which is user-defined
  9. num_classes = 2  # 1 class (person) + background
  10. # get number of input features for the classifier
  11. in_features = model.roi_heads.box_predictor.cls_score.in_features
  12. # replace the pre-trained head with a new one
  13. model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

2-修改模型以添加其他主干

  1. import torchvision
  2. from torchvision.models.detection import FasterRCNN
  3. from torchvision.models.detection.rpn import AnchorGenerator
  5. # load a pre-trained model for classification and return
  6. # only the features
  7. backbone = torchvision.models.mobilenet_v2(pretrained=True).features
  8. # FasterRCNN needs to know the number of
  9. # output channels in a backbone. For mobilenet_v2, it's 1280
  10. # so we need to add it here
  11. backbone.out_channels = 1280
  13. # let's make the RPN generate 5 x 3 anchors per spatial
  14. # location, with 5 different sizes and 3 different aspect
  15. # ratios. We have a Tuple[Tuple[int]] because each feature
  16. # map could potentially have different sizes and
  17. # aspect ratios
  18. anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512),),
  19.                                    aspect_ratios=((0.5, 1.0, 2.0),))
  21. # let's define what are the feature maps that we will
  22. # use to perform the region of interest cropping, as well as
  23. # the size of the crop after rescaling.
  24. # if your backbone returns a Tensor, featmap_names is expected to
  25. # be [0]. More generally, the backbone should return an
  26. # OrderedDict[Tensor], and in featmap_names you can choose which
  27. # feature maps to use.
  28. roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=[0],
  29.                                                 output_size=7,
  30.                                                 sampling_ratio=2)
  32. # put the pieces together inside a FasterRCNN model
  33. model = FasterRCNN(backbone,
  34.                    num_classes=2,
  35.                    rpn_anchor_generator=anchor_generator,
  36.                    box_roi_pool=roi_pooler)

PennFudan 数据集的实例细分模型

在我们的案例中，由于我们的数据集非常小，我们希望从预训练模型中进行微调，因此我们将遵循方法 1。

这里我们还想计算实例分割掩码，因此我们将使用 Mask R-CNN：

  1. import torchvision
  2. from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
  3. from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor
  5. def get_model_instance_segmentation(num_classes):
  6.     # load an instance segmentation model pre-trained pre-trained on COCO
  7.     model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True)
  9.     # get number of input features for the classifier
  10.     in_features = model.roi_heads.box_predictor.cls_score.in_features
  11.     # replace the pre-trained head with a new one
  12.     model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
  14.     # now get the number of input features for the mask classifier
  15.     in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
  16.     hidden_layer = 256
  17.     # and replace the mask predictor with a new one
  18.     model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask,
  19.                                                        hidden_layer,
  20.                                                        num_classes)
  22.     return model

就是这样，这将使model随时可以在您的自定义数据集上进行训练和评估。

将所有内容放在一起

在references/detection/中，我们提供了许多帮助程序功能来简化训练和评估检测模型。 在这里，我们将使用references/detection/engine.py，references/detection/utils.py和references/detection/transforms.py。 只需将它们复制到您的文件夹中，然后在此处使用它们即可。

让我们写一些辅助函数来进行数据扩充/转换：

  1. import transforms as T
  3. def get_transform(train):
  4.     transforms = []
  5.     transforms.append(T.ToTensor())
  6.     if train:
  7.         transforms.append(T.RandomHorizontalFlip(0.5))
  8.     return T.Compose(transforms)

测试forward()方法（可选）

在遍历数据集之前，最好先查看模型在训练过程中的期望值以及对样本数据的推断时间。

  1. model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
  2. dataset = PennFudanDataset('PennFudanPed', get_transform(train=True))
  3. data_loader = torch.utils.data.DataLoader(
  4.  dataset, batch_size=2, shuffle=True, num_workers=4,
  5.  collate_fn=utils.collate_fn)
  6. # For Training
  7. images,targets = next(iter(data_loader))
  8. images = list(image for image in images)
  9. targets = [{k: v for k, v in t.items()} for t in targets]
  10. output = model(images,targets)   # Returns losses and detections
  11. # For inference
  12. model.eval()
  13. x = [torch.rand(3, 300, 400), torch.rand(3, 500, 400)]
  14. predictions = model(x)           # Returns predictions

现在，我们编写执行训练和验证的main函数：

  1. from engine import train_one_epoch, evaluate
  2. import utils
  4. def main():
  5.     # train on the GPU or on the CPU, if a GPU is not available
  6.     device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
  8.     # our dataset has two classes only - background and person
  9.     num_classes = 2
  10.     # use our dataset and defined transformations
  11.     dataset = PennFudanDataset('PennFudanPed', get_transform(train=True))
  12.     dataset_test = PennFudanDataset('PennFudanPed', get_transform(train=False))
  14.     # split the dataset in train and test set
  15.     indices = torch.randperm(len(dataset)).tolist()
  16.     dataset = torch.utils.data.Subset(dataset, indices[:-50])
  17.     dataset_test = torch.utils.data.Subset(dataset_test, indices[-50:])
  19.     # define training and validation data loaders
  20.     data_loader = torch.utils.data.DataLoader(
  21.         dataset, batch_size=2, shuffle=True, num_workers=4,
  22.         collate_fn=utils.collate_fn)
  24.     data_loader_test = torch.utils.data.DataLoader(
  25.         dataset_test, batch_size=1, shuffle=False, num_workers=4,
  26.         collate_fn=utils.collate_fn)
  28.     # get the model using our helper function
  29.     model = get_model_instance_segmentation(num_classes)
  31.     # move model to the right device
  32.     model.to(device)
  34.     # construct an optimizer
  35.     params = [p for p in model.parameters() if p.requires_grad]
  36.     optimizer = torch.optim.SGD(params, lr=0.005,
  37.                                 momentum=0.9, weight_decay=0.0005)
  38.     # and a learning rate scheduler
  39.     lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
  40.                                                    step_size=3,
  41.                                                    gamma=0.1)
  43.     # let's train it for 10 epochs
  44.     num_epochs = 10
  46.     for epoch in range(num_epochs):
  47.         # train for one epoch, printing every 10 iterations
  48.         train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq=10)
  49.         # update the learning rate
  50.         lr_scheduler.step()
  51.         # evaluate on the test dataset
  52.         evaluate(model, data_loader_test, device=device)
  54.     print("That's it!")

您应该获得第一个周期的输出：

  1. Epoch: [0]  [ 0/60]  eta: 0:01:18  lr: 0.000090  loss: 2.5213 (2.5213)  loss_classifier: 0.8025 (0.8025)  loss_box_reg: 0.2634 (0.2634)  loss_mask: 1.4265 (1.4265)  loss_objectness: 0.0190 (0.0190)  loss_rpn_box_reg: 0.0099 (0.0099)  time: 1.3121  data: 0.3024  max mem: 3485
  2. Epoch: [0]  [10/60]  eta: 0:00:20  lr: 0.000936  loss: 1.3007 (1.5313)  loss_classifier: 0.3979 (0.4719)  loss_box_reg: 0.2454 (0.2272)  loss_mask: 0.6089 (0.7953)  loss_objectness: 0.0197 (0.0228)  loss_rpn_box_reg: 0.0121 (0.0141)  time: 0.4198  data: 0.0298  max mem: 5081
  3. Epoch: [0]  [20/60]  eta: 0:00:15  lr: 0.001783  loss: 0.7567 (1.1056)  loss_classifier: 0.2221 (0.3319)  loss_box_reg: 0.2002 (0.2106)  loss_mask: 0.2904 (0.5332)  loss_objectness: 0.0146 (0.0176)  loss_rpn_box_reg: 0.0094 (0.0123)  time: 0.3293  data: 0.0035  max mem: 5081
  4. Epoch: [0]  [30/60]  eta: 0:00:11  lr: 0.002629  loss: 0.4705 (0.8935)  loss_classifier: 0.0991 (0.2517)  loss_box_reg: 0.1578 (0.1957)  loss_mask: 0.1970 (0.4204)  loss_objectness: 0.0061 (0.0140)  loss_rpn_box_reg: 0.0075 (0.0118)  time: 0.3403  data: 0.0044  max mem: 5081
  5. Epoch: [0]  [40/60]  eta: 0:00:07  lr: 0.003476  loss: 0.3901 (0.7568)  loss_classifier: 0.0648 (0.2022)  loss_box_reg: 0.1207 (0.1736)  loss_mask: 0.1705 (0.3585)  loss_objectness: 0.0018 (0.0113)  loss_rpn_box_reg: 0.0075 (0.0112)  time: 0.3407  data: 0.0044  max mem: 5081
  6. Epoch: [0]  [50/60]  eta: 0:00:03  lr: 0.004323  loss: 0.3237 (0.6703)  loss_classifier: 0.0474 (0.1731)  loss_box_reg: 0.1109 (0.1561)  loss_mask: 0.1658 (0.3201)  loss_objectness: 0.0015 (0.0093)  loss_rpn_box_reg: 0.0093 (0.0116)  time: 0.3379  data: 0.0043  max mem: 5081
  7. Epoch: [0]  [59/60]  eta: 0:00:00  lr: 0.005000  loss: 0.2540 (0.6082)  loss_classifier: 0.0309 (0.1526)  loss_box_reg: 0.0463 (0.1405)  loss_mask: 0.1568 (0.2945)  loss_objectness: 0.0012 (0.0083)  loss_rpn_box_reg: 0.0093 (0.0123)  time: 0.3489  data: 0.0042  max mem: 5081
  8. Epoch: [0] Total time: 0:00:21 (0.3570 s / it)
  9. creating index...
  10. index created!
  11. Test:  [ 0/50]  eta: 0:00:19  model_time: 0.2152 (0.2152)  evaluator_time: 0.0133 (0.0133)  time: 0.4000  data: 0.1701  max mem: 5081
  12. Test:  [49/50]  eta: 0:00:00  model_time: 0.0628 (0.0687)  evaluator_time: 0.0039 (0.0064)  time: 0.0735  data: 0.0022  max mem: 5081
  13. Test: Total time: 0:00:04 (0.0828 s / it)
  14. Averaged stats: model_time: 0.0628 (0.0687)  evaluator_time: 0.0039 (0.0064)
  15. Accumulating evaluation results...
  16. DONE (t=0.01s).
  17. Accumulating evaluation results...
  18. DONE (t=0.01s).
  19. IoU metric: bbox
  20.  Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.606
  21.  Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.984
  22.  Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.780
  23.  Average Precision  (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.313
  24.  Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.582
  25.  Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.612
  26.  Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.270
  27.  Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.672
  28.  Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.672
  29.  Average Recall     (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.650
  30.  Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.755
  31.  Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.664
  32. IoU metric: segm
  33.  Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.704
  34.  Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.979
  35.  Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.871
  36.  Average Precision  (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.325
  37.  Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.488
  38.  Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.727
  39.  Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.316
  40.  Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.748
  41.  Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.749
  42.  Average Recall     (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.650
  43.  Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.673
  44.  Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.758

因此，经过一个周期的训练，我们获得了 60.6 的 COCO 风格 mAP 和 70.4 的遮罩 mAP。

经过 10 个周期的训练，我得到了以下指标

  1. IoU metric: bbox
  2.  Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.799
  3.  Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.969
  4.  Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.935
  5.  Average Precision  (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.349
  6.  Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.592
  7.  Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.831
  8.  Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.324
  9.  Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.844
  10.  Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.844
  11.  Average Recall     (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.400
  12.  Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.777
  13.  Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.870
  14. IoU metric: segm
  15.  Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.761
  16.  Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.969
  17.  Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.919
  18.  Average Precision  (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.341
  19.  Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.464
  20.  Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.788
  21.  Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.303
  22.  Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.799
  23.  Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.799
  24.  Average Recall     (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.400
  25.  Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.769
  26.  Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.818

但是这些预测是什么样的？ 让我们在数据集中拍摄一张图像并进行验证

intermediate/../../_static/img/tv_tutorial/tv_image05.png

经过训练的模型会在此图片中预测 9 个人物实例，让我们看看其中的几个：

intermediate/../../_static/img/tv_tutorial/tv_image06.png intermediate/../../_static/img/tv_tutorial/tv_image07.png

结果看起来还不错！

总结

在本教程中，您学习了如何在自定义数据集上为实例细分模型创建自己的训练管道。 为此，您编写了一个torch.utils.data.Dataset类，该类返回图像以及真实情况框和分段蒙版。 您还利用了在 COCO train2017 上预先训练的 Mask R-CNN 模型，以便对该新数据集执行迁移学习。

对于更完整的示例（包括多机/多 GPU 训练），请检查在torchvision存储库中存在的references/detection/train.py。

您可以在此处下载本教程的完整源文件。