《PyTorch》Part7 PyTorch之Mask R-CNN图像实例分割模型

《PyTorch》Part7 PyTorch之Mask R-CNN图像实例分割模型
GPU端运行参考:
手把手教你训练自己的Mask R-CNN图像实例分割模型(PyTorch官方教程)

运行环境:

Win10
torch 1.6.0+cu101
torchvision 0.7.0+cu101
显卡: NVIDIA1050 内存:2GB
特别注意需要安装:
Cython
pycocotools

相关注意项:

1.源码下载地址:

http://pytorch.org/tutorials/intermediate/torchvision_tutorial.html

数据集(下载完了放在当前工程文件夹下,解压出来):

网页版: https://www.cis.upenn.edu/~jshi/ped_html/
下载直达:https://www.cis.upenn.edu/~jshi/ped_html/PennFudanPed.zip

maskrcnn_resnet50_fpn模型地址:

https://download.pytorch.org/models/maskrcnn_resnet50_fpn_coco-bf2d0c1e.pth

2.下载依赖的辅助函数:https://github.com/pytorch/vision/tree/master/references/detection

在references/detection/中,有许多辅助函数来简化训练和评估检测模型。在这里,将使用
references/detection/engine.py,references/detection/utils.py和references/detection/transforms.py。
只需将它们复制到自己的文件夹并在此处使用它们。

3.由于缺少pycocotools包引发的错误,而pycocotools包需要cython和C编译器,所以在安装cython之前需要先安装Cython

pip install “git+https://github.com/philferriere/cocoapi.git#egg=pycocotools^&subdirectory=PythonAPI”

4.安装pycocotools
windows环境下安装pycocotools用:

pip install pycocotools

5.本人GTX1050,2G内存,运行时出现分配内存不足,导致无法训练,故设置num_epochs = 1,默认为10
6.同第4点原因,batch_size = 1
7.同第4点原因,num_workers = 1
8.解决:RuntimeError: CUDA out of memory. Tried to allocate 2.00 MiB
https://blog.csdn.net/weixin_42253689/article/details/109473245?utm_medium=distribute.pc_relevant.none-task-blog-title-2&spm=1001.2101.3001.4242
9.由于报显卡内存不足,暂时没管梯度反向传播,默认的改成了在CPU中运行(跑了好久…没办法)

PyTorch里的requires_grad、volatile及no_grad

requires_grad=True 要求计算梯度

requires_grad=False 不要求计算梯度

with torch.no_grad()或者@torch.no_grad()中的数据不需要计算梯度,也不会进行反向传播

(torch.no_grad()是新版本pytorch中volatile的替代)

1.数据集

注意项中的第1点操作,下载数据集
数据集组成:

PennFudanPed/
PedMasks/
FudanPed00001_mask.png
FudanPed00002_mask.png
FudanPed00003_mask.png
FudanPed00004_mask.png

PNGImages/
FudanPed00001.png
FudanPed00002.png
FudanPed00003.png
FudanPed00004.png
‘’’

查看:

from PIL import Image
 
Image.open('PennFudanPed/PNGImages/FudanPed00001.png')
 
mask = Image.open('PennFudanPed/PedMasks/FudanPed00001_mask.png')
 
mask.putpalette([
    0, 0, 0, # black background
    255, 0, 0, # index 1 is red
    255, 255, 0, # index 2 is yellow
    255, 153, 0, # index 3 is orange
])
 
mask

运行结果:

每一张图像都有对应的mask标注,不同的颜色表示不同的实例。在训练模型之前,需要写好数据集的载入接口。

# 1.为数据集编写类
import os
import numpy as np
import torch
from PIL import Image

class PennFudanDataset(object):
    def __init__(self, root, transforms=None):
        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"))))

    # 获取第idx张图片和对应的mask标签
    def __getitem__(self, idx):
        # Load images and 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:] # 取第1~N项

        # 将颜色编码的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])
        # 计算包围盒面积 deltaY * deltaX
        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) # 获取图片的数量

检查一下上面接口返回的dataset的内部结构

dataset = PennFudanDataset('PennFudanPed/')
dataset[0]

运行结果:

dataset返回了一个PIL.Image以及一个dictionary,包含boxes、labels和masks等域,这都是训练的时候网络需要用到的。

2.定义模型

Mask R-CNN是基于Faster R-CNN改造而来的。Faster R-CNN用于预测图像中潜在的目标框和分类得分,而Mask R-CNN在此基础上加了一个额外的分支,用于预测每个实例的分割mask。

有两种方式来修改torchvision modelzoo中的模型,以达到预期的目的。

第一种,采用预训练的模型,在修改网络最后一层后finetune。
第二种,根据需要替换掉模型中的骨干网络,如将ResNet替换成MobileNet等。

2.1 方法1:

从预先训练的模型开始,然后微调最后一层下面是微调已经训练的模型,假设想从一个在COCO上预先训练过的模型开始,并希望为你的特定类进行微调
场景:利用COCO上预训练的模型,为指定类别的任务进行finetune。

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 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 方法2:

当我们想用不同的模型替换模型的主干网络时,例如:用于更快的预测MobileNet



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*3个锚点
# 具有5种不同的大小和3种不同的宽高比
# 我们有一个tuple[tuple[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)

2.3 定义Mask R-CNN模型

本文的目的是在PennFudan数据集上训练Mask R-CNN实例分割模型,即上述第一种情况,直接使用torchvision.models.detection中已有的官方网络定义和接口的文件。

# 在这个例子中,我们希望从预先训练的模型中进行微调,因为我们的
# 数据集非常小,所以我们将遵循上述的第一种情况
# 这里我们还要计算实例分割Mask,因此我们将使用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
    hidder_layer = 256
    # 并用新的掩膜预测器替换掩膜预测器
    model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask,
                                                       hidder_layer,
                                                       num_classes)

    return model

模型定义好了,接下来可以在PennFudan数据集进行训练和测试了。

3.训练模型

3.1 准备工作

注意事项第2点,下载依赖的辅助函数。下载完后,将references/detection文件夹下的所有文件拷贝过来,官方说是eferences/detection/engine.py、references/detection/utils.py、references/detection/transforms.py,实际上整个文件夹下的文件都需要

3.2 数据增强/转换

# 为数据扩充/转换编写辅助函数
import utils
import transforms as T
from engine import train_one_epoch, evaluate

def get_transform(train):
    transforms = []
    # converts the image, a PIL image, into a PyTorch Tensor
    transforms.append(T.ToTensor())
    if train:
        # during training, randomly flip the training images
        # and ground-truth for data augmentation
        transforms.append(T.RandomHorizontalFlip(0.5))
    return T.Compose(transforms)

3.3 开始训练模型

至此,数据集、模型、数据增强的部分都已经写好。在模型初始化、优化器及学习率调整策略选定后,就可以开始训练了。这里,设置模型训练10个epochs(实际由于GPU不给力,转成了CPU运行,同时将epoch设置成了1),并且在每个epoch完成后在测试集上对模型的性能进行评价。
其中batch_size =1,num_workers=1

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

def main():
    # 在GPU上训练,若无GPU,可选择在CPU上训练
    device = torch.device('cpu') if torch.cuda.is_available() else torch.device('gpu')

    # 我们的数据集只有两个类 - 背景和人
    num_calsses = 2
    # 使用我们的数据集和定义的转换
    dataset = PennFudanDataset('PennFudanPed', get_transform(train=True))
    dataset_test = PennFudanDataset('PennFudanPad', 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=1, shuffle=True, num_workers=1, # 4
        collate_fn=utils.collate_fn)

    data_loader_test = torch.utils.data.DataLoader(
        dataset_test, batch_size=1, shuffle=False, num_workers=1, # 4
        collate_fn=utils.collate_fn)

    # 使用我们的辅助函数获取模型
    model = get_model_instance_segmentation(num_calsses)

    # 将我们的模型迁移到合适的设备(device)
    model.to(device)

    # 构造一个随机梯度下降SGD优化器
    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 = 1 #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!")

4.测试模型

模型已经训练好了,来检查一下模型在测试图像上预测的结果。

# pick one image from the test set
img, _ = dataset_test[0]
 
# put the model in evaluation mode
model.eval()
with torch.no_grad():
    prediction = model([img.to(device)])

这里输出的prediction中,包含了在图像中预测出的boxes、labels、masks和scores等信息。输出结果:

接下来,将测试图像及对应的预测结果可视化出来,看看效果如何:

Image.fromarray(img.mul(255).permute(1, 2, 0).byte().numpy())
 
Image.fromarray(prediction[0]['masks'][0, 0].mul(255).byte().cpu().numpy())

训练自己的Mask R-CNN模型就完成了。

相关推荐
©️2020 CSDN 皮肤主题: 编程工作室 设计师:CSDN官方博客 返回首页