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）