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)
预训练权重下载
- 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_laye
r设置成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
:默认设置为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_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)
- 然后将
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))
训练完成后,会绘制 以上是关于Mask-RCNN : 代码使用的主要内容,如果未能解决你的问题,请参考以下文章 「深度学习一遍过」必修27:基于Mask-RCNN的人体姿态估计的设计与实现