微调torchvision 0.3的目标检测模型

微调torchvision 0.3的目标检测模型

本文将微调在 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。

可以下载本教程的完整源文件。