Mask-RCNN(2) : 代码使用

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）

预训练权重下载

Resnet50预训练权重 https://download.pytorch.org/models/resnet50-0676ba61.pth (注意，下载预训练权重后要重命名，比如在train.py中读取的是resnet50.pth文件，不是resnet50-0676ba61.pth)
Mask R-CNN(Resnet50+FPN)预训练权重 https://download.pytorch.org/models/maskrcnn_resnet50_fpn_coco-bf2d0c1e.pth (注意，载预训练权重后要重命名，比如在train.py中读取的是maskrcnn_resnet50_fpn_coco.pth文件，不是maskrcnn_resnet50_fpn_coco-bf2d0c1e.pth)

数据集

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

COCO2017数据集

COCO官网地址：https://cocodataset.org/
对数据集不了解的可以参考博文：https://blog.csdn.net/qq_37541097/article/details/113247318
这里以下载coco2017数据集为例，主要下载三个文件：
- 2017 Train images [118K/18GB]：训练过程中使用到的所有图像文件
- 2017 Val images [5K/1GB]：验证过程中使用到的所有图像文件
- 2017 Train/Val annotations [241MB]：对应训练集和验证集的标注json文件
都解压到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数据集

数据集下载地址： http://host.robots.ox.ac.uk/pascal/VOC/voc2012/index.html#devkit
对数据集不了解的可参考博文：https://blog.csdn.net/qq_37541097/article/details/115787033
解压后得到的文件夹结构如下：

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

注意事项

在使用训练脚本时，注意要将--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_results{now}.txt"
seg_results_file = f"seg_results{now}.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)

然后将backbone 和num_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_grad为True的所有权重，训练时就会训练这些参数。
默认采用优化器是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))

训练完成后，会绘制训练损失loss以及lr的曲线以及map曲线

  # plot loss and lr curve
    if len(train_loss) != 0 and len(learning_rate) != 0:
        from plot_curve import plot_loss_and_lr
        plot_loss_and_lr(train_loss, learning_rate)

    # plot mAP curve
    if len(val_map) != 0:
        from plot_curve import plot_map
        plot_map(val_map)