mmdetection代码解析（1）train.py

待我称王封你为后i 2021-11-16 23:31 893阅读 0赞

#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()