Pytorch 目標檢測（Faster RCNN、Mask RCNN）

前言（必讀）

最近做目標檢測，然後記錄一下 Faster RCNN、Mask RCNN來做目標檢測踩得那些坑。

首先，本文並不是利用Pytorch從頭去實現Faster RCNN、Mask RCNN這兩個結構的文章。如果有意向去從頭實現並瞭解每一步細節可以看看下面這些視頻和博客：

來自B站的兩位大佬講解

1. 大佬一：視頻 博客 GitHub
2. 大佬二：視頻 博客 GitHub

上面都是利用pytorch從原理到具體實現的步驟。

不過本文主要還是利用Pytorch中的Torchvision.model中的Faster RCNN、Mask RCNN來實現遷移學習。

關於如何利用遷移學習來訓練自己的數據集，這裏也給出兩個超讚的教程：

1. 教程一：TORCHVISION 目標檢測網絡微調
2. 教程二：手把手教你訓練自己的Mask R-CNN圖像實例分割模型（PyTorch官方教程）

看完以上兩個教程，基本上利用Pytorch中的Torchvision.model中的Faster RCNN、Mask RCNN來實現遷移學習也基本上沒問題了。

下面介紹採坑：

1. 安裝 pycocotools

在整個過程中你要安裝 pycocotools，主要用到其中的IOU計算的庫來評價模型的性能。但是這個安裝並不簡單，看了很多文章相對來說只有這篇文章最符合。（Windows下安裝 pycocotools）

2. 定義 Faster RCNN、Mask RCNN 模型

爲什麼要說這個呢？因爲如果你不是很有錢，或者公司有點摳買不起一張8G以上的顯卡，不改動就訓練這兩個網絡你基本上不可能成功。懂？財大氣粗可以忽略……

因爲本人就用的垃圾顯卡（兩張P600，3G內存），訓練Faster RCNN、Mask RCNN 這兩個網絡不要想着使用多GPU運行，我看了GitHub說了Faster RCNN、Mask RCNN暫時不支持多GPU運行。

幸運的是，在改動一些參數之後就可以完美運行。

Mask R-CNN是基於Faster R-CNN改造而來的。Faster R-CNN用於預測圖像中潛在的目標框和分類得分.

而Mask R-CNN在此基礎上加了一個額外的分支，用於預測每個實例的分割mask。

有兩種方式來修改torchvision modelzoo中的模型，以達到預期的目的。第一種，採用預訓練的模型，在修改網絡最後一層後finetune。第二種，根據需要替換掉模型中的骨幹網絡，如將ResNet替換成MobileNet等。

2.1 微調一個預訓練好的Faster RCNN模型

假設你想從一個在COCO上預先訓練過的模型開始，並想針對你的特定類對它進行微調。下面有一種可行的方法：

import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
 
# 在COCO上加載經過預訓練的預訓練模型
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
#這個操作你是真的要有固定參數
for param in model.parameters():
    param.requires_grad = False
    
# 將分類器替換爲具有用戶定義的 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.2 替換 Faster RCNN 模型的骨幹網絡

場景：替換掉模型的骨幹網絡。舉例來說，默認的骨幹網絡（ResNet-50）對於某些應用來說可能參數過多不易部署，可以考慮將其替換成更輕量的網絡（如MobileNet）。

import torchvision
from torchvision.models.detection import FasterRCNN
from torchvision.models.detection.rpn import AnchorGenerator
 
# 加載用於分類的預先訓練的模型並僅返回features
backbone = torchvision.models.mobilenet_v2(pretrained=True).features
# FasterRCNN需要知道骨幹網中的輸出通道數量。對於mobilenet_v2，它是1280，所以我們需要在這裏添加它
backbone.out_channels = 1280
 
# 我們讓RPN在每個空間位置生成5 x 3個Anchors(具有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],#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)

2.3 微調一個預訓練好的Mask RCNN模型（本文使用）

import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor
 
# load an instance segmentation model pre-trained on COCO
model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True)
for param in model.parameters():
    param.requires_grad = False

num_classes = 2
# get the 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)

2.4 替換 Mask RCNN 模型的骨幹網絡

import torchvision
from torchvision.models.detection import FasterRCNN
from torchvision.models.detection.rpn import AnchorGenerator
 
# 加載用於分類的預先訓練的模型並僅返回features
backbone = torchvision.models.mobilenet_v2(pretrained=True).features
# FasterRCNN需要知道骨幹網中的輸出通道數量。對於mobilenet_v2，它是1280，所以我們需要在這裏添加它
backbone.out_channels = 1280
 
# 我們讓RPN在每個空間位置生成5 x 3個Anchors(具有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)
mask_roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=[0],
                                                     output_size=14,
                                                     sampling_ratio=2)

model = MaskRCNN(backbone,
                 num_classes=2,
                 rpn_anchor_generator=anchor_generator,
                 box_roi_pool=roi_pooler,
                 mask_roi_pool=mask_roi_pooler)

3. 訓練過程

訓練過程總的來說沒什麼，平平無奇，但還是總結一下幾個我遇到的坑。

1. batch_size不宜過大，過大GPU吃不消
2. num_workers設置爲0，我也不知道我設置成其他數會報錯
3. 學習率lr不宜設置太小，可以循序漸進
4. 不一定要使用lr_scheduler，可以一直保持一個固定的學習率（我的0.01）