在這篇演示中，我們將在資料集Penn-Fudan Database for Pedestrian Detection and Segmentation上微調一個預訓練模型Mask R-CNN。此資料集包含345個行人例項（instance）的170張影象，我們將使用它來闡述如何使用torchvision在自定義資料集上訓練例項分割模型的一些新特性。

資料集說明

這一資料集包含報告在實驗Object Detection Combining Recognition and Segmentation提出的行人檢測的影象。每張影象來自於校園以及城市街道。我們的目標是影象上的行人。每一張影象上將會有至少一個行人。

在資料集中的標註了的行人的高度在[180, 390]畫素。所有標註的行人都是直立狀態。

總共有170張影象標註了345個標註行人，其中96張影象拍攝於賓夕法尼亞大學附近，另外74張影象拍攝於復旦大學附近。

定義資料集

在用於目標檢測、例項分割以及人物關鍵點檢測的參考指令碼允許輕鬆地新增新的自定義資料集。資料集應該整合標準類torch.utils.data.Dataset，並且實現__len__和__getitem__方法。

唯一我們需要的要求就是__getitem__應該返回一下內容：

• boxes (FloatTensor[N, 4]：N個邊界邊框的座標以[x0, y0, x1, y1]的格式給出，並且範圍是從0
  到W和0到H
• labels (Int64Tensor[N]：每個邊界邊框的label，0總是代表背景類別
• image_id (Int64Tensor[1])：影象的識別符號，它在影象資料集中應該是唯一的，並且用於評估階段
• area (Tensor[N])：每個邊界框的面積（area），這將在評估階段計算COCO指標，並將其分為small、medium和large
• iscrowd (UInt8Tensor[N])：iscrowd=True的例項在評估階段將被忽略
• （可選）masks (UInt8Tensor[N, H, W])：對於每一個目標的分割遮罩
• （可選）keypoints (FloatTensor[N, K, 3])
  ：對於N個目標，它包含以[x, y, visibility]格式的K個關鍵點（keypoints）定義。visibility=0表示關鍵點不可見。注意：對於資料增強（data augmentation），翻轉一個關鍵點（keypoint）的概念是獨立於資料表示（data representation）的，並且你大概率要為你的新關鍵點表示（keypoint representation）調整references/detection/transforms.py檔案

如果你的模型返回上面的方法，那麼它將會在訓練和評估階段奏效，並且將使用pycocotools的評估指令碼，這個工具可以使用pip install pycocotools安裝。

注意：對於WIndows系統，請從gautamchitnis中安裝pycocotools，使用如下命令：

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

對於labels需要注意，模型會將類別0認作位背景類別。如果你的資料集沒有包含背景類，你不應該在你的labels中出現0。舉個例子，假設你只有2個類別，貓和狗，你可以定義1（而不是0）代表貓以及2代表狗。所以，如果有個影象包含兩個類別，你的labelstensor應該像這樣[1,2]。

另外，如果你想在訓練階段使用aspect ration grouping（也就是每一個數據批量（data batch）的圖片都有相似的長寬比），那麼推薦做法就是實現get_height_and_width方法，這個方法返回圖片的長度和寬度。如果這個方法沒有被提供，我們將使用__getitem__查詢資料集的所有元素，這將會載入圖片到記憶體，並且會比提供了自定義方法的速度慢。

為PennFudan自定義資料集

下面開始為PennFudan資料編寫一個Dataset類，在下載並提取出zip檔案後，我們將有如下所示的檔案結構：

PennFudanPed/
  PedMasks/
    FudanPed00001_mask.png
    FudanPed00002_mask.png
    FudanPed00003_mask.png
    FudanPed00004_mask.png
    ...
  PNGImages/
    FudanPed00001.png
    FudanPed00002.png
    FudanPed00003.png
    FudanPed00004.png

下面是一對影象和對應的分割遮罩的示例：

對於每一張圖片都有對應的分割遮罩，其中每個顏色對應不同的例項（instance），讓我們開始為此資料集編寫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
        # 載入所有圖片檔案，將它們排序
        # 確保圖片與遮罩一一對應
        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):
        # 載入圖片和遮罩
        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")
        # 注意我們並沒有把遮罩也轉換為RGB
        # 因為每一個顏色都對應不同的例項
        # 0表示背景
        mask = Image.open(mask_path)
        # 將PIL影象轉換成numpy的array
        mask = np.array(mask)
        # 例項以不同的顏色編碼
        obj_ids = np.unique(mask)
        # 第一個id是背景，所以移除它
        obj_ids = obj_ids[1:]
        
        # 將編碼的顏色遮罩分割為二元遮罩
        masks = mask == obj_ids[:, None, None]
        
        # 對於每個遮罩，獲取他們的邊界邊框座標
        num_objs = len(obj_ids)
        boxes = []
        for i in range(num_objs):
            # np.where 返回的為二維座標
            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)
        # 僅只有1個類別：行人類
        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])
        # 假設所有的例項都不是人群（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)

這就是資料集的所有內容。接下來讓我們定義可以在這個資料集上進行預測的模型。

定義模型

在本文，我們將使用基於Faster R-CNN模型之上的Mask R-CNN。Faster R-CNN是可以同時預測在影象中潛在的目標的邊界邊框和類別分數的模型。

Mask R-CNN添加了額外的分支到Faster R-CNN上，此模型對於每一個例項也預測了分割遮罩。

有兩種常見的情況，其中一種就是在torchvision樂園中使用想要修改可用的模型之一。首先，當我們從一個預訓練模型開始，簡單地微調一下最後一層即可。另一種方法便是當我們想使用不同的架構替換掉模型的骨幹（例如，為了更快的預測）。

讓我們在接下來的章節看一下我們將會如何做到以上兩點。

1 一個預訓練模型的微調

假設我們希望從一個在COCO上預訓練的模型開始，並且對於我們特定的類別對其進行微調。

以下是一種可能的做法：

import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor

# 在COCO上載入預訓練的模型
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)

# 替換分類器，由使用者定義num_classes
num_classes = 2  # 類別 1（行人）+背景類別
# 獲得分類器輸入特徵的數量
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

# 為分類任務載入預訓練模型，並且僅返回特徵（features）
backbone = torchvision.models.mobilenet_v2(pretrained=True).features
# Faster RCNN需要知道骨幹中輸出通道的數量
# 對於mobilenet_v2是1280，所以我們需要為其新增
backbone.out_channels = 1280

# 讓我們使用RPN在每一個空間位置生成5x3個錨點（anchor）
# 其中5種不同尺寸和3種不同長寬比
# 我們需要一個Tuple[Tuple[int]]，因為每一個特徵圖（feature map）潛在地會有不同的尺寸和長寬比
anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512),), aspect_ratios=((0.5, 1.0, 2.0),))

# 讓我們定義什麼是特徵圖（feature maps），我們將用其執行裁剪感興趣區域ROI（region of interest），
# 以及在重新縮放之後裁剪尺寸。
# 如果你的架構返回一個Tensor，featmap_names被期待為[0]。通常來講，骨幹應該返回一個
# OrderedDict[Tensor]，並且在featmap_names中，你可以選擇使用哪一個features maps。
roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=['0'], output_size=7, sampling_ratio=2)

# 將部件放進FasterRCNN模型中
model = FasterRCNN(backbone, num_classes=2,
                   rpn_anchor_generator=anchor_generator,
                   box_roi_pool=roi_pooler)

使用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):
    # 載入一個在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

至此，這將使得model將在你自定義的訓練集上準備開始訓練和評估了。

整合

在references/detection中，我們有很多幫助函式來簡化訓練和評估檢測模型。在這裡，我們將使用references/detection/engin.py，references/detection/utils.py以及references/detection/transforms.py。直接將這些檔案複製到你的檔案目錄下並使用它。

讓我們編寫用於資料增強（data augmentation）\transformation的幫助函式：

import transforms as T
from torchvision.transforms import ToTensor

def get_transform(train):
    transforms = []
    transforms.append(ToTensor())
    if train:
        transforms.append(T.RandomHorizontalFlip(0.5))
    return T.Compose(transforms)

測試forward()方法

在迭代整個資料集之前，讓我們看一下在部分資料上模型期待的訓練和測試時間。

import utils
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)
# 訓練
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)  # 返回損失值（loss）和檢測
# 測試
model.eval()
x = [torch.rand(3, 300, 400), torch.rand(3, 500, 400)]
predictions = model(x)

現在，讓我們編寫主函式，用來執行訓練和測試：

from engine import train_one_epoch, evaluate
import utils


def main():
    # 若GPU可用便在GPU上訓練
    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

    # 我們的資料集只有2個類別：背景和行人
    num_classes = 2
    # 使用我們的資料集和定義的transformations
    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)
    # 使用learning rate scheduler
    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):
        # 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!")

最終的預測效果，可以使用資料集中的一張圖片來驗證

訓練的模型預測了9個例項，讓我們看下其中的幾個例項：

在本文，對於例項分割模型，你已經學會了如何建立你自己在自定義資料集上的訓練流程。為此，你編寫了使用者反悔圖片和真實邊界以及分割遮罩的torch.utils.data.Dataset類。為了在執行遷移學習和新資料集上，你還利用了預訓練在COCO train2017上的Mask R-CNN模型。

對於更完整的示例，包含多機器、多顯示卡訓練，檢視torchvision倉庫中的references/detection/train.py即可。