动作识别0-06：mmaction2(SlowFast)-源码无死角解析（2）-训练架构总览-1

以下链接是个人关于mmaction2(SlowFast-动作识别) 所有见解，如有错误欢迎大家指出，我会第一时间纠正。有兴趣的朋友可以加微信：17575010159 相互讨论技术。若是帮助到了你什么，一定要记得点赞！因为这是对我最大的鼓励。$文末附带$$公众号-$$海量资源。$

动作识别0-00：mmaction2(SlowFast)-目录-史上最新无死角讲解

$极度推荐的商业级项目：$这是本人落地的行为分析项目，主要包含（1.行人检测，2.行人追踪，3.行为识别三大模块）：行为分析(商用级别)00-目录-史上最新无死角讲解

前言

通过前面的介绍，我相信大家已经把模型训练起来了，我知道你是冰雪聪明的。那么下面我们就来分析其源码吧，不放过任何的细节，这样我们在做项目的时候心里才有底。废话不多说了，我们直接开始吧

代码注释

我们在训练的时候执行了这样一个指令：

python tools/train.py configs/recognition/slowfast/my_slowfast_r50_4x16x1_256e_ucf101_rgb.py   --work-dir work_dirs/my_slowfast_r50_4x16x1_256e_ucf101_rgb    --validate --seed 0 --deterministic

可以看到其运行 python 文件为 tools/train.py，本人注释如下：

import argparse
import copy
import os
import os.path as osp
import time

import mmcv
import torch
from mmcv import Config
from mmcv.runner import init_dist, set_random_seed

from mmaction import __version__
from mmaction.apis import train_model
from mmaction.datasets import build_dataset
from mmaction.models import build_model
from mmaction.utils import collect_env, get_root_logger


def parse_args():
    """
    对参数进行解析
    :return:
    """
    parser = argparse.ArgumentParser(description='Train a recognizer')
    # 配置文件路径
    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')
    # 在训练的过程中是否对数据进行验证
    parser.add_argument(
        '--validate',
        action='store_true',
        help='whether to evaluate the checkpoint during training')

    # 创建一个GPU参数组
    group_gpus = parser.add_mutually_exclusive_group()
    # 对GPU使用的数目进行设定(该参数只适用于非分布式的训练)
    group_gpus.add_argument(
        '--gpus',
        type=int,
        help='number of gpus to use '
        '(only applicable to non-distributed training)')

    # 对GPU使用的数目进行设定(该参数只适用于非分布式的训练)
    group_gpus.add_argument(
        '--gpu-ids',
        type=int,
        nargs='+',
        help='ids of gpus to use '
        '(only applicable to non-distributed training)')
    # 设定一个训练的随机种子，在做实验对比测时候，能保证训练测试数据一致
    parser.add_argument('--seed', type=int, default=None, help='random seed')
    # 是否为CUDNN后端设置确定性选项(暂时没有看明白是什么意思)
    parser.add_argument(
        '--deterministic',
        action='store_true',
        help='whether to set deterministic options for CUDNN backend.')
    # 进行训练的作者
    parser.add_argument(
        '--launcher',
        choices=['none', 'pytorch', 'slurm', 'mpi'],
        default='none',
        help='job launcher')

    parser.add_argument('--local_rank', type=int, default=0)
    args = parser.parse_args()

    # 分布式训练的rank号
    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

    # 命令行的参数会覆盖config文件中的参数
    # work_dir is determined in this priority:
    # CLI > config file > default (base filename)

    # 如果args.work_dir不为空，则命令行的work_dir替换掉cfg.work_dir
    if args.work_dir is not None:
        # update configs according to CLI args if args.work_dir is not None
        cfg.work_dir = args.work_dir
    # 如果args.work_dir为空,则使用args.config文件中的工作目录
    elif cfg.get('work_dir', None) is None:
        # use config filename as default work_dir if cfg.work_dir is None
        cfg.work_dir = osp.join('./work_dirs',
                                osp.splitext(osp.basename(args.config))[0])
    # 进行一些相关参数覆盖
    if args.resume_from is not None:
        cfg.resume_from = args.resume_from
    if args.gpu_ids is not None:
        cfg.gpu_ids = args.gpu_ids
    else:
        cfg.gpu_ids = range(1) if args.gpus is None else range(args.gpus)

    # 首先init分布式env，因为logger依赖于dist信息。
    # init distributed env first, since logger depends on the dist info.
    if args.launcher == 'none':
        distributed = False
    else:
        distributed = True
        init_dist(args.launcher, **cfg.dist_params)


    # create work_dir，创建work_dir目录
    mmcv.mkdir_or_exist(osp.abspath(cfg.work_dir))
    # dump config，转存储config文件到work_dir目录
    cfg.dump(osp.join(cfg.work_dir, osp.basename(args.config)))
    # init logger before other steps，在其他步骤之前初始化日志程序
    timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime())
    log_file = osp.join(cfg.work_dir, f'{timestamp}.log')
    logger = get_root_logger(log_file=log_file, log_level=cfg.log_level)

    # 初始化元数据以记录一些重要信息，例如，环境信息和种子，将被记录
    # init the meta dict to record some important information such as
    # environment info and seed, which will be logged
    meta = dict()
    # log env info
    env_info_dict = collect_env()
    env_info = '\n'.join([f'{k}: {v}' for k, v in env_info_dict.items()])
    dash_line = '-' * 60 + '\n'
    logger.info('Environment info:\n' + dash_line + env_info + '\n' +
                dash_line)
    meta['env_info'] = env_info


    # log some basic info，记录一些基本信息
    logger.info(f'Distributed training: {distributed}')
    logger.info(f'Config: {cfg.text}')

    # set random seeds，设置随机种子
    if args.seed is not None:
        logger.info('Set random seed to {}, deterministic: {}'.format(
            args.seed, args.deterministic))
        set_random_seed(args.seed, deterministic=args.deterministic)
    cfg.seed = args.seed
    meta['seed'] = args.seed

    # 构建模型
    model = build_model(
        cfg.model, train_cfg=cfg.train_cfg, test_cfg=cfg.test_cfg)
    # 创建数据迭代器
    datasets = [build_dataset(cfg.data.train)]
    # 如果提供了数据验证集
    if len(cfg.workflow) == 2:
        val_dataset = copy.deepcopy(cfg.data.val)
        datasets.append(build_dataset(val_dataset))

    # 根据配置加载模型
    if cfg.checkpoint_config is not None:
        # save mmaction version, config file content and class names in
        # checkpoints as meta data
        cfg.checkpoint_config.meta = dict(
            mmaction_version=__version__, config=cfg.text)

    # 对模型进行训练
    train_model(
        model, #
        datasets,
        cfg,
        distributed=distributed,
        validate=args.validate,
        timestamp=timestamp,
        meta=meta)


if __name__ == '__main__':
    main()

可以看到，其最后调用了train_model函数，该函数的注释如下：

def train_model(model, # 根据配置参数创建的模型
                dataset, # 根据配置参数创建的数据迭代器
                cfg, # 配置参数
                distributed=False, # 是否进行分布式训练
                validate=False, # 是否在训练过程中对模型进行验证
                timestamp=None, # 时间戳信息
                meta=None): # 一些额外的信息，如环境变量，随机种子等
    """Train model entry function.

    Args:
        model (nn.Module): The model to be trained.
        dataset (:obj:`Dataset`): Train dataset.
        cfg (dict): The config dict for training.
        distributed (bool): Whether to use distributed training.
            Default: False.
        validate (bool): Whether to do evaluation. Default: False.
        timestamp (str | None): Local time for runner. Default: None.
        meta (dict | None): Meta dict to record some important information.
            Default: None
    """
    # log相关设置
    logger = get_root_logger(log_level=cfg.log_level)

    # prepare data loaders，准备以及加载数据
    dataset = dataset if isinstance(dataset, (list, tuple)) else [dataset]

    # 获取数据迭代器的相关设置
    dataloader_setting = dict(
        videos_per_gpu=cfg.data.get('videos_per_gpu', {}),
        workers_per_gpu=cfg.data.get('workers_per_gpu', {}),
        # cfg.gpus will be ignored if distributed
        num_gpus=len(cfg.gpu_ids),
        dist=distributed,
        seed=cfg.seed)
    dataloader_setting = dict(dataloader_setting,
                              **cfg.data.get('train_dataloader', {}))

    # 对数据迭代器进行设置
    data_loaders = [
        build_dataloader(ds, **dataloader_setting) for ds in dataset
    ]

    # put model on gpus,如果是进行分布式训练
    if distributed:
        find_unused_parameters = cfg.get('find_unused_parameters', False)
        # Sets the `find_unused_parameters` parameter in
        # torch.nn.parallel.DistributedDataParallel
        model = MMDistributedDataParallel(
            model.cuda(),
            device_ids=[torch.cuda.current_device()],
            broadcast_buffers=False,
            find_unused_parameters=find_unused_parameters)
    #如果没有进行分布式训练
    else:
        model = MMDataParallel(
            model.cuda(cfg.gpu_ids[0]), device_ids=cfg.gpu_ids)

    # build runner，创建优化器
    optimizer = build_optimizer(model, cfg.optimizer)

    # 创建模型训练的类
    runner = EpochBasedRunner(
        model,
        optimizer=optimizer,
        work_dir=cfg.work_dir,
        logger=logger,
        meta=meta)

    # an ugly workaround to make .log and .log.json filenames the same
    # 获得时间戳
    runner.timestamp = timestamp

    # fp16 setting，是否使用浮点型16位
    fp16_cfg = cfg.get('fp16', None)
    if fp16_cfg is not None:
        optimizer_config = Fp16OptimizerHook(
            **cfg.optimizer_config, **fp16_cfg, distributed=distributed)

    # 获得优化器的配置
    elif distributed and 'type' not in cfg.optimizer_config:
        optimizer_config = OptimizerHook(**cfg.optimizer_config)
    else:
        optimizer_config = cfg.optimizer_config

    # register hooks，注册训练模型的相关组件，如学习率，优化器，预训练模型等
    runner.register_training_hooks(cfg.lr_config, optimizer_config,
                                   cfg.checkpoint_config, cfg.log_config,
                                   cfg.get('momentum_config', None))

    # 如果是进行分布式训练,则注册一个DistSamplerSeedHook()
    if distributed:
        runner.register_hook(DistSamplerSeedHook())

    # 对模型进行评估验证
    if validate:
        # 根据测试数据的配置参数，创建数据迭代器
        eval_cfg = cfg.get('evaluation', {})
        val_dataset = build_dataset(cfg.data.val, dict(test_mode=True))
        dataloader_setting = dict(
            videos_per_gpu=cfg.data.get('videos_per_gpu', {}),
            workers_per_gpu=cfg.data.get('workers_per_gpu', {}),
            # cfg.gpus will be ignored if distributed
            num_gpus=len(cfg.gpu_ids),
            dist=distributed,
            shuffle=False)
        dataloader_setting = dict(dataloader_setting,
                                  **cfg.data.get('val_dataloader', {}))
        val_dataloader = build_dataloader(val_dataset, **dataloader_setting)

        # 注释训练评估的钩子
        eval_hook = DistEvalHook if distributed else EvalHook
        runner.register_hook(eval_hook(val_dataloader, **eval_cfg))

    # 加载预训练模型
    if cfg.resume_from:
        runner.resume(cfg.resume_from)
    elif cfg.load_from:
        runner.load_checkpoint(cfg.load_from)

    # 正式开始模型训练
    runner.run(data_loaders, cfg.workflow, cfg.total_epochs)

结语

从上面的结构来看，对于模型的训练，主要分为步：
1.模型训练前期准备
2.进行模型训练

但是呢万变不离其宗，几乎所有的套路都是一样的，无非就是：

1.解析参数
2.构建网络模型
3.加载训练测试数据集迭代器
4.迭代训练
5.模型评估保存

好了，总体的结构就简单的介绍到这里，下小结为大家开始讲解代码的每一个细节，我相信大家已经注意到了，最后调用了如下函数：

    # 正式开始模型训练
    runner.run(data_loaders, cfg.workflow, cfg.total_epochs)

该函数才是真正的对数据进行训练，下小节我会对：

    # 创建模型训练的类
    runner = EpochBasedRunner(model,optimizer=optimizer,work_dir=cfg.work_dir,logger=logger,meta=meta)

进行详细的讲解。好好分析其实现过程。