MaskrcnnBenchmark 源码解析-训练/推演工具脚本(tools)

发表于   |   更新于 | 分类于 | 评论数： | 阅读次数：
本文字数： 15k | 阅读时长 ≈ 13 分钟

源码文件

train_net.py 文件概览

为了更好的解读 MaskrcnnBenchmark 的源代码, 我们首先来看看执行模型训练代码的脚本文件都使用了哪些类和函数, 该脚本可以训练 MaskrcnnBenchmark 中的所有模型, 因此, 我们可以从该文件出发, 顺藤摸瓜的探索, 以期最终能够对整个 MaskrcnnBenchmark 框架有一个全面系统的了解, 那么接下来就先来改一下该文件的大致结构, 如下所示:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

# 导入各种包
from maskrcnn_benchmark.config import cfg
#...

def train(cfg, local_rank, distributed):
    # ... 训练脚本核心代码

def test(cfg, model, distributed):
    # ... 推演代码

def main():
    # ... 主函数, 会在内部调用上述函数

if __name__ == "__main__":
    main()

train_net 导入的各种包和函数

首先来看看该文件导入了那些包和函数, 同时我们会针对性的简单介绍一下它们的主要作用和功能.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
# ./tools/train_net.py

# 下面的这个必须在导入所有包之前导入, 不能放到其他位置. TODO 原因
from maskrcnn_benchmark.utils.env import setup_environment

# 常规包
import argparse
import os

import torch
from maskrcnn_benchmark.config import cfg # 导入默认配置信息
from maskrcnn_benchmark.data import make_data_loader # 数据集载入
from maskrcnn_benchmark.solver import make_lr_scheduler # 学习率更新策略
from maskrcnn_benchmark.solver import make_optimizer # 设置优化器, 封装了PyTorch的SGD类
from maskrcnn_benchmark.engine.inference import inference # 推演代码
from maskrcnn_benchmark.engine.trainer import do_train # 模型训练的核心逻辑. 会重点解析

# 调用了 ./maskrcnn_benchmark/modeling/detector/ 中的 build_detection_model() 函数
# 该函数和 Detectron 中的类似, 都是用来创建目标检测模型的, 这也是创建模型的入口函数, 十分重要
from maskrcnn_benchmark.modeling.detector import build_detection_model

from maskrcnn_benchmark.checkpoint import DetectronCheckpointer

# 封装了 PyTorch 的 torch.utils.collect_env.get_preety_env_info 函数, 同时附加了 PIL.__version__ 版本新
from maskrcnn_benchmark.utils.collect_env import collect_env_info

# 分布式训练相关设置, 由于我的gpu个数为1, 因此 get_rank() 会返回 0
from maskrcnn_benchmark.utils.comm import synchronize, get_rank

from maskrcnn_benchmark.utils.imports import import_file

# 封装了 logging 模块, 用于向屏幕输出一些日志信息
from maskrcnn_benchmark.utils.logger import set_up_logger

# 封装了 os.mkdirs 函数, 当文件夹已存在时会自动略过, 不会提示错误
from maskrcnn_benchmark.utils.miscellaneous import mkdir

相比于 Detectron 来说, MaskrcnnBenchmark 的默认配置文件显得相当 “清爽”, 定义的配置项也很精简, 下面就是 cfg 的部分配置清单, 输入 print(cfg) 即可看到全部配置项.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
DATALOADER:
  ASPECT_RATIO_GROUPING: True
  NUM_WORKERS: 4
  SIZE_DIVISIBILITY: 0
DATASETS:
  TEST: ()
  TRAIN: ()
INPUT:
  MAX_SIZE_TEST: 1333
  MAX_SIZE_TRAIN: 1333
  MIN_SIZE_TEST: 800
  MIN_SIZE_TRAIN: 800
  PIXEL_MEAN: [102.9801, 115.9465, 122.7717]
  PIXEL_STD: [1.0, 1.0, 1.0]
  TO_BGR255: True
MODEL:
  BACKBONE:
    CONV_BODY: R-50-C4
    FREEZE_CONV_BODY_AT: 2
    OUT_CHANNELS: 1024
  DEVICE: cuda
# ...

train_net.main() 主函数

下面我们根据脚本的执行顺序, 先来看看主函数的代码:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
# ./tools/train_net.py

def main():
    parser = argparse.ArgumentParser(description="PyTorch Object Detection Training")
    parser.add_argument(
        # 这里虽然是 config-file, 但要用 args.config_file来访问, 不能用 args.config-file 访问.
        "--config-file",
        default="",
        metavar="FILE",
        help="path to config file",
        type=str,
    )
    parser.add_argument("--local_rank", type=int, default=0)
    parser.add_argument(
        "--skip-test",
        dest="skip_test",
        help="Do not test the final model",
        action="store_true",
    )
    parser.add_argument(
        "opts",
        help="Modify config options using the command-line",
        default=None,
        nargs=argparse.REMAINDER, # 这一行不能少
    )

    args = parser.parse_args()

    # 获取 gpu 的数量, 我电脑只有一个GPU, 故 num_gpus = 1
    num_gpus = int(os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1

    # 是否进行多GPU的分布式训练
    args.distributed = num_gpus > 1

    # 设置分布式训练时的一些初始化信息
    if args.distributed:
        torch.cuda.set_device(args.local_rank)
        torch.distributed.init_process_group(
            backend="nccl", init_method="env://"
        )

    # 将 config_file 指定的配置项覆盖到默认配置项当中
    cfg.merge_from_file(args.config_file)
    cfg.merge_from_list(args.opts)
    cfg.freeze() # 冻结所有的配置项, 防止修改

    output_dir = cfg.OUTPUT_DIR
    if output_dir: # 这里封装了 os.mkdir, 使得当output_dir存在时, 会直接略过, 不会返回错误
        mkdir(output_dir)

    # 下面的多行代码都是一些向屏幕上输出相关信息的日志代码
    # 同时也会保存在 output_dir 文件夹下的 log.txt 文件内
    # 输出的信息来源于自身系统以及配置文件中的信息
    logger = setup_logger("maskrcnn_benchmark", output_dir, get_rank())
    logger.info("Using {} GPUs".format(num_gpus))
    logger.info(args)

    logger.info("Collecting env info (might take some time)")
    logger.info("\n" + collect_env_info())

    logger.info("Loaded configuration file {}".format(args.config_file))

    # 打开指定的配置文件, 并读取其中的相关信息, 将值存储在 config_str 中, 然后输出到屏幕上
    with open(args.config_file, "r") as cf:
        config_str = "\n" + cf.read()
        logger.info(config_str)
    logger.info("Running with config:\n{}".format(cfg))

    # 调用 train 函数, 该函数会执行模型训练的代码逻辑, 详解可看后文
    model = train(cfg, args.local_rank, args.distributed)

    if not args.skip_test:

        # 同理, 该函数会执行模型推演的代码逻辑, 详情可看后文
        test(cfg, model, args.distributed)

train_net.train() 训练脚本

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
def train(cfg, local_rank, distributed):

    # 该语句调用了 ./maskrcnn_benchmark/modeling/detector/ 中的 build_detection_model() 函数
    # 该函数和 Detectron 中的类似, 都是用来创建目标检测模型的, 这也是创建模型的入口函数, 十分重要
    # 这里我们只需要知道该函数会根据我们的配置文件返回一个网络模型就可以了
    model = build_detection_model(cfg)

    # 默认为 "cuda"
    device = torch.device(cfg.MODEL.DEVICE)
    model.to(device) # 将模型移到指定设备上

    # 封装了 torch.optiom.SGD() 函数, 根据tensor的requires_grad属性构成需要更新的参数列表
    optimizer = make_optimizer(cfg, model)

    # 根据配置信息设置 optimizer 的学习率更新策略
    scheduler = make_lr_scheduler(cfg, optimizer)

    # 分布式训练情况下, 并行处理数据
    if distributed:
        model = torch.nn.parallel.DistributedDataParallel(
            model, device_ids=[local_rank], output_device=local_rank,
            # this should be removed if we update BatchNorm stats
            broadcast_buffers=False,
        )

    # 创建一个参数字典, 并将迭代次数置为0
    arguments = {}
    arguments["iteration"] = 0

    # 获取输出的文件夹路径, 默认为 '.', 这里建议将 cfg.OUTPUT_DIR 提前设置成自己的路径.
    output_dir = cfg.OUTPUT_DIR

    # 因为我只有一个gpu, 所以这里 save_to_disk=True
    save_to_disk = get_rank() == 0

    # DetectronCheckpointer 对象, 后面会用在 do_train() 函数的参数
    checkpointer = DetectronCheckpointer(
        cfg, model, optimizer, scheduler, output_dir, save_to_disk
    )
    extra_checkpoint_data = checkpointer.load(cfg.MODEL.WEIGHT) # 加载指定权重文件
    arguments.update(extra_checkpoint_data) # 字典的update方法, 对字典的键值进行更新

    # data_loader 的类型为列表, 内部元素类型为 torch.utils.data.DataLoader
    # 注意, 当is_train=True时, 要确保 cfg.DATASETS.TRAIN 的值为一个列表
    # 并且必须指向一个或多个存在的anns文件, 默认情况下该值为空, 所以必须在配置文件中指定
    data_loader = make_data_loader(
        cfg,
        is_train=True,
        is_distributed=distributed,
        start_iter=arguments["iteration"],
    )

    checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD # 默认值为2500

    do_train(
        model,
        data_loader,
        optimizer,
        scheduler,
        checkpointer,
        device,
        checkpoint_period,
        arguments,
    )

    return model

在执行模型训练逻辑时, 该函数调用了 ./maskrcnn_benchmark/engine/trainer.py 文件中的 do_train()函数, 该函数是执行训练逻辑的核心代码, 具体的解析请看do_train

我们注意到, 在代码的第一句开头使用了 ./maskrcnn_benchmark/modeling/detector/ 文件夹下面的用来创建目标检测模型的函数, 这也是创建模型的入口函数, 十分重要. 关于该函数的详细解析可以查看讲解模型创建的博文, 这里只需要知道该函数会根据我们的配置文件返回一个模型就可以了, 如果想查看该模型的具体网络结构及其参数, 可以利用 print(model) 查看, 如下所示为当前当前配置信息下的部分结构信息:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
GeneralizedRCNN(
  (backbone): Sequential(
    (body): ResNet(
      (stem): StemWithFixedBatchNorm(
        (conv1): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
        (bn1): FrozenBatchNorm2d()
      )
      (layer1): Sequential(
        (0): BottleneckWithFixedBatchNorm(
          (downsample): Sequential(
            (0): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
            (1): FrozenBatchNorm2d()
          )
          (conv1): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)
          (bn1): FrozenBatchNorm2d()
          (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
          (bn2): FrozenBatchNorm2d()
          (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
          (bn3): FrozenBatchNorm2d()
        )
#...

train_net.test() 推演脚本

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
def test(cfg, model, distributed):
    if distributed:
        model = model.module
    torch.cuda.empty_cache()  # TODO check if it helps. releases unoccupied memory
    iou_types = ("bbox",)
    if cfg.MODEL.MASK_ON:# 如果mask为Ture, 则添加分割信息
        iou_types = iou_types + ("segm",)
    output_folders = [None] * len(cfg.DATASETS.TEST) # 根据标签文件数确定输出文件夹数
    dataset_names = cfg.DATASETS.TEST
    if cfg.OUTPUT_DIR:
        for idx, dataset_name in enumerate(dataset_names):# 遍历标签文件
            output_folder = os.path.join(cfg.OUTPUT_DIR, "inference", dataset_name)
            mkdir(output_folder) # 创建输出文件夹
            output_folders[idx] = output_folder # 将文件夹的路径名放入列表

    # 根据配置文件信息创建数据集
    data_loaders_val = make_data_loader(cfg, is_train=False, is_distributed=distributed)

    # 遍历每个标签文件, 执行 inference 过程.
    for output_folder, dataset_name, data_loader_val in zip(output_folders, dataset_names, data_loaders_val):
        inference(
            model,
            data_loader_val,
            dataset_name=dataset_name,
            iou_types=iou_types,
            box_only=cfg.MODEL.RPN_ONLY,
            device=cfg.MODEL.DEVICE,
            expected_results=cfg.TEST.EXPECTED_RESULTS,
            expected_results_sigma_tol=cfg.TEST.EXPECTED_RESULTS_SIGMA_TOL,
            output_folder=output_folder,
        )
        synchronize() # 多GPU推演时的同步函数

在执行模型的推演逻辑时, 该函数调用了 ./maskrcnn_benchmark/engine/inference.py 文件中的 inference()函数, 该函数是执行推演逻辑的核心代码, 具体的解析请看inference

test_net.py 文件概览

1
2
3
4
5
6
7
8
9
10
11
12
# ./tools/test_net.py

# 导入各种包及函数
from maskrcnn_benchmark.config import cfg
# ...

# 主程序
def main():
    # ...

if __name__ == "__main__":
    main()

test_net 导入的各种包及函数

我们首先看看该文件导入的包及函数

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
# ./tools/test_net.py

from maskrcnn_benchmark.utils.env import setup_environment

# 常规包
import argparse
import os

import torch

from maskrcnn_benchmark.config import cfg # 导入全局配置

from maskrcnn_benchmark.data import make_data_loader # 数据载入器

from maskrcnn_benchmark.engine.inference import inference # 模型推演函数

from maskrcnn_benchmark.modeling.detector import build_detection_model # 创建模型

from maskrcnn_benchmark.utils.checkpoint import DetectronCheckpointer

from maskrcnn_benchmark.utils.collect_env import collect_env_info

from maskrcnn_benchmark.utils.comm import synchronize, get_rank

from maskrcnn_benchmark.utils.logger import setup_logger

from maskrcnn_benchmark.utils.miscellaneous import mkdir

可以看出, 在 ./tools/test_net.py 文件中导入的包和函数和 ./tools/train_net.py 差不多, 我们已经在之前简单介绍了这些包的功能和用途, 并给出了详细解析的链接, 这里我们就不再重复介绍, 有疑惑的可以翻到上面去看关于 train_net.py 导入的包及函数的解析.

test_net.main() 主函数

由于在进行模型推演时, 我们只需要准备好预训练文件, 数据集, 以及模型结构就可以完成整个推演过程, 因此在 test_net.py 脚本中只用了一个主函数来完成这些功能, 下面我们就来看看这个主函数的具体实现吧.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
# ./tools/test_net.py

def main():
    parser = argparse.ArgumentParser(description="PyTorch Object Detection Inference")

    # 权重文件路径
    parser.add_argument(
        "--config-file",
        default="/private/home/fmassa/github/detectron.pytorch_v2/configs/e2e_faster_rcnn_R_50_C4_1x_caffe2.yaml",
        metavar="FILE",
        help="path to config file",
    )

    # local_rank
    parser.add_argument("--local_rank", type=int, default=0)

    # 其他的配置选项
    parser.add_argument(
        "opts",
        help="Modify config options using the command-line",
        default=None,
        nargs=argparse.REMAINDER,
    )

    args = parser.parse_args()

    # 获取 gpu 数目
    num_gpus = int(os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1

    # 根据gpu数目设置distributed布尔变量
    distributed = num_gpus > 1

    if distributed:
        torch.cuda.set_device(args.local_rank)
        torch.distributed.init_process_group(
            backend="nccl", init_method="env://"
        )

        # 将指定的配置文件的设置覆盖到全局设置中
        cfg.merge_from_file(args.config_file)
        cfg.merge_from_list(args.opts)
        cfg.freeze() # 冻结配置信息, 防止更改

        save_dir = ""
        logger = setup_logger("maskrcnn_benchmark", save_dir, get_rank())
        logger.info("Using {} GPUs".format(num_gpus))
        logger.info(cfg)

        logger.info("Collecting env info (might take some time)")
        logger.info("\n" + collect_env_info())

        # 根据配置信息创建模型
        model = build_detection_model(cfg)

        # 将模型移动到指定设备上
        model.to(cfg.MODEL.DEVICE)

        # 获取输出文件夹父路径
        output_dir = cfg.OUTPUT_DIR

        # 加载权重
        checkpointer = DetectronCheckpointer(cfg, model, save_dir=output_dir)
        _ = checkpointer.load(cfg.MODEL.WEIGHT)

        # 设置 iou 类型
        iou_types = ("bbox",)
        if cfg.MODEL.MASK_ON:
            iou_types = iou_types + ("segm",)

        # 根据数据集的数量定义输出文件夹
        output_folders = [None] * len(cfg.DATASETS.TEST)

        dataset_names = cfg.DATASETS.TEST

        # 创建输出文件夹
        if cfg.OUTPUT_DIR:
            for idx, dataset_name in enumerate(dataset_names):
                output_folder = os.path.join(cfg.OUTPUT_DIR, "inference", dataset_name)
                mkdir(output_folder)
                output_folders[idx] = output_folder

        # 加载测试数据集
        data_loaders_val = make_data_loader(cfg, is_train=False, is_distributed=distributed)

        # 对数据集中的数据按批次调用inference函数
        for output_folder, dataset_name, data_loader_val in zip(output_folders, dataset_names, data_loaders_val):
            inference(
                model,
                data_loader_val,
                dataset_name=dataset_name,
                iou_types=iou_types,
                box_only=cfg.MODEL.RPN_ONLY,
                device=cfg.MODEL.DEVICE,
                expected_results=cfg.TEST.EXPECTED_RESULTS,
                expected_results_sigma_tol=cfg.TEST.EXPECTED_RESULTS_SIGMA_TOL,
                output_folder=output_folder,
                )
            synchronize()