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:0data_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.1batch-size: 建议设置大于4，需要根据GPU显存进行设置。aspect-ratio-group-factor:默认设置为3pretrain:默认设置为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)
其实训练的时候，不一定非要每个Epoch都去验证一次，为了节约时间，大家可以指定在某些epoch进行验证，比如指定 epoch % 10 ==0每隔10个epoch进行验证, 不过这样就没法获得保存每个epoch的验证结果，这样就没法根据每个epoch它的验证结果去绘制对应的曲线了。
参考: https://www.bilibili.com/video/BV1hY411E7wD