深度学习Pytorch（十）——基于torchvision的目标检测模型

深度学习Pytorch（十）——基于torchvision的目标检测模型

本节，将微调Penn-Fudan数据库中对行人检测和分割的已预先训练的Mask R-CNN模型。该数据集包含170个图像和345个行人实例。用该数据集说明如何在torchvision中使用新功能，以便在自定义数据集上训练实例分割模型

一、定义数据集

对于训练对象检测的脚本，实例分割和人员关键点检测要能够轻松支持添加新的自定义的数据，数据集应该从标准的类torch.utils.data.Dataset继承，并实现__len__和__getitem__
对于__getitem__应该返回：

1. 图像：PIL图像大小（h,w）
2. target：包含如下字段的字典
   （1）boxes(FloatTensor[N,4])：N边框坐标的格式[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以用于新的关键点表示
   如果要在training阶段使用宽高比分组，还需要实现get_high_width方法，返回图像的高度和宽度。若没有该function，将通过__getitem__查询数据集的所有元素，这种function会将图像加载到内存中，但是比提供的自定义的方法要慢

二、为PennFudan编写自定义数据集

1、下载数据集

数据集——>我在这里
下载后解压放在工作路径下，先查看一下数据集吧，结构如下：



（黑压压一大片，这没问题哈，是掩膜，哈哈哈哈哈~不要以为自己下载的不对！）
文件夹中每个文件的详细说明如下（readme中详细说明）：

2、为数据集编写类

在这里插入代码片

三、定义模型

定义一个可以预测上述数据集的模型。本节，使用Mask R-CNN，它基于Faster R-CNN。Faster R-CNN是一个模型，可以预测图像中潜在对象的边界框和类别得分。（此处贴一个Faster R-CNN的详解——>我在这里）

Mask R-CNN在Faster R-CNN中添加了一个额外的分支，预测每个实例的分割蒙版

有两种常见情况可能需要修改torchvision modelzoo中的一个可用模型。

• 想要从预先训练的模型开始的时候，微调最后一层
• 想用不同的模型替换模型的主干时（用于更快的预测）

以下是读这两种情况的处理：

Ⅰ 微调已经预训练的模型

假设你想从一个在CoCo（其实它叫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 (person)+background
# 获取分类器的输入参数的数量
in_features=model.roi_heads.box_predictor.cls_score.in_features
# 用新的头部替换预先训练好的头部
model.roi_heads.box_predictor=FastRCNNPredictor(in_features, num_classes)

Ⅱ 修改模型以添加不同的主干

#%% 修改模型以添加不同的主干
import torchvision
from torchvision.models.detection import FasterRCNN
from torchvision.models.detection.rpn import AnchorGenerator

# 加载预先训练的模型进行分类和返回
backbone=torchvision.models.mobilenet_v2(pretrained=True).features
# faster r-cnn需要知道骨干网络中的输出通道数量，对于mobilenet_v2,是1280，这里添加backbone.out_channels=1280
backbone.out_channels=1280
# 让RPN在每个空间位置生成5*3个锚点，具有5种不同的大小和3种不同的宽高比
# 有一个元组，每个特征映射可能具有不同的大小和宽高比
anchor_generator=AnchorGenerator(sizes=((32,64,128,256,512),),
                                 aspect_ratios=((0.5,1.0,2.0),))
# 定义用于执行区域裁剪的特征映射，以及重新缩放后裁剪的大小。如果主干返回Tensor,featmap_names应为[0]更一般的，主干应该返回OrderedDict[Tensor],并在feature_names中，选择要使用的功能映射
roi_pooler=torchvision.ops.MultiScaleRoIAlign(featmap_names=[0], 
                                              output_size=7,
                                              sampling_ratio=2)
# 将这些pieces放在FasterRCNN模型中
model=FastRCNNPredictor(backbone,
                        num_classes=2,
                        rpn_anchor_generator=anchor_generator,
                        box_roi_pool=roi_pooler)

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

四、整合

在**reference/detection/**中有许多辅助函数来简化训练和评估检测模型。此处将使用reference/detection/engine.py，reference/detection/util.py和reference/detection/transforms.py。只需要将他们复制到文件夹并在工作路径使用——>我在这里

1、为数据扩充/转换编写辅助函数

#为数据扩充/转换编写辅助函数
import sys
sys.path.append("D:\Python\Pytorch")
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)

2、编写执行训练和验证的主要功能

#编写执行训练和验证的主要功能
from engine import train_one_epoch,evaluate
import utils

def main():
    # 在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,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个epoch
    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 all")

在第一个epoch训练后可以得到一下结果：

因此，在一个epoch训练之后，获得了CoCo-style mAP为64.8，mask mAP为67.6——如下图

经过训练10个epoch后，得到如下指标：

小结明天再写，太晚了，我要赶紧回宿舍洗洗睡了。明天见，猿友们~