mmdetection代码解析（1）train.py

#coding=utf-8

from __future__ import division

import argparse
import os
from mmcv import Config

from mmdet import __version__
from mmdet.datasets import get_dataset
from mmdet.apis import (train_detector, init_dist, get_root_logger,
                        set_random_seed)
from mmdet.models import build_detector
import torch


def parse_args():
    parser = argparse.ArgumentParser(description='Train a detector')

    # 模型配置文件
    parser.add_argument('config', help='train config file path')

    # log文件和模型文件存储路径
    parser.add_argument('--work_dir', help='the dir to save logs and models')

    # 恢复训练模型的路径
    parser.add_argument('--resume_from', help='the checkpoint file to resume from')

    # 是否在训练过程中评估checkpoint
    parser.add_argument(
        '--validate',
        action='store_true',
        help='whether to evaluate the checkpoint during training')

    # 使用的GPU数量，默认是1,，仅适用于非分布式训练
    parser.add_argument(
        '--gpus',
        type=int,
        default=1,
        help='number of gpus to use '
        '(only applicable to non-distributed training)')

    # 为当前GPU设置随机种子，
    # 在训练开始时，参数的初始化是随机的，
    # 为了让每次的结果一致，我们需要设置随机种子。
    parser.add_argument('--seed', type=int, default=None, help='random seed')

    # 决定是否使用分布式训练，
    # 分布式训练支持多种 launch 方法，
    # 如果launcher为空，怎为非分布式训练
    parser.add_argument(
        '--launcher',
        choices=['none', 'pytorch', 'slurm', 'mpi'],
        default='none',
        help='job launcher')

    # local_rank代表当前程序进程使用的GPU标号
    parser.add_argument('--local_rank', type=int, default=0)

    # 是否使用mixup进行数据增强
    parser.add_argument('--mixup', action='store_true', help='use mixup training.')

    args = parser.parse_args()

    # 如果环境变量中没有指定当前进程使用的GPU标号，则使用参数里指定的
    if 'LOCAL_RANK' not in os.environ:
        os.environ['LOCAL_RANK'] = str(args.local_rank)

    return args


def main():
    args = parse_args()

    # 设置模型配置文件
    cfg = Config.fromfile(args.config)

    # set cudnn_benchmark
    # 设置这个 flag 可以让内置的cuDNN的auto-tuner
    # 自动寻找最适合当前配置的高效算法，
    # 来达到优化运行效率的问题
    if cfg.get('cudnn_benchmark', False):
        torch.backends.cudnn.benchmark = True

    # 根据命令行参数更新模型配置文件中的如下两个配置选项
    if args.work_dir is not None:
        cfg.work_dir = args.work_dir
    if args.resume_from is not None:
        cfg.resume_from = args.resume_from

    # 训练使用的GPU数量
    cfg.gpus = args.gpus

    # 初始化分布式环境，日志初始化的时候需要依赖分布式信息
    if args.launcher == 'none':
        distributed = False
    else:
        distributed = True
        init_dist(args.launcher, **cfg.dist_params)

    # 初始化日志
    logger = get_root_logger(cfg.log_level)
    logger.info('Distributed training: {}'.format(distributed))

    # 为当前GPU设置随机种子，
    # 在训练开始时，参数的初始化是随机的，
    # 为了让每次的结果一致，我们需要设置随机种子
    if args.seed is not None:
        logger.info('Set random seed to {}'.format(args.seed))
        set_random_seed(args.seed)

    # 创建检测器：
    # 模型参数配置
    # 训练参数配置
    # 测试参数配置
    # 得到的model是一个模型类的对象
    # 比如mmdet/detectors/目录下的类的对象
    model = build_detector(cfg.model,
                           train_cfg=cfg.train_cfg,
                           test_cfg=cfg.test_cfg)
    # print(type(model))
    # 输出：<class 'mmdet.models.detectors.mask_rcnn.MaskRCNN'>

    # 根据配置文件加载训练数据
    # 得到一个数据集对象，比如：
    # mmdet.datasets.coco.CocoDataset的对象
    train_dataset = get_dataset(cfg.data.train)


    if cfg.checkpoint_config is not None:
        # 保存mmdet版本信息，模型配置文件内容，目标分类信息到
        # checkpoints作为元数据???
        # print(type(cfg.checkpoint_config))
        cfg.checkpoint_config.meta = dict(
            mmdet_version=__version__,
            config=cfg.text,
            CLASSES=train_dataset.CLASSES)

    # 为便于可视化添加类别属性
    model.CLASSES = train_dataset.CLASSES

    # 训练检测器
    train_detector(
        model, # 模型
        train_dataset, # 训练集数据
        cfg, # 模型配置文件
        distributed=distributed, # 分布式配置开关
        validate=args.validate, # 是否在训练过程中评估checkpoint
        logger=logger # 日志
        )


if __name__ == '__main__':
    main()