目录

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 注意事项

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

1. 如果倍增batch_size，建议学习率也跟着倍增。假设将batch_size从4设置成8，那么学习率lr从0.004设置成0.008
2. 如果使用Batch Normalization模块时，batch_size不能小于4，否则效果会变差。如果显存不够，batch_size必须小于4时，建议在创建resnet50_fpn_backbone时， 将norm_layer设置成FrozenBatchNorm2d或将trainable_layers设置成0(即冻结整个backbone)
3. 训练过程中保存的det_results.txt(目标检测任务)以及seg_results.txt(实例分割任务)是每个epoch在验证集上的COCO指标，前12个值是COCO指标，后面两个值是训练平均损失以及学习率
4. 在使用预测脚本时，要将weights_path设置为你自己生成的权重路径。
5. 使用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文件。

         效果还行！！