專欄目錄:
-
PyTorch簡介
-
PyTorch環境搭建
-
PyTorch自動微分
-
PyTorch神經網絡
第三章:PyTorch之入門強化
第四章:PyTorch之圖像篇
-
微調基於torchvision 0.3的目標檢測模型
-
微調TorchVision模型
-
空間變換器網絡
-
使用PyTorch進行神經傳遞
-
生成對抗示例
-
使用ONNX將模型轉移至Caffe2和移動端
第五章:PyTorch之文本篇
-
聊天機器人教程
-
使用字符級RNN生成名字
-
使用字符級RNN進行名字分類
-
在深度學習和NLP中使用Pytorch
-
使用Sequence2Sequence網絡和注意力進行翻譯
第六章:PyTorch之生成對抗網絡
第七章:PyTorch之強化學習
在這篇文章中,我們將微調在 Penn-Fudan 數據庫中對行人檢測和分割的已預先訓練的 Mask R-CNN 模型。它包含170個圖像和345個行人實例,我們 將用它來說明如何在 torchvision 中使用新功能,以便在自定義數據集上訓練實例分割模型。
1.定義數據集
對於訓練對象檢測的引用腳本,實例分割和人員關鍵點檢測要求能夠輕松支持添加新的自定義數據。數據集應該從標准的類torch.utils.data.Dataset
繼承而來,並實現_len
和_getitem_
我們要求的唯一特性是數據集的__getitem__
應該返回:* 圖像:PIL圖像大小(H,W) * 目標:包含以下字段的字典
<1> boxes(FloatTensor[N,4])
:N邊框(bounding boxes)坐標的格式[x0,x1,y0,y1],取值范圍是0到W,0到H。
<2> labels(Int64Tensor[N])
:每個邊框的標簽。
<3> image_id(Int64Tensor[1])
:圖像識別器,它應該在數據集中的所有圖像中是唯一的,並在評估期間使用。
<4> area(Tensor[N])
:邊框的面積,在使用COCO指標進行評估時使用此項來分隔小、中和大框之間的度量標准得分。
<5> iscrowed(UInt8Tensor[N,H,W])
:在評估期間屬性設置為iscrowed=True
的實例會被忽略。
<6> (可選)masks(UInt8Tesor[N,H,W])
:每個對象的分段掩碼。
<7> (可選)keypoints (FloatTensor[N, K, 3]
:對於N個對象中的每一個,它包含[x,y,visibility]格式的K個關鍵點,用 於定義對象。visibility = 0
表示關鍵點不可見。請注意,對於數據擴充,翻轉關鍵點的概念取決於數據表示,您應該調整 reference/detection/transforms.py 以用於新的關鍵點表示。
如果你的模型返回上述方法,它們將使其適用於培訓和評估,並將使用 pycocotools 的評估腳本。
此外,如果要在訓練期間使用寬高比分組(以便每個批次僅包含具有相似寬高比的圖像),則建議還實現get_height_and_width
方法, 該方法返回圖像的高度和寬度。如果未提供此方法,我們將通過__getitem__
查詢數據集的所有元素,這會將圖像加載到內存中,但比提供自定義方法時要慢。
2.為 PennFudan 編寫自定義數據集
2.1 下載數據集
下載並解壓縮zip文件后,我們有以下文件夾結構:
PennFudanPed/
PedMasks/
FudanPed00001_mask.png
FudanPed00002_mask.png
FudanPed00003_mask.png
FudanPed00004_mask.png
...
PNGImages/
FudanPed00001.png
FudanPed00002.png
FudanPed00003.png
FudanPed00004.png
下面是一個圖像以及其分割掩膜的例子:
因此每個圖像具有相應的分割掩膜,其中每個顏色對應於不同的實例。讓我們為這個數據集寫一個torch.utils.data.Dataset
類。
2.2 為數據集編寫類
import os
import numpy as np
import torch
from PIL import Image
class PennFudanDataset(object):
def __init__(self, root, transforms):
self.root = root
self.transforms = transforms
# 下載所有圖像文件,為其排序
# 確保它們對齊
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 ad 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")
# 請注意我們還沒有將mask轉換為RGB,
# 因為每種顏色對應一個不同的實例
# 0是背景
mask = Image.open(mask_path)
# 將PIL圖像轉換為numpy數組
mask = np.array(mask)
# 實例被編碼為不同的顏色
obj_ids = np.unique(mask)
# 第一個id是背景,所以刪除它
obj_ids = obj_ids[1:]
# 將顏色編碼的mask分成一組
# 二進制格式
masks = mask == obj_ids[:, None, None]
# 獲取每個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])
# 將所有轉換為torch.Tensor
boxes = torch.as_tensor(boxes, dtype=torch.float32)
# 這里僅有一個類
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])
# 假設所有實例都不是人群
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)
3.定義模型
現在我們需要定義一個可以上述數據集執行預測的模型。在本教程中,我們將使用 Mask R-CNN, 它基於 Faster R-CNN。Faster R-CNN 是一種模型,可以預測圖像中潛在對象的邊界框和類別得分。
Mask R-CNN 在 Faster R-CNN 中添加了一個額外的分支,它還預測每個實例的分割蒙版。
有兩種常見情況可能需要修改torchvision modelzoo
中的一個可用模型。第一個是我們想要從預先訓練的模型開始,然后微調最后一層。另一種是當我們想要用不同的模型替換模型的主干時(例如,用於更快的預測)。
下面是對這兩種情況的處理。
-
1 微調已經預訓練的模型 讓我們假設你想從一個在COCO上已預先訓練過的模型開始,並希望為你的特定類進行微調。這是一種可行的方法:
import torchvision from torchvision.models.detection.faster_rcnn import FastRCNNPredictor # 在COCO上加載經過預訓練的預訓練模型 model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True) # replace the classifier with a new one, that has # 將分類器替換為具有用戶定義的 num_classes的新分類器 num_classes = 2 # 1 class (person) + background # 獲取分類器的輸入參數的數量 in_features = model.roi_heads.box_predictor.cls_score.in_features # 用新的頭部替換預先訓練好的頭部 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 # 加載預先訓練的模型進行分類和返回 # 只有功能 backbone = torchvision.models.mobilenet_v2(pretrained=True).features # FasterRCNN需要知道骨干網中的輸出通道數量。對於mobilenet_v2,它是1280,所以我們需要在這里添加它 backbone.out_channels = 1280 # 我們讓RPN在每個空間位置生成5 x 3個錨點 # 具有5種不同的大小和3種不同的寬高比。 # 我們有一個元組[元組[int]] # 因為每個特征映射可能具有不同的大小和寬高比 anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512),), aspect_ratios=((0.5, 1.0, 2.0),)) # 定義一下我們將用於執行感興趣區域裁剪的特征映射,以及重新縮放后裁剪的大小。 # 如果您的主干返回Tensor,則featmap_names應為[0]。 # 更一般地,主干應該返回OrderedDict [Tensor] # 並且在featmap_names中,您可以選擇要使用的功能映射。 roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=[0], output_size=7, sampling_ratio=2) # 將這些pieces放在FasterRCNN模型中 model = FasterRCNN(backbone, num_classes=2, rpn_anchor_generator=anchor_generator, box_roi_pool=roi_pooler)
3.1 PennFudan 數據集的實例分割模型
在我們的例子中,我們希望從預先訓練的模型中進行微調,因為我們的數據集非常小,所以我們將遵循上述第一種情況。
這里我們還要計算實例分割掩膜,因此我們將使用 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):
# 加載在COCO上預訓練的預訓練的實例分割模型
model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True)
# 獲取分類器的輸入特征數
in_features = model.roi_heads.box_predictor.cls_score.in_features
# 用新的頭部替換預先訓練好的頭部
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
# 現在獲取掩膜分類器的輸入特征數
in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
hidden_layer = 256
# 並用新的掩膜預測器替換掩膜預測器
model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask,
hidden_layer,
num_classes)
return model
就是這樣,這將使模型准備好在您的自定義數據集上進行訓練和評估。
4.整合
在references/detection/
中,我們有許多輔助函數來簡化訓練和評估檢測模型。在這里,我們將使用 references/detection/engine.py
,references/detection/utils.py
和references/detection/transforms.py
。只需將它們復制到您的文件夾並在此處使用它們。
4.1 為數據擴充/轉換編寫輔助函數:
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)
4.2 編寫執行訓練和驗證的主要功能
from engine import train_one_epoch, evaluate
import utils
def main():
# 在GPU上訓練,若無GPU,可選擇在CPU上訓練
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
# 我們的數據集只有兩個類 - 背景和人
num_classes = 2
# 使用我們的數據集和定義的轉換
dataset = PennFudanDataset('PennFudanPed', get_transform(train=True))
dataset_test = PennFudanDataset('PennFudanPed', get_transform(train=False))
# 在訓練和測試集中拆分數據集
indices = torch.randperm(len(dataset)).tolist()
dataset = torch.utils.data.Subset(dataset, indices[:-50])
dataset_test = torch.utils.data.Subset(dataset_test, indices[-50:])
# 定義訓練和驗證數據加載器
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)
# 使用我們的輔助函數獲取模型
model = get_model_instance_segmentation(num_classes)
# 將我們的模型遷移到合適的設備
model.to(device)
# 構造一個優化器
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)
# 和學習率調度程序
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
step_size=3,
gamma=0.1)
# 訓練10個epochs
num_epochs = 10
for epoch in range(num_epochs):
# 訓練一個epoch,每10次迭代打印一次
train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq=10)
# 更新學習速率
lr_scheduler.step()
# 在測試集上評價
evaluate(model, data_loader_test, device=device)
print("That's it!")
在第一個epoch訓練后可以得到下面的結果:
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
因此,在一個epoch訓練之后,我們獲得了COCO-style mAP為60.6,並且mask mAP為70.4。
經過訓練10個epoch后,我得到了以下指標:
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
但預測結果如何呢?讓我們在數據集中拍攝一張圖像並進行驗證。
訓練的模型預測了此圖像中的9個人物,讓我們看看其中的幾個,由下圖可以看到預測效果很好。
5.總結
在本教程中,您學習了如何在自定義數據集上為實例分段模型創建自己的訓練管道。為此,您編寫了一個torch.utils.data.Dataset
類, 它返回圖像以及地面實況框和分割掩碼。您還利用了在COCO train2017上預訓練的Mask R-CNN模型,以便對此新數據集執行傳輸學習。
有關包含multi-machine / multi-gpu training的更完整示例,請檢查 torchvision 存儲庫中的references/detection/train.py
。
歡迎關注磐創博客資源匯總站:
http://docs.panchuang.net/
歡迎關注PyTorch官方中文教程站:
http://pytorch.panchuang.net/