評價此頁

★ ★ ★ ★ ★

intermediate/torchvision_tutorial

在 Google Colab 中執行

注意

跳轉到最後下載完整的示例程式碼。

TorchVision 物件檢測微調教程#

創建於：2023 年 12 月 14 日 | 最後更新：2025 年 9 月 5 日 | 最後驗證：2024 年 11 月 5 日

在本教程中，我們將對預訓練的 Mask R-CNN 模型進行微調，並在賓夕法尼亞大學行人檢測和分割資料庫 (Penn-Fudan Database for Pedestrian Detection and Segmentation) 上進行訓練。該資料庫包含 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：N 個邊界框的座標，格式為 [x0, y0, x1, y1]，範圍從 0 到 W 和 0 到 H
- labels, 形狀為 [N] 的整數 torch.Tensor：每個邊界框的標籤。0 始終代表背景類。
- image_id, int: 影像識別符號。它應該在資料集的所有影像中是唯一的，並在評估期間使用。
- area, 形狀為 [N] 的浮點數 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))

<matplotlib.image.AxesImage object at 0x7f8955babee0>

因此，每張影像都有一個對應的分割掩碼，其中每種顏色代表一個不同的例項。讓我們為該資料集編寫一個 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 是一個同時預測影像中潛在物件的邊界框和類別分數的模型。

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

Mask R-CNN 在 Faster R-CNN 的基礎上增加了一個額外的分支，該分支還可以預測每個例項的分割掩碼。

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

在兩種常見情況下，我們可能需要修改 TorchVision Model Zoo 中提供的模型之一。第一種情況是我們想從預訓練模型開始，只微調最後一層。另一種情況是我們想用不同的骨幹網路替換模型的骨幹網路（例如，為了更快的預測）。

在接下來的章節中，我們將分別介紹如何實現這兩種方法。

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]
 26%|██▌       | 41.0M/160M [00:00<00:00, 429MB/s]
 52%|█████▏    | 83.1M/160M [00:00<00:00, 436MB/s]
 78%|███████▊  | 125M/160M [00:00<00:00, 419MB/s]
100%|██████████| 160M/160M [00:00<00:00, 426MB/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, 282MB/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.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.0396, grad_fn=<NllLossBackward0>), 'loss_box_reg': tensor(0.0487, grad_fn=<DivBackward0>), 'loss_objectness': tensor(0.0076, grad_fn=<BinaryCrossEntropyWithLogitsBackward0>), 'loss_rpn_box_reg': tensor(0.0048, 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]
 25%|██▍       | 42.0M/170M [00:00<00:00, 439MB/s]
 49%|████▉     | 84.0M/170M [00:00<00:00, 435MB/s]
 75%|███████▍  | 127M/170M [00:00<00:00, 442MB/s]
100%|██████████| 170M/170M [00:00<00:00, 444MB/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.

Epoch: [0]  [ 0/60]  eta: 0:00:39  lr: 0.000090  loss: 3.9297 (3.9297)  loss_classifier: 0.8942 (0.8942)  loss_box_reg: 0.2116 (0.2116)  loss_mask: 2.7965 (2.7965)  loss_objectness: 0.0242 (0.0242)  loss_rpn_box_reg: 0.0032 (0.0032)  time: 0.6557  data: 0.0122  max mem: 1753
Epoch: [0]  [10/60]  eta: 0:00:12  lr: 0.000936  loss: 1.4103 (1.9439)  loss_classifier: 0.4974 (0.5230)  loss_box_reg: 0.2421 (0.2619)  loss_mask: 0.7706 (1.1278)  loss_objectness: 0.0242 (0.0272)  loss_rpn_box_reg: 0.0032 (0.0039)  time: 0.2470  data: 0.0153  max mem: 2759
Epoch: [0]  [20/60]  eta: 0:00:09  lr: 0.001783  loss: 0.8459 (1.3453)  loss_classifier: 0.2110 (0.3548)  loss_box_reg: 0.2549 (0.2574)  loss_mask: 0.3053 (0.7064)  loss_objectness: 0.0116 (0.0229)  loss_rpn_box_reg: 0.0030 (0.0039)  time: 0.2098  data: 0.0157  max mem: 2759
Epoch: [0]  [30/60]  eta: 0:00:06  lr: 0.002629  loss: 0.5936 (1.0881)  loss_classifier: 0.1046 (0.2662)  loss_box_reg: 0.2762 (0.2534)  loss_mask: 0.1992 (0.5460)  loss_objectness: 0.0042 (0.0178)  loss_rpn_box_reg: 0.0041 (0.0047)  time: 0.2117  data: 0.0158  max mem: 2759
Epoch: [0]  [40/60]  eta: 0:00:04  lr: 0.003476  loss: 0.4908 (0.9438)  loss_classifier: 0.0671 (0.2164)  loss_box_reg: 0.2265 (0.2521)  loss_mask: 0.1678 (0.4537)  loss_objectness: 0.0051 (0.0161)  loss_rpn_box_reg: 0.0056 (0.0054)  time: 0.2133  data: 0.0168  max mem: 2759
Epoch: [0]  [50/60]  eta: 0:00:02  lr: 0.004323  loss: 0.4616 (0.8495)  loss_classifier: 0.0546 (0.1871)  loss_box_reg: 0.1976 (0.2361)  loss_mask: 0.1711 (0.4063)  loss_objectness: 0.0051 (0.0139)  loss_rpn_box_reg: 0.0066 (0.0061)  time: 0.2162  data: 0.0167  max mem: 2863
Epoch: [0]  [59/60]  eta: 0:00:00  lr: 0.005000  loss: 0.3605 (0.7698)  loss_classifier: 0.0450 (0.1650)  loss_box_reg: 0.1146 (0.2153)  loss_mask: 0.1818 (0.3713)  loss_objectness: 0.0024 (0.0121)  loss_rpn_box_reg: 0.0066 (0.0061)  time: 0.2098  data: 0.0150  max mem: 2863
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.0796 (0.0796)  evaluator_time: 0.0029 (0.0029)  time: 0.0939  data: 0.0108  max mem: 2863
Test:  [49/50]  eta: 0:00:00  model_time: 0.0416 (0.0582)  evaluator_time: 0.0036 (0.0058)  time: 0.0644  data: 0.0102  max mem: 2863
Test: Total time: 0:00:03 (0.0744 s / it)
Averaged stats: model_time: 0.0416 (0.0582)  evaluator_time: 0.0036 (0.0058)
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.684
 Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.978
 Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.829
 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.432
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.700
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.306
 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.764
 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.657
 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.669
 Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.978
 Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.880
 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.409
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.690
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.282
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.724
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.724
 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.643
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.729
Epoch: [1]  [ 0/60]  eta: 0:00:12  lr: 0.005000  loss: 0.3263 (0.3263)  loss_classifier: 0.0701 (0.0701)  loss_box_reg: 0.1182 (0.1182)  loss_mask: 0.1341 (0.1341)  loss_objectness: 0.0015 (0.0015)  loss_rpn_box_reg: 0.0024 (0.0024)  time: 0.2075  data: 0.0233  max mem: 2863
Epoch: [1]  [10/60]  eta: 0:00:11  lr: 0.005000  loss: 0.3105 (0.3021)  loss_classifier: 0.0435 (0.0443)  loss_box_reg: 0.0975 (0.1014)  loss_mask: 0.1341 (0.1485)  loss_objectness: 0.0013 (0.0020)  loss_rpn_box_reg: 0.0036 (0.0059)  time: 0.2218  data: 0.0173  max mem: 3233
Epoch: [1]  [20/60]  eta: 0:00:08  lr: 0.005000  loss: 0.3105 (0.3143)  loss_classifier: 0.0435 (0.0441)  loss_box_reg: 0.0975 (0.1026)  loss_mask: 0.1612 (0.1593)  loss_objectness: 0.0013 (0.0023)  loss_rpn_box_reg: 0.0036 (0.0061)  time: 0.2152  data: 0.0171  max mem: 3233
Epoch: [1]  [30/60]  eta: 0:00:06  lr: 0.005000  loss: 0.2894 (0.2990)  loss_classifier: 0.0423 (0.0424)  loss_box_reg: 0.0826 (0.0972)  loss_mask: 0.1506 (0.1519)  loss_objectness: 0.0008 (0.0020)  loss_rpn_box_reg: 0.0042 (0.0056)  time: 0.2079  data: 0.0165  max mem: 3233
Epoch: [1]  [40/60]  eta: 0:00:04  lr: 0.005000  loss: 0.2337 (0.2834)  loss_classifier: 0.0383 (0.0409)  loss_box_reg: 0.0672 (0.0904)  loss_mask: 0.1217 (0.1447)  loss_objectness: 0.0007 (0.0019)  loss_rpn_box_reg: 0.0035 (0.0055)  time: 0.2063  data: 0.0149  max mem: 3233
Epoch: [1]  [50/60]  eta: 0:00:02  lr: 0.005000  loss: 0.2160 (0.2717)  loss_classifier: 0.0297 (0.0389)  loss_box_reg: 0.0580 (0.0850)  loss_mask: 0.1183 (0.1406)  loss_objectness: 0.0007 (0.0018)  loss_rpn_box_reg: 0.0036 (0.0054)  time: 0.2046  data: 0.0142  max mem: 3233
Epoch: [1]  [59/60]  eta: 0:00:00  lr: 0.005000  loss: 0.2160 (0.2682)  loss_classifier: 0.0288 (0.0384)  loss_box_reg: 0.0578 (0.0817)  loss_mask: 0.1196 (0.1406)  loss_objectness: 0.0009 (0.0022)  loss_rpn_box_reg: 0.0036 (0.0052)  time: 0.2044  data: 0.0147  max mem: 3233
Epoch: [1] Total time: 0:00:12 (0.2084 s / it)
creating index...
index created!
Test:  [ 0/50]  eta: 0:00:02  model_time: 0.0398 (0.0398)  evaluator_time: 0.0023 (0.0023)  time: 0.0532  data: 0.0107  max mem: 3233
Test:  [49/50]  eta: 0:00:00  model_time: 0.0400 (0.0402)  evaluator_time: 0.0025 (0.0039)  time: 0.0548  data: 0.0102  max mem: 3233
Test: Total time: 0:00:02 (0.0544 s / it)
Averaged stats: model_time: 0.0400 (0.0402)  evaluator_time: 0.0025 (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.819
 Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.990
 Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.945
 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.487
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.827
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.362
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.854
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.854
 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.857
IoU metric: segm
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.757
 Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.990
 Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.935
 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.359
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.773
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.331
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.794
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.794
 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.657
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.802
That's it!

因此，經過一個 epoch 的訓練後，我們獲得了超過 50 的 COCO 風格 mAP，以及 65 的掩碼 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 0x7f8953ee74c0>

結果看起來不錯！

總結#

在本教程中，您學習瞭如何為自定義資料集上的物件檢測模型建立自己的訓練管道。為此，您編寫了一個 torch.utils.data.Dataset 類，該類返回影像和真實邊界框及分割掩碼。您還利用了在 COCO train2017 上預訓練的 Mask R-CNN 模型，以便在新資料集上執行遷移學習。

有關更完整的示例，包括多機/多 GPU 訓練，請檢視 torchvision 儲存庫中的 references/detection/train.py。

指令碼總執行時間： (0 分 47.300 秒)