Mask-RCNN : 代码使用

Posted @BangBang

tags:

篇首语:本文由小常识网(cha138.com)小编为大家整理,主要介绍了Mask-RCNN : 代码使用相关的知识,希望对你有一定的参考价值。

Mask- RCNN原理及网络详解,参见:Mask- RCNN网络详解

1. Mask-RCNN代码使用

该项目参考自pytorch官方torchvision模块中的源码(使用pycocotools处略有不同)

环境配置

  • Python3.6/3.7/3.8
  • Pytorch1.10或以上
  • pycocotools(Linux:pip install pycocotools; Windows:pip install pycocotools-windows(不需要额外安装vs))
  • Ubuntu或Centos(不建议Windows)
  • 最好使用GPU训练
  • 详细环境配置见requirements.txt

文件结构

  ├── backbone: 特征提取网络
  ├── network_files: Mask R-CNN网络
  ├── train_utils: 训练验证相关模块(包括coco验证相关)
  ├── my_dataset_coco.py: 自定义dataset用于读取COCO2017数据集
  ├── my_dataset_voc.py: 自定义dataset用于读取Pascal VOC数据集
  ├── train.py: 单GPU/CPU训练脚本
  ├── train_multi_GPU.py: 针对使用多GPU的用户使用
  ├── predict.py: 简易的预测脚本,使用训练好的权重进行预测
  ├── validation.py: 利用训练好的权重验证/测试数据的COCO指标,并生成record_mAP.txt文件
  └── transforms.py: 数据预处理(随机水平翻转图像以及bboxes、将PIL图像转为Tensor)

预训练权重下载

数据集

本项目使用的有COCO2017数据集和Pascal VOC2012数据集

COCO2017数据集

├── coco2017: 数据集根目录
     ├── train2017: 所有训练图像文件夹(118287)
     ├── val2017: 所有验证图像文件夹(5000)
     └── annotations: 对应标注文件夹
              ├── instances_train2017.json: 对应目标检测、分割任务的训练集标注文件
              ├── instances_val2017.json: 对应目标检测、分割任务的验证集标注文件
              ├── captions_train2017.json: 对应图像描述的训练集标注文件
              ├── captions_val2017.json: 对应图像描述的验证集标注文件
              ├── person_keypoints_train2017.json: 对应人体关键点检测的训练集标注文件
              └── person_keypoints_val2017.json: 对应人体关键点检测的验证集标注文件夹

Pascal VOC2012数据集

VOCdevkit
    └── VOC2012
         ├── Annotations               所有的图像标注信息(XML文件)
         ├── ImageSets
         │   ├── Action                人的行为动作图像信息
         │   ├── Layout                人的各个部位图像信息
         │   │
         │   ├── Main                  目标检测分类图像信息
         │   │     ├── train.txt       训练集(5717)
         │   │     ├── val.txt         验证集(5823)
         │   │     └── trainval.txt    训练集+验证集(11540)
         │   │
         │   └── Segmentation          目标分割图像信息
         │         ├── train.txt       训练集(1464)
         │         ├── val.txt         验证集(1449)
         │         └── trainval.txt    训练集+验证集(2913)
         │
         ├── JPEGImages                所有图像文件
         ├── SegmentationClass         语义分割png图(基于类别)
         └── SegmentationObject        实例分割png图(基于目标)

训练方法

  • 确保提前准备好数据集
  • 确保提前下载好对应预训练模型权重
  • 确保设置好--num-classes--data-path (Pascal voc不包含背景,num_class=20, 对于coco数据集,不包含背景num_class =90)
  • 若要使用单GPU训练直接使用train.py训练脚本
  • 若要使用多GPU训练,使用torchrun --nproc_per_node=8 train_multi_GPU.py指令,nproc_per_node参数为使用GPU数量 (相比于之前的多GPU运行的命令,torchrun指令可以更好的管理进程)
  • 如果想指定使用哪些GPU设备可在指令前加上CUDA_VISIBLE_DEVICES=0,3(例如我只要使用设备中的第1块和第4块GPU设备)
  • CUDA_VISIBLE_DEVICES=0,3 torchrun --nproc_per_node=2 train_multi_GPU.py

注意事项

  1. 在使用训练脚本时,注意要将--data-path设置为自己存放数据集的根目录:
# 假设要使用COCO数据集,启用自定义数据集读取CocoDetection并将数据集解压到成/data/coco2017目录下
python train.py --data-path /data/coco2017

# 假设要使用Pascal VOC数据集,启用自定义数据集读取VOCInstances并数据集解压到成/data/VOCdevkit目录下
python train.py --data-path /data/VOCdevkit
  • 如果倍增batch_size,建议学习率也跟着倍增。假设将batch_size从4设置成8,那么学习率lr从0.004设置成0.008
  • 如果使用Batch Normalization模块时,batch_size不能小于4,否则效果会变差。如果显存不够,batch_size必须小于4时,建议在创建resnet50_fpn_backbone时, 将norm_layer设置成FrozenBatchNorm2d或将trainable_layers设置成0(即冻结整个backbone)
  • 训练过程中保存的det_results.txt(目标检测任务)以及seg_results.txt(实例分割任务)是每个epoch在验证集上的COCO指标,前12个值是COCO指标,后面两个值是训练平均损失以及学习率
  • 在使用预测脚本时,要将weights_path设置为你自己生成的权重路径。
  • 使用validation文件时,注意确保你的验证集或者测试集中必须包含每个类别的目标,并且使用时需要修改–num-classes、–data-path、–weights-path以及 --label-json-path(该参数是根据训练的数据集设置的)。其他代码尽量不要改动

复现结果

在COCO2017数据集上进行复现,训练过程中仅载入Resnet50的预训练权重,训练26个epochs。训练采用指令如下:

torchrun --nproc_per_node=8 train_multi_GPU.py --batch-size 8 --lr 0.08 --pretrain False --amp True

训练得到权重下载地址: https://pan.baidu.com/s/1qpXUIsvnj8RHY-V05J-mnA 密码: 63d5

在COCO2017验证集上的mAP(目标检测任务):

 Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.381
 Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.588
 Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.411
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.215
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.420
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.492
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.315
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.499
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.523
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.319
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.565
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.666

在COCO2017验证集上的mAP(实例分割任务):

 Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.340
 Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.552
 Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.361
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.151
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.369
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.500
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.290
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.449
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.468
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.266
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.509
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.619

2. 训练讲解

对应代码在train.py

训练参数

  • device: 默认指定为第一块GPU:cuda:0
  • data_path: 指向数据集的跟目录,默认使用coco数据集
  • num-classes: 由于默认使用的是coco数据集,所以设置num-classes=90 (不包含背景有90个类别)。
  • 为啥不是80个类别呢?其实在coco数据集中包含背景的话一共91个类别。不包含背景的90个类别,有10个列表标注为N/A,即不可用,将这10个类别排除后,你会发现它其实只有80个类别。不可用的类别是用来做全景分割使用的,在目标检测中是没有使用到的。类别信息在coco91_indices.json文件中

  • out_dir:设置保存权重的路径
  • resume: 在训练中中断时,可以将resume设置为最后一次保存的权重路径,那么它就会默认读取最新的训练权重,接着之前的训练结果往后训练。
  • lr-steps:训练过程中,迭代到哪些Epoch后我们会对学习率进行衰减。比如这里设置为default =[16,22],即训练到16和22epochs,会对学习率进行衰减。
  • lr-gamma: 衰减的倍率因子,默认为0.1,比如初始学习率lr为0.004,当迭代到第16个epoch时,那么学习率会衰减到0.004*0.1,那么到第22个epoch时会衰减到0.004*0.1*0.1
  • batch-size: 建议设置大于4,需要根据GPU显存进行设置。
  • aspect-ratio-group-factor:默认设置为3
  • pretrain:默认设置为True,即默认会载入官方在coco数据集上训练好的权重,如果你想从头训练的话,将pretrain设置为false即可,它就不会载入预训练权重。
  • amp: 是否去使用混合精度,默认设置为False,需要根据GPU设备而定,如果设备支持混合精度,就可以将它设置为True,采用混合精度它不仅能够提升训练速度,还能减少GPU显存占用。

main 函数

设置保存结果的文件

检测和实例分割的验证的结果

now = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
det_results_file = f"det_resultsnow.txt"
seg_results_file = f"seg_resultsnow.txt"

读取数据集

 # coco2017 -> annotations -> instances_train2017.json
train_dataset = CocoDetection(data_root, "train", data_transform["train"])
# VOCdevkit -> VOC2012 -> ImageSets -> Main -> train.txt
# train_dataset = VOCInstances(data_root, year="2012", txt_name="train.txt", transforms=data_transform["train"])

默认读取coco数据集,指定数据集的根目录,如果是读取训练集就传入train,如果读取验证集就传入val, 然后再传入对应的预处理方式。如果不想使用coco数据集,而是使用voc数据集,那么就用VOCInstances类。

aspect_ration_group_factor

  • 训练之前会按照图片的长宽比进行归类,这样做的目的是为了减少训练过程中GPU所占用的显存。当然如果GPU显存比较大,可以完全不考虑使用它,将aspect_ration_group_factor设置为-1就可以了。

  • 比如将batch-size设置为2,由于训练时需要将图片打包为一个batch,所以输入网络时,它真正的大小应该为黑色矩形框的大小(这样就可以装下这个batch的图片)。对应右边同样是两种图片,输入网络的大小也是外面的黑色矩形框。很明显到一个batch的所有图片比率比较相近的时候,黑色的矩形框要小一点。当一个batch的长宽比率相差比较大的时候,我们真正输入网络的图片其实是很大的。

  • 所以按照长宽比率进行排序之后,我们对它分组,然后再组内进行采样,这样我们输入网络的图片打包成batch它的大小会更小些,从而对GPU的显存占用也会更小些。

# 是否按图片相似高宽比采样图片组成batch
# 使用的话能够减小训练时所需GPU显存,默认使用
if args.aspect_ratio_group_factor >= 0:
    train_sampler = torch.utils.data.RandomSampler(train_dataset)
    # 统计所有图像高宽比例在bins区间中的位置索引
    group_ids = create_aspect_ratio_groups(train_dataset, k=args.aspect_ratio_group_factor)
    # 每个batch图片从同一高宽比例区间中取
    train_batch_sampler = GroupedBatchSampler(train_sampler, group_ids, args.batch_size)

Dataloader

  • 训练集的dataloader
if train_sampler:
        # 如果按照图片高宽比采样图片,dataloader中需要使用batch_sampler
        train_data_loader = torch.utils.data.DataLoader(train_dataset,
                                                        batch_sampler=train_batch_sampler,
                                                        pin_memory=True,
                                                        num_workers=nw,
                                                        collate_fn=train_dataset.collate_fn)
    else:
        train_data_loader = torch.utils.data.DataLoader(train_dataset,
                                                        batch_size=batch_size,
                                                        shuffle=True,
                                                        pin_memory=True,
                                                        num_workers=nw,
                                                        collate_fn=train_dataset.collate_fn)

train_data_loader根据我们是否启用了aspect_ratio_group_factor来进行相应设置的。

  • 验证集的dataloader
# load validation data set
 # coco2017 -> annotations -> instances_val2017.json
 val_dataset = CocoDetection(data_root, "val", data_transform["val"])
 # VOCdevkit -> VOC2012 -> ImageSets -> Main -> val.txt
 # val_dataset = VOCInstances(data_root, year="2012", txt_name="val.txt", transforms=data_transform["val"])
 val_data_loader = torch.utils.data.DataLoader(val_dataset,
                                               batch_size=1,
                                               shuffle=False,
                                               pin_memory=True,
                                               num_workers=nw,
                                               collate_fn=train_dataset.collate_fn)

这里注意下,验证的时候batch-size一定要设置为1,如果在验证的时候单张输入图片的话,那么我们输入网络的图片大小其实是和当前图片大小是一样的,此时我们对图片不需要太多处理。但是如果将多张图片同时输入到网络中,需要对每张图片进行padding调整多张图片到统一尺寸,这样就会补很多0,会对验证的结果产生影响。所以最准确的验证方法就是单张预测,需要将batch-size设置为1就可以了。

  • 实例化模型 create_model
def create_model(num_classes, load_pretrain_weights=True):
    # 如果GPU显存很小,batch_size不能设置很大,建议将norm_layer设置成FrozenBatchNorm2d(默认是nn.BatchNorm2d)
    # FrozenBatchNorm2d的功能与BatchNorm2d类似,但参数无法更新
    # trainable_layers包括['layer4', 'layer3', 'layer2', 'layer1', 'conv1'], 5代表全部训练
    # backbone = resnet50_fpn_backbone(norm_layer=FrozenBatchNorm2d,
    #                                  trainable_layers=3)
    # resnet50 imagenet weights url: https://download.pytorch.org/models/resnet50-0676ba61.pth
    backbone = resnet50_fpn_backbone(pretrain_path="resnet50.pth", trainable_layers=3)

    model = MaskRCNN(backbone, num_classes=num_classes)

    if load_pretrain_weights:
        # coco weights url: "https://download.pytorch.org/models/maskrcnn_resnet50_fpn_coco-bf2d0c1e.pth"
        weights_dict = torch.load("./maskrcnn_resnet50_fpn_coco.pth", map_location="cpu")
        for k in list(weights_dict.keys()):
            if ("box_predictor" in k) or ("mask_fcn_logits" in k):
                del weights_dict[k]

        print(model.load_state_dict(weights_dict, strict=False))

    return model
  • 传入num_classes,和load_pretrain_weights(是否载入预训练权重)
  • 实例化backbone:resnet50_fpn_backbone, 传入预训练模型pretrain_path,以及trainable_layers, 其中trainable_layers包括['layer4','layer3','layer2','layer1','conv1'], trainable_layers=3表示只训练layer4,layer3,layer2。如果想冻结所有权重,可将trainable_layers设置为0。
  • 如果由于GPU显存的限制,batch-size比较小的话,可以传入norm_layer=FrozenBatchNorm2d,此时在backbone中会将所有BN层替换为FrozenBatchNorm2d,这样的话即使你讲batch-size设置为较小的值,也不会影响你最终的效果。
backbone = resnet50_fpn_backbone(norm_layer=FrozenBatchNorm2d,
                                 trainable_layers=3)
  • 然后将backbonenum_classes传入到MaskRcnn就实例化了model
  • 如果load_pretain_weights为True,就会载入官方在coco数据集中预训练好的maskrcnn模型, 将有关类别的权重删除掉,然后载入剩余的模型权重。
 if load_pretrain_weights:
      # coco weights url: "https://download.pytorch.org/models/maskrcnn_resnet50_fpn_coco-bf2d0c1e.pth"
      weights_dict = torch.load("./maskrcnn_resnet50_fpn_coco.pth", map_location="cpu")
      for k in list(weights_dict.keys()):
          if ("box_predictor" in k) or ("mask_fcn_logits" in k):
              del weights_dict[k]

      print(model.load_state_dict(weights_dict, strict=False))
  • 训练参数及优化器
# define optimizer
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params, lr=args.lr,
                            momentum=args.momentum,
                            weight_decay=args.weight_decay)

scaler = torch.cuda.amp.GradScaler() if args.amp else None
  • 遍历模型的所有参数,获取requires_gradTrue的所有权重,训练时就会训练这些参数。
  • 默认采用优化器是SGD,加上momentum以及weight_decay
  • 如果将混合精度amp设置为true的话,就会实例化GradScaler,不采用的话就会等于None。

设置学习策略

 # learning rate scheduler
   lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer,
                                                       milestones=args.lr_steps,
                                                       gamma=args.lr_gamma)

指定在哪些Epoch的时候,将权重衰减,衰减因子为lr_gamma

  • 断点续练

根据args.resume获得最新一次的权重,载入包括模型权重、优化器optimizer、学习率lr_scheduler等等,然后接着最新的权重继续训练。

   if args.resume:
        # If map_location is missing, torch.load will first load the module to CPU
        # and then copy each parameter to where it was saved,
        # which would result in all processes on the same machine using the same set of devices.
        checkpoint = torch.load(args.resume, map_location='cpu')  # 读取之前保存的权重文件(包括优化器以及学习率策略)
        model.load_state_dict(checkpoint['model'])
        optimizer.load_state_dict(checkpoint['optimizer'])
        lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
        args.start_epoch = checkpoint['epoch'] + 1
        if args.amp and "scaler" in checkpoint:
            scaler.load_state_dict(checkpoint["scaler"])
  • 训练每个Epoch
 mean_loss, lr = utils.train_one_epoch(model, optimizer, train_data_loader,
                                              device, epoch, print_freq=50,
                                              warmup=True, scaler=scaler)
        train_loss.append(mean_loss.item())
        learning_rate.append(lr)

        # update the learning rate
        lr_scheduler.step()

        # evaluate on the test dataset
        det_info, seg_info = utils.evaluate(model, val_data_loader, device=device)

一个个Epoch训练,每迭代完一个Epoch就会调用lr_scheduler.step()方法,训练完一个Epoch就需要对模型进行验证,验证完之后会保存有关目标检测以及实例分割的信息。

    with open(det_results_file, "a") as f:
            # 写入的数据包括coco指标还有loss和learning rate
            result_info = [f"i:.4f" for i in det_info + [mean_loss.item()]] + [f"lr:.6f"]
            txt = "epoch: ".format(epoch, '  '.join(result_info))
            f.write(txt + "\\n")

        # write seg into txt
        with open(seg_results_file, "a") as f:
            # 写入的数据包括coco指标还有loss和learning rate
            result_info = [f"i:.4f" for i in seg_info + [mean_loss.item()]] + [f"lr:.6f"]
            txt = "epoch: ".format(epoch, '  '.join(result_info))
            f.write(txt + "\\n")
  • 保存训练的模型,这里需要注意的是为什么我们保存的模型那么大?,因为我们这里不光保存了模型的权重,还保存了优化器的权重学习率的参数,以及训练的epoch。如果只保存模型的权重,模型的大小就会小很多。但是如果只保存模型的权重,就没法去使用resume参数了,因为如果只载入模型权重,但不知道之前的优化器的信息,所以在保存的时候将这些信息都保存起来了
 save_files = 
            'model': model.state_dict(),
            'optimizer': optimizer.state_dict(),
            'lr_scheduler': lr_scheduler.state_dict(),
            'epoch': epoch
        if args.amp:
            save_files["scaler"] = scaler.state_dict()
        torch.save(save_files, "./save_weights/model_.pth".format(epoch))

训练完成后,会绘制

以上是关于Mask-RCNN : 代码使用的主要内容,如果未能解决你的问题,请参考以下文章

「深度学习一遍过」必修27:基于Mask-RCNN的人体姿态估计的设计与实现

『计算机视觉』RCNN学习_其二:Mask-RCNN

深度学习与计算机视觉任务应用综述

Mask-RCNN : 代码使用

mask-rcnn代码解读:display(self)函数的解析

深度学习在计算机视觉各项任务中的应用