Mask RCNN网络源码解读（Ⅴ） --- Mask R-CNN论文解读环境配置以及训练脚本解析
目录
1.源码地址
2.项目配置
2.1 环境配置
2.2 文件结构
2.3 预训练权重下载地址（下载后放入当前文件夹中）
2.4 数据集：本例程使用的有COCO2017数据集和Pascal VOC2012数据集
2.4.1 COCO2017数据集
2.4.2 Pascal VOC2012数据集
2.5 训练方法
2.6 注意事项
3.train.py 解析
1.源码地址
Mask R-CNN源码地址
https://pan.baidu.com/s/1ZDU0aM5DIpzCLxlnxvS8TA?pwd=5lwg
2.项目配置
2.1 环境配置
- 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
2.2 文件结构
├── 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）
2.3 预训练权重下载地址（下载后放入当前文件夹中）
- 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)
2.4 数据集：本例程使用的有COCO2017数据集和Pascal VOC2012数据集
2.4.1 COCO2017数据集
- COCO官网地址：COCO - Common Objects in Context
- 对数据集不了解的可以看下博文：MS COCO数据集介绍以及pycocotools简单使用_太阳花的小绿豆的博客-CSDN博客_coco数据集文章目录1. MS COCO数据集简介2. MS COCO数据集下载3. MS COCO标注文件格式3.1 使用Python的json库查看3.2 使用官方cocoAPI查看1. MS COCO数据集简介官网地址https://cocodataset.org/简介MS COCO是一个非常大型且常用的数据集，其中包括了目标检测，分割，图像描述等。其主要特性如下：Object segmentation: 目标分割Recognition in context: 图像情景识别Superphttps://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: 对应人体关键点检测的验证集标注文件夹
2.4.2 Pascal VOC2012数据集
- 数据集下载地址： The PASCAL Visual Object Classes Challenge 2012 (VOC2012)
- 对数据集不了解的可以看下博文：PASCAL VOC2012数据集介绍_太阳花的小绿豆的博客-CSDN博客_voc2012数据集之前有在Bilibili上简单介绍过这个数据集，但一直没有写博文，今天抽空总结下，如果不想看文章的，可以看下我在Bilibili上的讲解视频。Pascal VOC2012数据集详解视频: https://b23.tv/F1kSCKPascal VOC2012官网地址：http://host.robots.ox.ac.uk/pascal/VOC/voc2012/官方发表关于介绍数据集的文章 《The PASCALVisual Object Classes Challenge: A Retrospectihttps://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图（基于目标）
2.5 训练方法
- 确保提前准备好数据集
- 确保提前下载好对应预训练模型权重
- 确保设置好
--num-classes和--data-path- 若要使用单GPU训练直接使用train.py训练脚本
- 若要使用多GPU训练，使用
torchrun --nproc_per_node=8 train_multi_GPU.py指令,nproc_per_node参数为使用GPU数量- 如果想指定使用哪些GPU设备可在指令前加上
CUDA_VISIBLE_DEVICES=0,3(例如我只要使用设备中的第1块和第4块GPU设备)CUDA_VISIBLE_DEVICES=0,3 torchrun --nproc_per_node=2 train_multi_GPU.py
2.6 注意事项
- 在使用训练脚本时，注意要将
--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（该参数是根据训练的数据集设置的）。其他代码尽量不要改动
3.train.py 解析
if __name__ == "__main__": import argparse parser = argparse.ArgumentParser( description=__doc__) # 训练设备类型 parser.add_argument('--device', default='cuda:0', help='device') # 训练数据集的根目录 parser.add_argument('--data-path', default='/data/coco2017', help='dataset') # 检测目标类别数(不包含背景) parser.add_argument('--num-classes', default=90, type=int, help='num_classes') # 文件保存地址 parser.add_argument('--output-dir', default='./save_weights', help='path where to save') # 若需要接着上次训练，则指定上次训练保存权重文件地址 parser.add_argument('--resume', default='', type=str, help='resume from checkpoint') # 指定接着从哪个epoch数开始训练 parser.add_argument('--start_epoch', default=0, type=int, help='start epoch') # 训练的总epoch数 parser.add_argument('--epochs', default=26, type=int, metavar='N', help='number of total epochs to run') # 学习率 parser.add_argument('--lr', default=0.004, type=float, help='initial learning rate, 0.02 is the default value for training ' 'on 8 gpus and 2 images_per_gpu') # SGD的momentum参数 parser.add_argument('--momentum', default=0.9, type=float, metavar='M', help='momentum') # SGD的weight_decay参数 parser.add_argument('--wd', '--weight-decay', default=1e-4, type=float, metavar='W', help='weight decay (default: 1e-4)', dest='weight_decay') # 针对torch.optim.lr_scheduler.MultiStepLR的参数 parser.add_argument('--lr-steps', default=[16, 22], nargs='+', type=int, help='decrease lr every step-size epochs') # 针对torch.optim.lr_scheduler.MultiStepLR的参数 parser.add_argument('--lr-gamma', default=0.1, type=float, help='decrease lr by a factor of lr-gamma') # 训练的batch size(如果内存/GPU显存充裕，建议设置更大) parser.add_argument('--batch_size', default=2, type=int, metavar='N', help='batch size when training.') parser.add_argument('--aspect-ratio-group-factor', default=3, type=int) parser.add_argument("--pretrain", type=bool, default=True, help="load COCO pretrain weights.") # 是否使用混合精度训练(需要GPU支持混合精度) parser.add_argument("--amp", default=False, help="Use torch.cuda.amp for mixed precision training") args = parser.parse_args() print(args) # 检查保存权重文件夹是否存在，不存在则创建 if not os.path.exists(args.output_dir): os.makedirs(args.output_dir) main(args)训练设备我们默认是机器上的GPU，如果GPU环境配置好的话默认使用电脑的GPU。
训练集我们默认使用COCO，如果换成自己的数据集记得更改路径。
检测目标类别数(不包含背景)设置为90，VOC要改成21。
--output-dir：权重保存地址
--resume：如果训练中断了，比如我要训练26个epoch，当我训练到第十个epoch的时候突然断电了，我想接着来，则将resume设置为最后一次保存的权重地址。
----start_epoch：指定接着从哪个epoch数开始训练
--epochs：训练的总epoch数
--lr：学习率
--momentum： SGD的momentum参数
--wd：SGD的weight_decay参数
--lr-steps：在训练过程中迭代到哪些epoch时对学习率进行衰减（16，22）
--lr-gamma：衰减的倍率因子（0.1），比如这里初始学习率是0.00=04，当迭代到16个epoch时学习率就会衰减到0.0004，到第22个epoch时，会衰减到0.00004.
--batch_size：训练的batch size(如果内存/GPU显存充裕，建议设置更大)。
--pretrain：一般为true，载入官方训练好的训练权重。如果要从头训练设置成false即可。
--amp：是否使用混合精度训练(需要GPU支持混合精度)，最好设置成为true。
def main(args): device = torch.device(args.device if torch.cuda.is_available() else "cpu") print("Using {} device training.".format(device.type)) # 用来保存coco_info的文件 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" data_transform = { "train": transforms.Compose([transforms.ToTensor(), transforms.RandomHorizontalFlip(0.5)]), "val": transforms.Compose([transforms.ToTensor()]) } data_root = args.data_path # load train data set # 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"]) train_sampler = None # 是否按图片相似高宽比采样图片组成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) # 注意这里的collate_fn是自定义的，因为读取的数据包括image和targets，不能直接使用默认的方法合成batch batch_size = args.batch_size nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, 8]) # number of workers print('Using %g dataloader workers' % nw) 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) # 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) # create model num_classes equal background + classes model = create_model(num_classes=args.num_classes + 1, load_pretrain_weights=args.pretrain) model.to(device) train_loss = [] learning_rate = [] val_map = [] # 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 # learning rate scheduler lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=args.lr_steps, gamma=args.lr_gamma) # 如果传入resume参数，即上次训练的权重地址，则接着上次的参数训练 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"]) for epoch in range(args.start_epoch, args.epochs): # train for one epoch, printing every 50 iterations 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) # write detection into txt 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") val_map.append(det_info[1]) # pascal mAP # save weights 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)) # 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)首先判断GPU环境是否配置正确，GPU环境配置正确则使用GPU，否则使用CPU。
创建两个文件det_results_file 、seg_results_file ：分别用来保存目标检测的验证结果以及实例分割的验证结果。
data_transform模块实现对数据的预处理，官方仅仅将图片转化成Tensor并实现水平方向翻转，这个之前再Faster R-CNN部分讲过，不再赘述。
train_dataset实现数据读取，我们默认读取COCO数据集，CocoDetection类是在my_dataset_coco.py文件下的（后续博客会说）。
aspect_ratio_group_factor参数如果启用的话，我们会在训练之前对数据按照图片长宽比例进行归类，这么做的目的是为了减少GPU占用显存。
这里简单说明一下为什么将图片按高宽比归类能减少显存占用：
我们这里假设batch_size设置为2，蓝色对应第一张图片黄色对应第二张图片，由于我们训练时会将图片打包成一个batch，因此真正输入网络的大小是黑色虚线部分的大小。很明显，当两张图片比例比较相近的时候黑色矩形框要较小一点（如图左）否则黑色矩形框即一个batch中的图片的大小会更大一点，因此我们根据长宽比例进行排序之后，我们输入到网络中的batch大小会更小一些，这样对GPU的显存占用也会更少一些。
train_data_loader也是根据我们是否使用aspect_ratio_group_factor来设置的。
val_data_loader实例化了验证集的数据，val_data_loader实例化了验证集的数据读取器。在验证的时候，记得将batch_size设置为1。验证时如果是单张输入图片的话，那么我们输入网络的图片大小其实就是当前图片大小（不会进行batch归一化）。
用create_model方法实例化模型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这里的输入参数是训练模型的类别以及是否载入官方的权重。
首先我们实例化我们的backbone，这里我们使用resnet50_fpn_backbone作为backbone。这里我们就会载入resnet50的权重并指定训练哪些层。（'layer4', 'layer3', 'layer2'），如果不想训练这个则设置为0，即冻结backbone所有权重。如果batch_size设置的比较小的话，可以启用norm_layer=FrozenBatchNorm2d这个选项，就会将网络的BN层替换成FrozenBatchNorm层，这样即使将batchsize设置成较小的值也不会影响最终的预测效果。
用MaskRCNN（后面博客会说）网络初始化model。
如果我们的构造参数load_pretrain_weights=True时，我们就会载入官方提供的预训练权重，载入的过程中会将有关于类别的权重删除掉，删除掉后将剩余的权重载入到网络中。
回到我们的main函数中：
遍历模型的所有参数，将requires_grad为true的权重记录下来，在训练的时候我们就会训练这些权重，我们记录在变量params中。
这里默认的优化器是SGD + Momentum。
如果我们使用混合精度就会实例化一个scaler。
实例化一个学习率下降的变量lr_scheduler，指定在哪些epoch中进行学习率的衰减。
如果传入resume参数，即上次训练的权重地址，则接着上次的参数训练。如果不用接着上次的权重进行训练（从头开始训练）：
我们从args.start_epoch循环次数开始训练到args.epochs次数为止，我们每循环一个epoch就会调用lr_scheduler方法更新学习率，调用evaluate对模型进行验证，保存有关于目标检测和实例分割的信息，最后保存模型的权重。
最后会画出训练曲线。
接着利用predict.py脚本读取训练好的权重看下效果。这里需要修改下类别数、读取的权重路径以及索引对应类别的json文件。
效果还行！！
