PyTorch TorchVision 對象檢測微調(diào)教程

2022-04-06 18:29 更新

提示
為了充分利用本教程,我們建議使用此Colab版本。這將允許您嘗試下面提供的信息。

在本教程中，我們將在賓夕法尼亞復旦數(shù)據(jù)庫中微調(diào)預先訓練的?Mask R-CNN?模型，以進行行人檢測和分割。它包含 170 張包含 345 個行人實例的圖像，我們將使用它來說明如何使用 torchvision 中的新功能，以便在自定義數(shù)據(jù)集上訓練實例分割模型。

定義數(shù)據(jù)集

用于訓練對象檢測、實例分段和人員關鍵點檢測的參考腳本可以輕松支持添加新的自定義數(shù)據(jù)集。數(shù)據(jù)集應繼承自標準類torch.utils.data.Dataset，并實現(xiàn) __len__和__getitem__。

我們要求的唯一特異性是數(shù)據(jù)集應返回：__getitem__

  • 圖像：PIL 大小的圖像(H,?W)
  • 目標：包含以下字段的字典
    • boxe（FloatTensor[N，4]）：N個邊界框的坐標，格式為[x0，y0，x1，y1]，范圍從0到W，從0到H.
    • labels（Int64Tensor[N]）：每個邊界框的標簽。0始終代表后臺類。
    • image_id（Int64Tensor[1]）：圖像標識符。它在數(shù)據(jù)集中的所有圖像之間應該是唯一的，并在評估期間使用
    • area (Tensor[N]):邊界框的面積。在使用COCO度量進行評估時，會使用它來區(qū)分小盒子、中盒子和大盒子之間的度量分數(shù)。
    • iscrowd（UInt8Tensor[N]）：iscrowd=True的實例將在計算期間被忽略。
    • （可選）masks（UInt8Tensor[N，H，W]）：每個對象的分割遮罩.
    • （可選）keypoints（FloatTensor[N，K，3]）：對于N個對象中的每一個，它包含[x，y，visibility]格式的K個關鍵點，定義對象。可見性=0表示關鍵點不可見。請注意，對于數(shù)據(jù)擴充，翻轉(zhuǎn)關鍵點的概念取決于數(shù)據(jù)表示，您應該調(diào)整references/detection/transforms.py用于新的關鍵點表示

如果您的模型返回上述方法，它們將使培訓和評估都有效，并將使用Pycocotools中的評估腳本，該腳本可以與pip install Pycools一起安裝。

注意
對于Windows，請使用命令從gautamchitnis安裝pycocotools
pip?install?git+https://github.com/gautamchitnis/cocoapi.git@cocodataset-master#subdirectory=PythonAPI

labels上有一個注釋。該模型將0類視為背景。如果數(shù)據(jù)集不包含后臺類，則labels中不應包含0。例如，假設只有兩個類，cat和dog，可以定義1（而不是0）來表示cat，定義2來表示dogs。例如，如果其中一個圖像同時具有兩個類，那么labels張量應該類似于[1,2]。

此外，如果希望在訓練期間使用縱橫比分組（以便每個批次僅包含具有類似縱橫比的圖像），則建議還實現(xiàn)get_height_和_width方法，該方法返回圖像的高度和寬度。如果不提供此方法，我們將通過__getitem__查詢數(shù)據(jù)集的所有元素，__getitem__將圖像加載到內(nèi)存中，速度比提供自定義方法慢。

為PennFudan編寫自定義數(shù)據(jù)集

讓我們?yōu)镻ennFudan數(shù)據(jù)集編寫一個數(shù)據(jù)集。下載并解壓縮zip文件后，我們有以下文件夾結(jié)構(gòu)：

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

下面是一對圖像和分割蒙版的一個示例

?

因此，每個圖像都有一個相應的分割遮罩，每個顏色對應一個不同的實例。讓我們?yōu)檫@個數(shù)據(jù)集編寫一個torch.utils.data.Dataset類。

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

這就是數(shù)據(jù)集的全部內(nèi)容?，F(xiàn)在，讓我們定義一個可以對此數(shù)據(jù)集執(zhí)行預測的模型。

定義模型

在本教程中，我們將使用Mask R-CNN，它基于更快的R-CNN。更快的R-CNN模型可以預測圖像中潛在對象的邊界框和類分數(shù)。

Mask R-CNN為更快的R-CNN添加了一個額外分支，該分支還可以預測每個實例的分段掩碼。

在兩種常見情況下，可能需要修改torchvision modelzoo中的一個可用模型。第一個是當我們想從一個預先訓練好的模型開始，只需微調(diào)最后一層。另一個是當我們想用另一個模型替換模型的主干時（例如，為了更快的預測）。

讓我們來看看在接下來的部分中我們將如何做一個或另一個。

1 - 從預訓練模型微調(diào)

假設您要從在 COCO 上預先訓練的模型開始，并希望針對您的特定類對其進行微調(diào)。以下是一種可能的方法：

  1. import torchvision
  2. from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
  5. ## load a model pre-trained on COCO
  6. model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
  9. ## replace the classifier with a new one, that has
  10. ## num_classes which is user-defined
  11. num_classes = 2  # 1 class (person) + background
  12. ## get number of input features for the classifier
  13. in_features = model.roi_heads.box_predictor.cls_score.in_features
  14. ## replace the pre-trained head with a new one
  15. 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
  6. ## load a pre-trained model for classification and return
  7. ## only the features
  8. backbone = torchvision.models.mobilenet_v2(pretrained=True).features
  9. ## FasterRCNN needs to know the number of
  10. ## output channels in a backbone. For mobilenet_v2, it's 1280
  11. ## so we need to add it here
  12. backbone.out_channels = 1280
  15. ## let's make the RPN generate 5 x 3 anchors per spatial
  16. ## location, with 5 different sizes and 3 different aspect
  17. ## ratios. We have a Tuple[Tuple[int]] because each feature
  18. ## map could potentially have different sizes and
  19. ## aspect ratios
  20. anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512),),
  21.                                    aspect_ratios=((0.5, 1.0, 2.0),))
  24. ## let's define what are the feature maps that we will
  25. ## use to perform the region of interest cropping, as well as
  26. ## the size of the crop after rescaling.
  27. ## if your backbone returns a Tensor, featmap_names is expected to
  28. ## be [0]. More generally, the backbone should return an
  29. ## OrderedDict[Tensor], and in featmap_names you can choose which
  30. ## feature maps to use.
  31. roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=['0'],
  32.                                                 output_size=7,
  33.                                                 sampling_ratio=2)
  36. ## put the pieces together inside a FasterRCNN model
  37. model = FasterRCNN(backbone,
  38.                    num_classes=2,
  39.                    rpn_anchor_generator=anchor_generator,
  40.                    box_roi_pool=roi_pooler)

一種PennFudan數(shù)據(jù)集的實例分割模型

在我們的例子中，我們希望從預先訓練的模型進行微調(diào)，因為我們的數(shù)據(jù)集非常小，所以我們將遵循方法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
  8. def get_model_instance_segmentation(num_classes):
  9.     # load an instance segmentation model pre-trained on COCO
  10.     model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True)
  13.     # get number of input features for the classifier
  14.     in_features = model.roi_heads.box_predictor.cls_score.in_features
  15.     # replace the pre-trained head with a new one
  16.     model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
  19.     # now get the number of input features for the mask classifier
  20.     in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
  21.     hidden_layer = 256
  22.     # and replace the mask predictor with a new one
  23.     model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask,
  24.                                                        hidden_layer,
  25.                                                        num_classes)
  28.     return model

就這樣，這將使model準備好在自定義數(shù)據(jù)集上進行訓練和評估。

將所有內(nèi)容放在一起

在references/detection/中，我們有許多幫助函數(shù)來簡化訓練和評估檢測模型。在這里，我們將使用references/detection/engine.py，references/detection/utils.py和references/detection/transforms.py。只需將references/detection下的所有內(nèi)容復制到您的文件夾中，并在此處使用它們。

讓我們?yōu)閿?shù)據(jù)擴充/轉(zhuǎn)換編寫一些幫助函數(shù)：

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

測試forward()方法（可選）

在迭代數(shù)據(jù)集之前，最好先查看模型在樣本數(shù)據(jù)上的訓練和推理時間內(nèi)的預期內(nèi)容。

  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

現(xiàn)在，讓我們編寫執(zhí)行訓練和驗證的 main 函數(shù)：

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

因此，經(jīng)過一個時期的訓練，我們獲得了60.6的COCO式mAP和70.4的掩模mAP。 經(jīng)過 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

但是預測是什么樣的呢？讓我們在數(shù)據(jù)集中獲取一個圖像并進行驗證

經(jīng)過訓練的模型預測了此圖像中的 9 個人員實例，讓我們看一下其中的幾個實例：

?

結(jié)果看起來相當不錯！

結(jié)束語

在本教程中，您學習了如何在自定義數(shù)據(jù)集上創(chuàng)建自己的培訓管道，例如分段模型。為此，你寫了一個torch.utils.data.Dataset類，該類返回圖像、地面真值框和分割遮罩。為了在這個新數(shù)據(jù)集上執(zhí)行轉(zhuǎn)移學習，您還利用了COCO train2017上預先培訓的Mask R-CNN模型。

有關更完整的示例，包括多機/多gpu培訓，請查看references/detection/train.py，它出現(xiàn)在torchvision回購協(xié)議中。

在此處下載本教程的完整源文件。

以上內(nèi)容是否對您有幫助：
在線筆記
App下載
App下載

掃描二維碼

下載編程獅App

公眾號
微信公眾號

編程獅公眾號