免費(fèi)AI編程助手
PyTorch TorchVision 對象檢測微調(diào)教程
2022-04-06 18:29 更新
提示
為了充分利用本教程,我們建議使用此Colab版本。這將允許您嘗試下面提供的信息。
在本教程中，我們將在賓夕法尼亞復(fù)旦數(shù)據(jù)庫中微調(diào)預(yù)先訓(xùn)練的?Mask R-CNN?模型，以進(jìn)行行人檢測和分割。它包含 170 張包含 345 個(gè)行人實(shí)例的圖像，我們將使用它來說明如何使用 torchvision 中的新功能，以便在自定義數(shù)據(jù)集上訓(xùn)練實(shí)例分割模型。
定義數(shù)據(jù)集
用于訓(xùn)練對象檢測、實(shí)例分段和人員關(guān)鍵點(diǎn)檢測的參考腳本可以輕松支持添加新的自定義數(shù)據(jù)集。數(shù)據(jù)集應(yīng)繼承自標(biāo)準(zhǔn)類torch.utils.data.Dataset，并實(shí)現(xiàn) __len__和__getitem__。
我們要求的唯一特異性是數(shù)據(jù)集應(yīng)返回：__getitem__
- 圖像：PIL 大小的圖像
(H,?W) - 目標(biāo)：包含以下字段的字典
boxe（FloatTensor[N，4]）：N個(gè)邊界框的坐標(biāo)，格式為[x0，y0，x1，y1]，范圍從0到W，從0到H.labels（Int64Tensor[N]）：每個(gè)邊界框的標(biāo)簽。0始終代表后臺類。image_id（Int64Tensor[1]）：圖像標(biāo)識符。它在數(shù)據(jù)集中的所有圖像之間應(yīng)該是唯一的，并在評估期間使用area (Tensor[N]):邊界框的面積。在使用COCO度量進(jìn)行評估時(shí)，會使用它來區(qū)分小盒子、中盒子和大盒子之間的度量分?jǐn)?shù)。iscrowd（UInt8Tensor[N]）：iscrowd=True的實(shí)例將在計(jì)算期間被忽略。- （可選）
masks（UInt8Tensor[N，H，W]）：每個(gè)對象的分割遮罩. - （可選）
keypoints（FloatTensor[N，K，3]）：對于N個(gè)對象中的每一個(gè)，它包含[x，y，visibility]格式的K個(gè)關(guān)鍵點(diǎn)，定義對象。可見性=0表示關(guān)鍵點(diǎn)不可見。請注意，對于數(shù)據(jù)擴(kuò)充，翻轉(zhuǎn)關(guān)鍵點(diǎn)的概念取決于數(shù)據(jù)表示，您應(yīng)該調(diào)整references/detection/transforms.py用于新的關(guān)鍵點(diǎn)表示
如果您的模型返回上述方法，它們將使培訓(xùn)和評估都有效，并將使用Pycocotools中的評估腳本，該腳本可以與pip install Pycools一起安裝。
注意
對于Windows，請使用命令從gautamchitnis安裝pycocotoolspip?install?git+https://github.com/gautamchitnis/cocoapi.git@cocodataset-master#subdirectory=PythonAPI
labels上有一個(gè)注釋。該模型將0類視為背景。如果數(shù)據(jù)集不包含后臺類，則labels中不應(yīng)包含0。例如，假設(shè)只有兩個(gè)類，cat和dog，可以定義1（而不是0）來表示cat，定義2來表示dogs。例如，如果其中一個(gè)圖像同時(shí)具有兩個(gè)類，那么labels張量應(yīng)該類似于[1,2]。
此外，如果希望在訓(xùn)練期間使用縱橫比分組（以便每個(gè)批次僅包含具有類似縱橫比的圖像），則建議還實(shí)現(xiàn)get_height_和_width方法，該方法返回圖像的高度和寬度。如果不提供此方法，我們將通過__getitem__查詢數(shù)據(jù)集的所有元素，__getitem__將圖像加載到內(nèi)存中，速度比提供自定義方法慢。
為PennFudan編寫自定義數(shù)據(jù)集
讓我們?yōu)镻ennFudan數(shù)據(jù)集編寫一個(gè)數(shù)據(jù)集。下載并解壓縮zip文件后，我們有以下文件夾結(jié)構(gòu)：
PennFudanPed/
PedMasks/
FudanPed00001_mask.png
FudanPed00002_mask.png
FudanPed00003_mask.png
FudanPed00004_mask.png
...
PNGImages/
FudanPed00001.png
FudanPed00002.png
FudanPed00003.png
FudanPed00004.png下面是一對圖像和分割蒙版的一個(gè)示例
?
因此，每個(gè)圖像都有一個(gè)相應(yīng)的分割遮罩，每個(gè)顏色對應(yīng)一個(gè)不同的實(shí)例。讓我們?yōu)檫@個(gè)數(shù)據(jù)集編寫一個(gè)torch.utils.data.Dataset類。
import os
import numpy as np
import torch
from PIL import Image
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 = Image.open(img_path).convert("RGB")
# note that we haven't converted the mask to RGB,
# because each color corresponds to a different instance
# with 0 being background
mask = Image.open(mask_path)
# convert the PIL Image into a numpy array
mask = np.array(mask)
# instances are encoded as different colors
obj_ids = np.unique(mask)
# first id is the background, so remove it
obj_ids = obj_ids[1:]
# split the color-encoded mask into a set
# of binary masks
masks = mask == obj_ids[:, None, None]
# get bounding box coordinates for each mask
num_objs = len(obj_ids)
boxes = []
for i in range(num_objs):
pos = np.where(masks[i])
xmin = np.min(pos[1])
xmax = np.max(pos[1])
ymin = np.min(pos[0])
ymax = np.max(pos[0])
boxes.append([xmin, ymin, xmax, ymax])
# convert everything into a torch.Tensor
boxes = torch.as_tensor(boxes, dtype=torch.float32)
# there is only one class
labels = torch.ones((num_objs,), dtype=torch.int64)
masks = torch.as_tensor(masks, dtype=torch.uint8)
image_id = torch.tensor([idx])
area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])
# suppose all instances are not crowd
iscrowd = torch.zeros((num_objs,), dtype=torch.int64)
target = {}
target["boxes"] = boxes
target["labels"] = labels
target["masks"] = masks
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)這就是數(shù)據(jù)集的全部內(nèi)容?，F(xiàn)在，讓我們定義一個(gè)可以對此數(shù)據(jù)集執(zhí)行預(yù)測的模型。
定義模型
在本教程中，我們將使用Mask R-CNN，它基于更快的R-CNN。更快的R-CNN模型可以預(yù)測圖像中潛在對象的邊界框和類分?jǐn)?shù)。
Mask R-CNN為更快的R-CNN添加了一個(gè)額外分支，該分支還可以預(yù)測每個(gè)實(shí)例的分段掩碼。
在兩種常見情況下，可能需要修改torchvision modelzoo中的一個(gè)可用模型。第一個(gè)是當(dāng)我們想從一個(gè)預(yù)先訓(xùn)練好的模型開始，只需微調(diào)最后一層。另一個(gè)是當(dāng)我們想用另一個(gè)模型替換模型的主干時(shí)（例如，為了更快的預(yù)測）。
讓我們來看看在接下來的部分中我們將如何做一個(gè)或另一個(gè)。
1 - 從預(yù)訓(xùn)練模型微調(diào)
假設(shè)您要從在 COCO 上預(yù)先訓(xùn)練的模型開始，并希望針對您的特定類對其進(jìn)行微調(diào)。以下是一種可能的方法：
import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
## load a model pre-trained on COCO
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
## 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)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(pretrained=True).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 FasterRCNN model
model = FasterRCNN(backbone,
num_classes=2,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler)
一種PennFudan數(shù)據(jù)集的實(shí)例分割模型
在我們的例子中，我們希望從預(yù)先訓(xùn)練的模型進(jìn)行微調(diào)，因?yàn)槲覀兊臄?shù)據(jù)集非常小，所以我們將遵循方法1。 在這里，我們還想計(jì)算實(shí)例分割掩碼，因此我們將使用 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(pretrained=True)
# 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準(zhǔn)備好在自定義數(shù)據(jù)集上進(jìn)行訓(xùn)練和評估。
將所有內(nèi)容放在一起
在references/detection/中，我們有許多幫助函數(shù)來簡化訓(xùn)練和評估檢測模型。在這里，我們將使用references/detection/engine.py，references/detection/utils.py和references/detection/transforms.py。只需將references/detection下的所有內(nèi)容復(fù)制到您的文件夾中，并在此處使用它們。
讓我們?yōu)閿?shù)據(jù)擴(kuò)充/轉(zhuǎn)換編寫一些幫助函數(shù)：
import transforms as T
def get_transform(train):
transforms = []
transforms.append(T.ToTensor())
if train:
transforms.append(T.RandomHorizontalFlip(0.5))
return T.Compose(transforms)
測試forward()方法（可選）
在迭代數(shù)據(jù)集之前，最好先查看模型在樣本數(shù)據(jù)上的訓(xùn)練和推理時(shí)間內(nèi)的預(yù)期內(nèi)容。
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
dataset = PennFudanDataset('PennFudanPed', get_transform(train=True))
data_loader = torch.utils.data.DataLoader(
dataset, batch_size=2, shuffle=True, num_workers=4,
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
## For inference
model.eval()
x = [torch.rand(3, 300, 400), torch.rand(3, 500, 400)]
predictions = model(x) # Returns predictions現(xiàn)在，讓我們編寫執(zhí)行訓(xùn)練和驗(yàn)證的 main 函數(shù)：
from engine import train_one_epoch, evaluate
import utils
def main():
# train on the GPU or on the CPU, if a GPU is not available
device = torch.device('cuda') if torch.cuda.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('PennFudanPed', get_transform(train=True))
dataset_test = PennFudanDataset('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, num_workers=4,
collate_fn=utils.collate_fn)
data_loader_test = torch.utils.data.DataLoader(
dataset_test, batch_size=1, shuffle=False, num_workers=4,
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 for 10 epochs
num_epochs = 10
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!")
您應(yīng)該獲得第一個(gè)歷元的輸出：
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
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
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
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
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
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
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
Epoch: [0] Total time: 0:00:21 (0.3570 s / it)
creating index...
index created!
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
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
Test: Total time: 0:00:04 (0.0828 s / it)
Averaged stats: model_time: 0.0628 (0.0687) evaluator_time: 0.0039 (0.0064)
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.606
Average Precision (AP) @[ IoU=0.50 | area= all | maxDets=100 ] = 0.984
Average Precision (AP) @[ IoU=0.75 | area= all | maxDets=100 ] = 0.780
Average Precision (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.313
Average Precision (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.582
Average Precision (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.612
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 1 ] = 0.270
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 10 ] = 0.672
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.672
Average Recall (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.650
Average Recall (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.755
Average Recall (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.664
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.979
Average Precision (AP) @[ IoU=0.75 | area= all | maxDets=100 ] = 0.871
Average Precision (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.325
Average Precision (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.488
Average Precision (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.727
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 1 ] = 0.316
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 10 ] = 0.748
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.749
Average Recall (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.650
Average Recall (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.673
Average Recall (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.758
因此，經(jīng)過一個(gè)時(shí)期的訓(xùn)練，我們獲得了60.6的COCO式mAP和70.4的掩模mAP。 經(jīng)過 10 個(gè)紀(jì)元的訓(xùn)練，我得到了以下指標(biāo)
IoU metric: bbox
Average Precision (AP) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.799
Average Precision (AP) @[ IoU=0.50 | area= all | maxDets=100 ] = 0.969
Average Precision (AP) @[ IoU=0.75 | area= all | maxDets=100 ] = 0.935
Average Precision (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.349
Average Precision (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.592
Average Precision (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.831
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.844
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.844
Average Recall (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.400
Average Recall (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.777
Average Recall (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.870
IoU metric: segm
Average Precision (AP) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.761
Average Precision (AP) @[ IoU=0.50 | area= all | maxDets=100 ] = 0.969
Average Precision (AP) @[ IoU=0.75 | area= all | maxDets=100 ] = 0.919
Average Precision (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.341
Average Precision (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.464
Average Precision (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.788
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 1 ] = 0.303
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 10 ] = 0.799
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.799
Average Recall (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.400
Average Recall (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.769
Average Recall (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.818但是預(yù)測是什么樣的呢？讓我們在數(shù)據(jù)集中獲取一個(gè)圖像并進(jìn)行驗(yàn)證
經(jīng)過訓(xùn)練的模型預(yù)測了此圖像中的 9 個(gè)人員實(shí)例，讓我們看一下其中的幾個(gè)實(shí)例：
?
結(jié)果看起來相當(dāng)不錯！
結(jié)束語
在本教程中，您學(xué)習(xí)了如何在自定義數(shù)據(jù)集上創(chuàng)建自己的培訓(xùn)管道，例如分段模型。為此，你寫了一個(gè)torch.utils.data.Dataset類，該類返回圖像、地面真值框和分割遮罩。為了在這個(gè)新數(shù)據(jù)集上執(zhí)行轉(zhuǎn)移學(xué)習(xí)，您還利用了COCO train2017上預(yù)先培訓(xùn)的Mask R-CNN模型。
有關(guān)更完整的示例，包括多機(jī)/多gpu培訓(xùn)，請查看references/detection/train.py，它出現(xiàn)在torchvision回購協(xié)議中。
在此處下載本教程的完整源文件。
以上內(nèi)容是否對您有幫助：
更多建議：