热门标签
热门文章
- 1零基础学习Alfred(二)——常用功能和设置_aflred 设置迁移
- 2Flink 是如何保证 Exactly-once语义的？_flink 的 exactlyonce 语义怎么保证
- 3TimesNet处理UEA数据集，用在InceptionTime上
- 4第01课：中文自然语言处理的完整机器处理流程
- 5Android MediaPlayer整体架构源码分析 -【MediaCodec编解码器插件模块化注册和创建处理流程】【Part 2-B】_blocks-per-second
- 6Twin-Builder — 系统级多物理域数字孪生平台_twim builder 模块拼接
- 7python用于人工智能的例子,python人工智能有趣例子_python人工智能案例
- 8爬虫-基于python的旅游数据可视化-计算机毕业设计源码81319_python旅游数据可视化
- 9生产制造园区数字孪生3D大屏展示提升运营效益_可视化大屏实现3d效果的技术有哪些?
- 10启信宝商业大数据助力全国经济普查
当前位置: article > 正文
mmdetection2.0 | train.py相关源码详解（一）_mmdetect 2.0 coco.py解析
作者：IT小白 | 2024-03-29 23:12:16
赞
踩
mmdetect 2.0 coco.py解析
本文章从源码入手，详细的剖析 MMDetection 的训练细节。理解它是如何实现通用，灵活的训练的。
MMDetection 支持单机单卡、多卡训练。而且对于训练部分的代码和模型、数据集的代码的耦合性极低。真正实现了支持各种训练模式、各种模型和数据集的通用训练模板。
- 1
- 2
一、tools/train.py
tools/train.py 是单机单卡训练时需要运行的文件，其主要使用方法为：
python tools/train.py ${CONFIG_FILE} [optional arguments]
- 1
如果为单机多卡，需要运行 tools/dist_train.sh，注意此文件不支持多机多卡训练：
./tools/dist_train.sh ${CONFIG_FILE} ${GPU_NUM} [optional arguments]
- 1
可选参数如下：
=========== optional arguments ===========
# --work-dir 存储日志和模型的目录
# --resume-from 加载 checkpoint 的目录
# --no-validate 是否在训练的时候进行验证
# 互斥组：
# --gpus 使用的 GPU 数量
# --gpu_ids 使用指定 GPU 的 id
# --seed 随机数种子
# --deterministic 是否设置 cudnn 为确定性行为
# --options 其他参数
# --launcher 分布式训练使用的启动器，可以为：['none', 'pytorch', 'slurm', 'mpi']
# none：不启动分布式训练，dist_train.sh 中默认使用 pytorch 启动。
# --local_rank 本地进程编号，此参数 torch.distributed.launch 会自动传入。
- 1
- 2
- 3
- 4
- 5
- 6
- 7
- 8
- 9
- 10
- 11
- 12
- 13
我们来看一下 dist_train.sh 里面究竟是什么？
可以看出 dist_train.sh 的本质就是使用 torch.distributed.launch（这是分布式的辅助启动工具） 运行 tools/train.py。
torch.distributed.launch 需要使用 python -m 来运行，-m 是把一个模块当做脚本来运行的一个参数。一般情况下可以给 torch.distributed.launch 传如下的参数
--nproc_per_node：表示每台机器的 GPU 数量
--nnodes：表示机器的数量
--node_rank：机器的排名，如果为 0 代表是 master 节点（机器）。
--master_addr：master 节点的 IP 地址
--master_port：master 节点开放的端口号
- 1
- 2
- 3
- 4
- 5
可以看到，因为 dist_train.sh 只支持单机多卡训练，所以只传参了 --nproc_per_node（GPU 个数）和 --master_port（开放的端口号）
#!/usr/bin/env bash
CONFIG=$1
GPUS=$2
PORT=${PORT:-29500}
PYTHONPATH="$(dirname $0)/..":$PYTHONPATH \
python -m torch.distributed.launch --nproc_per_node=$GPUS --master_port=$PORT \
$(dirname "$0")/train.py $CONFIG --launcher pytorch ${@:3}
- 1
- 2
- 3
- 4
- 5
- 6
- 7
- 8
- 9
对于 tools/train.py 其主要的流程为：
（一）从命令行和配置文件获取参数配置
（二）构建模型
# 构建模型: 需要传入 cfg.model，cfg.train_cfg，cfg.test_cfg
model = build_detector(
cfg.model, train_cfg=cfg.train_cfg, test_cfg=cfg.test_cfg)
- 1
- 2
- 3
（三）构建数据集
# 构建数据集: 需要传入 cfg.data.train，表明是训练集
datasets = [build_dataset(cfg.data.train)]
- 1
- 2
（四）训练模型
# 训练检测器：需要传入模型、数据集、配置参数等
train_detector(
model,
datasets,
cfg,
distributed=distributed,
validate=(not args.no_validate),
timestamp=timestamp,
meta=meta)
- 1
- 2
- 3
- 4
- 5
- 6
- 7
- 8
- 9
所以对于 train.py 来说，首先从命令行和配置文件读取配置，然后分别用 build_detector、build_dataset 构建模型和数据集，最后将模型和数据集传入 train_detector 进行训练。
下面我们来看一下源码：
import argparse import copy import os import os.path as osp import time import warnings import mmcv import torch # Config 用于读取配置文件, DictAction 将命令行字典类型参数转化为 key-value 形式 from mmcv import Config, DictAction from mmcv.runner import get_dist_info, init_dist from mmcv.utils import get_git_hash from mmdet import __version__ from mmdet.apis import set_random_seed, train_detector from mmdet.datasets import build_dataset from mmdet.models import build_detector from mmdet.utils import collect_env, get_root_logger # python tools/train.py ${CONFIG_FILE} [optional arguments] # =========== optional arguments =========== # --work-dir 存储日志和模型的目录 # --resume-from 加载 checkpoint 的目录 # --no-validate 是否在训练的时候进行验证 # 互斥组： # --gpus 使用的 GPU 数量 # --gpu_ids 使用指定 GPU 的 id # --seed 随机数种子 # --deterministic 是否设置 cudnn 为确定性行为 # --options 其他参数 # --launcher 分布式训练使用的启动器，可以为：['none', 'pytorch', 'slurm', 'mpi'] # none：不启动分布式训练，dist_train.sh 中默认使用 pytorch 启动。 # --local_rank 本地进程编号，此参数 torch.distributed.launch 会自动传入。 def parse_args(): parser = argparse.ArgumentParser(description='Train a detector') parser.add_argument('config', help='train config file path') 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') # action: store (默认, 表示保存参数) # action: store_true, store_false (如果指定参数, 则为 True, False) parser.add_argument( '--no-validate', action='store_true', help='whether not to evaluate the checkpoint during training') # --------- 创建一个互斥组. argparse 将会确保互斥组中的参数只能出现一个 --------- group_gpus = parser.add_mutually_exclusive_group() group_gpus.add_argument( '--gpus', type=int, help='number of gpus to use ' '(only applicable to non-distributed training)') # 可以使用 python train.py --gpu-ids 0 1 2 3 指定使用的 GPU id # 参数结果：[0, 1, 2, 3] # nargs = '*'：参数个数可以设置0个或n个 # nargs = '+'：参数个数可以设置1个或n个 # nargs = '?'：参数个数可以设置0个或1个 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') parser.add_argument( '--deterministic', action='store_true', help='whether to set deterministic options for CUDNN backend.') # 其他参数: 可以使用 --options a=1,2,3 指定其他参数 # 参数结果: {'a': [1, 2, 3]} parser.add_argument( '--options', nargs='+', action=DictAction, help='override some settings in the used config, the key-value pair ' 'in xxx=yyy format will be merged into config file (deprecate), ' 'change to --cfg-options instead.') parser.add_argument( '--cfg-options', nargs='+', action=DictAction, help='override some settings in the used config, the key-value pair ' 'in xxx=yyy format will be merged into config file. If the value to ' 'be overwritten is a list, it should be like key="[a,b]" or key=a,b ' 'It also allows nested list/tuple values, e.g. key="[(a,b),(c,d)]" ' 'Note that the quotation marks are necessary and that no white space ' 'is allowed.') # 如果使用 dist_utils.sh 进行分布式训练, launcher 默认为 pytorch parser.add_argument( '--launcher', choices=['none', 'pytorch', 'slurm', 'mpi'], default='none', help='job launcher') # 本地进程编号，此参数 torch.distributed.launch 会自动传入。 parser.add_argument('--local_rank', type=int, default=0) args = parser.parse_args() # 如果环境中没有 LOCAL_RANK，就设置它为当前的 local_rank if 'LOCAL_RANK' not in os.environ: os.environ['LOCAL_RANK'] = str(args.local_rank) if args.options and args.cfg_options: raise ValueError( '--options and --cfg-options cannot be both ' 'specified, --options is deprecated in favor of --cfg-options') if args.options: warnings.warn('--options is deprecated in favor of --cfg-options') args.cfg_options = args.options return args def main(): args = parse_args() cfg = Config.fromfile(args.config) # 从文件读取配置 # 从命令行读取额外的配置 if args.cfg_options is not None: cfg.merge_from_dict(args.cfg_options) # import modules from string list. if cfg.get('custom_imports', None): from mmcv.utils import import_modules_from_strings import_modules_from_strings(**cfg['custom_imports']) # set cudnn_benchmark，设置True 可以加速输入大小固定的模型. 如：SSD300 if cfg.get('cudnn_benchmark', False): torch.backends.cudnn.benchmark = True # work_dir is determined in this priority: CLI > segment in file > filename # 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 # 当 work_dir 为 None 的时候, 使用 ./work_dir/配置文件名 作为默认工作目录 elif cfg.get('work_dir', None) is None: # use config filename as default work_dir if cfg.work_dir is None # os.path.basename(path) 返回文件名 # os.path.splitext(path) 分割路径, 返回路径名和文件扩展名的元组 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 # gpu id 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 distributed env first, since logger depends on the dist info. # 如果 launcher 为 none，不启用分布式训练。不使用 dist_train.sh 默认参数为 none. if args.launcher == 'none': distributed = False # launcher 不为 none，启用分布式训练。使用 dist_train.sh，会传 ‘pytorch’ else: distributed = True # 初始化 dist 里面会调用 init_process_group init_dist(args.launcher, **cfg.dist_params) # re-set gpu_ids with distributed training mode _, world_size = get_dist_info() cfg.gpu_ids = range(world_size) # create work_dir mmcv.mkdir_or_exist(osp.abspath(cfg.work_dir)) # dump config cfg.dump(osp.join(cfg.work_dir, osp.basename(args.config))) # init the 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 meta['config'] = cfg.pretty_text # log some basic info logger.info(f'Distributed training: {distributed}') logger.info(f'Config:\n{cfg.pretty_text}') # 设置随机化种子 if args.seed is not None: logger.info(f'Set random seed to {args.seed}, ' f'deterministic: {args.deterministic}') set_random_seed(args.seed, deterministic=args.deterministic) cfg.seed = args.seed meta['seed'] = args.seed meta['exp_name'] = osp.basename(args.config) # 构建模型: 需要传入 cfg.model, cfg.train_cfg, cfg.test_cfg model = build_detector( cfg.model, train_cfg=cfg.get('train_cfg'), test_cfg=cfg.get('test_cfg')) model.init_weights() # 构建数据集: 需要传入 cfg.data.train datasets = [build_dataset(cfg.data.train)] # workflow 代表流程： # [('train', 2), ('val', 1)] 就代表，训练两个 epoch 验证一个 epoch if len(cfg.workflow) == 2: val_dataset = copy.deepcopy(cfg.data.val) val_dataset.pipeline = cfg.data.train.pipeline datasets.append(build_dataset(val_dataset)) if cfg.checkpoint_config is not None: # save mmdet version, config file content and class names in # checkpoints as meta data cfg.checkpoint_config.meta = dict( mmdet_version=__version__ + get_git_hash()[:7], CLASSES=datasets[0].CLASSES) # add an attribute for visualization convenience model.CLASSES = datasets[0].CLASSES # 训练检测器, 传入：模型, 数据集, config 等 train_detector( model, datasets, cfg, distributed=distributed, validate=(not args.no_validate), timestamp=timestamp, meta=meta) if __name__ == '__main__': 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
- 76
- 77
- 78
- 79
- 80
- 81
- 82
- 83
- 84
- 85
- 86
- 87
- 88
- 89
- 90
- 91
- 92
- 93
- 94
- 95
- 96
- 97
- 98
- 99
- 100
- 101
- 102
- 103
- 104
- 105
- 106
- 107
- 108
- 109
- 110
- 111
- 112
- 113
- 114
- 115
- 116
- 117
- 118
- 119
- 120
- 121
- 122
- 123
- 124
- 125
- 126
- 127
- 128
- 129
- 130
- 131
- 132
- 133
- 134
- 135
- 136
- 137
- 138
- 139
- 140
- 141
- 142
- 143
- 144
- 145
- 146
- 147
- 148
- 149
- 150
- 151
- 152
- 153
- 154
- 155
- 156
- 157
- 158
- 159
- 160
- 161
- 162
- 163
- 164
- 165
- 166
- 167
- 168
- 169
- 170
- 171
- 172
- 173
- 174
- 175
- 176
- 177
- 178
- 179
- 180
- 181
- 182
- 183
- 184
- 185
- 186
- 187
- 188
- 189
- 190
- 191
- 192
- 193
- 194
- 195
- 196
- 197
- 198
- 199
- 200
- 201
- 202
- 203
- 204
- 205
- 206
- 207
- 208
- 209
- 210
- 211
- 212
- 213
- 214
- 215
- 216
- 217
- 218
- 219
- 220
- 221
- 222
- 223
- 224
- 225
- 226
- 227
- 228
- 229
- 230
- 231
- 232
- 233
- 234
- 235
- 236
- 237
我们分别来看看与 train.py 相关的核心函数：
1、init_dist：
此函数负责调用 init_process_group，完成分布式的初始化。在运行 dist_train.py 训练时，默认传递的 launcher 是 ‘pytorch’。所以此函数会进一步调用 _init_dist_pytorch 来完成初始化。
因为 torch.distributed 可以采用单进程控制多 GPU，也可以一个进程控制一个 GPU。一个进程控制一个 GPU 是目前Pytorch中，无论是单节点还是多节点，进行数据并行训练最快的方式。在 mmdet 中也是这么实现的。既然是单个进程控制单个 GPU，那么我么就需要绑定当前进程控制的是哪个 GPU。可以理解为在使用 torch.distributed.launch 运行 py 文件时。 它会多次调用 py 文件，每个 py 文件控制一个 GPU。并向每个 py 文件传参 --local_rank。（local_rank 是在这台机器上的本地进程编号）这样对于每个 py 文件，都能拿到传入的本地进程编号，我们只需要把当前进程绑定到指定的 GPU 即可。
在 _init_dist_pytorch 中就会设置当前进程控制的默认 GPU（torch.cuda.set_device），再使用 dist.init_process_group 初始化，初始化的方式为默认的 env://，即环境变量的方式。使用 env:// 方式初始化就需要用 torch.distributed.launch 运行 py 文件，torch.distributed.launch 会根据传入的参数设置环境变量，并运行 py 文件。
# Copyright (c) Open-MMLab. All rights reserved. import functools import os import subprocess import torch import torch.distributed as dist import torch.multiprocessing as mp from mmcv.utils import TORCH_VERSION def init_dist(launcher, backend='nccl', **kwargs): if mp.get_start_method(allow_none=True) is None: mp.set_start_method('spawn') # 默认进到这里 if launcher == 'pytorch': # 调用下面的 _init_dist_pytorch 函数来初始化。 _init_dist_pytorch(backend, **kwargs) elif launcher == 'mpi': _init_dist_mpi(backend, **kwargs) elif launcher == 'slurm': _init_dist_slurm(backend, **kwargs) else: raise ValueError(f'Invalid launcher type: {launcher}') def _init_dist_pytorch(backend, **kwargs): # rank 是所有进程的总编号，算上本地进程，从 0 开始。 rank = int(os.environ['RANK']) num_gpus = torch.cuda.device_count() # 这里也可以使用命令行传递来的 --local_rank。 torch.cuda.set_device(rank % num_gpus) dist.init_process_group(backend=backend, **kwargs)
- 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
2、set_random_seed：
此函数会对 python、numpy、torch 都设置随机数种子。
保持随机数种子相同时，卷积的结果在CPU上相同，在GPU上仍然不相同。这是因为，cudnn卷积行为的不确定性。使用 torch.backends.cudnn.deterministic = True 可以解决。
cuDNN 使用非确定性算法，并且可以使用 torch.backends.cudnn.enabled = False 来进行禁用。如果设置为 torch.backends.cudnn.enabled = True，说明设置为使用非确定性算法（即会自动寻找最适合当前配置的高效算法，来达到优化运行效率的问题）
一般来讲，应该遵循以下准则：
- 如果网络的输入数据维度或类型上变化不大，设置torch.backends.cudnn.benchmark = true 可以增加运行效率
- 如果网络的输入数据在每次 iteration 都变化的话，会导致 cnDNN 每次都会去寻找一遍最优配置，这样反而会降低运行效率。设置
torch.backends.cudnn.benchmark = False 避免重复搜索。
def set_random_seed(seed, deterministic=False):
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
# manual_seed_all 是为所有 GPU 都设置随机数种子。
torch.cuda.manual_seed_all(seed)
if deterministic:
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
- 1
- 2
- 3
- 4
- 5
- 6
- 7
- 8
- 9
- 10
3、get_root_logger：
get_root_logger 调用 get_logger 函数获取 logger 对象。
import logging
from mmcv.utils import get_logger
def get_root_logger(log_file=None, log_level=logging.INFO):
logger = get_logger(name='mmdet', log_file=log_file, log_level=log_level)
return logger
- 1
- 2
- 3
- 4
- 5
- 6
- 7
- 8
- 9
这里实现的 get_logger 函数非常灵活，如果传入相同的 log 的 name，会返回配置相同的 log。传入以点分割的日志名称的子模块，也会返回相同的 log。如：a 和 a.b 会返回相同的 log。
如果传入 log_file 会保存 log 的输出到 log_file 指定的路径，如果不传入 log_file，不保存日志的输出。只在控制台输出。
下面我们来分析一下源码：
import logging import torch.distributed as dist # 记录是否创建过 name 对应的 log，如果创建过设置为 True logger_initialized = {} def get_logger(name, log_file=None, log_level=logging.INFO): # 获取 log 对象。 logger = logging.getLogger(name) # 如果已经创建过，直接返回 if name in logger_initialized: return logger # 如果是创建过的以 ‘.’ 分割的子模块，也直接返回 for logger_name in logger_initialized: if name.startswith(logger_name): return logger stream_handler = logging.StreamHandler() handlers = [stream_handler] # 获取当前的 rank（总进程编号） if dist.is_available() and dist.is_initialized(): rank = dist.get_rank() else: rank = 0 # 只有 rank 0（master 节点的 local_rank 为 0 的进程）的主机才保存日志 if rank == 0 and log_file is not None: file_handler = logging.FileHandler(log_file, 'w') handlers.append(file_handler) formatter = logging.Formatter( '%(asctime)s - %(name)s - %(levelname)s - %(message)s') for handler in handlers: handler.setFormatter(formatter) handler.setLevel(log_level) logger.addHandler(handler) # 对于非 rank 为 0 的进程，只有 error 以上的信息才会显示 if rank == 0: logger.setLevel(log_level) else: logger.setLevel(logging.ERROR) # 将 log name 对应的值设为 True，表示创建过。 logger_initialized[name] = True return logger
- 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
因为在 train.py 中主要调用：构建模型（build_detector），构建数据集（build_dataset）,训练模型（train_detector）的函数，我们下来分别看看这三个函数的源码。
二、build_detector（mmdet/models/builder.py）
build_detector 函数将配置文件config中的：model、train_cfg 和 test_cfg 三部分传入参数。
下面以 faster_rcnn_r50_fpn_1x_coco.py 配置文件来举例：
具体在faster_rcnn_r50_fpn.py文件中
model
model = dict( type='FasterRCNN', pretrained='torchvision://resnet50', backbone=dict( type='ResNet', depth=50, num_stages=4, out_indices=(0, 1, 2, 3), frozen_stages=1, norm_cfg=dict(type='BN', requires_grad=True), norm_eval=True, style='pytorch'), neck=dict( type='FPN', in_channels=[256, 512, 1024, 2048], out_channels=256, num_outs=5), rpn_head=dict( type='RPNHead', in_channels=256, feat_channels=256, anchor_generator=dict( type='AnchorGenerator', scales=[8], ratios=[0.5, 1.0, 2.0], strides=[4, 8, 16, 32, 64]), bbox_coder=dict( type='DeltaXYWHBBoxCoder', target_means=[.0, .0, .0, .0], target_stds=[1.0, 1.0, 1.0, 1.0]), loss_cls=dict( type='CrossEntropyLoss', use_sigmoid=True, loss_weight=1.0), loss_bbox=dict(type='L1Loss', loss_weight=1.0)), roi_head=dict( type='StandardRoIHead', bbox_roi_extractor=dict( type='SingleRoIExtractor', roi_layer=dict(type='RoIAlign', output_size=7, sampling_ratio=0), out_channels=256, featmap_strides=[4, 8, 16, 32]), bbox_head=dict( type='Shared2FCBBoxHead', in_channels=256, fc_out_channels=1024, roi_feat_size=7, num_classes=80, bbox_coder=dict( type='DeltaXYWHBBoxCoder', target_means=[0., 0., 0., 0.], target_stds=[0.1, 0.1, 0.2, 0.2]), reg_class_agnostic=False, loss_cls=dict( type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0), loss_bbox=dict(type='L1Loss', loss_weight=1.0))))
- 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
train_cfg
train_cfg = dict( rpn=dict( assigner=dict( type='MaxIoUAssigner', pos_iou_thr=0.7, neg_iou_thr=0.3, min_pos_iou=0.3, match_low_quality=True, ignore_iof_thr=-1), sampler=dict( type='RandomSampler', num=256, pos_fraction=0.5, neg_pos_ub=-1, add_gt_as_proposals=False), allowed_border=-1, pos_weight=-1, debug=False), rpn_proposal=dict( nms_across_levels=False, nms_pre=2000, nms_post=1000, max_num=1000, nms_thr=0.7, min_bbox_size=0), rcnn=dict( assigner=dict( type='MaxIoUAssigner', pos_iou_thr=0.5, neg_iou_thr=0.5, min_pos_iou=0.5, match_low_quality=False, ignore_iof_thr=-1), sampler=dict( type='RandomSampler', num=512, pos_fraction=0.25, neg_pos_ub=-1, add_gt_as_proposals=True), pos_weight=-1, debug=False))
- 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
test_cfg
test_cfg = dict(
rpn=dict(
nms_across_levels=False,
nms_pre=1000,
nms_post=1000,
max_num=1000,
nms_thr=0.7,
min_bbox_size=0),
rcnn=dict(
score_thr=0.05,
nms=dict(type='nms', iou_threshold=0.5),
max_per_img=100)
# soft-nms is also supported for rcnn testing
# e.g., nms=dict(type='soft_nms', iou_threshold=0.5, min_score=0.05)
)
- 1
- 2
- 3
- 4
- 5
- 6
- 7
- 8
- 9
- 10
- 11
- 12
- 13
- 14
- 15
运行时会将上面的三个值作为参数传入 build_detector 函数，build_detector 函数会调用 build 函数，build 函数调用 build_from_cfg 函数构建检测器对象。其中 train_cfg 和 test_cfg 作为默认参数用于构建 detector 对象。
def build(cfg, registry, default_args=None):
if isinstance(cfg, list):
modules = [
build_from_cfg(cfg_, registry, default_args) for cfg_ in cfg
]
return nn.Sequential(*modules)
else:
# 调用 build_from_cfg 用来根据 config 字典构建 registry 里面的对象
return build_from_cfg(cfg, registry, default_args)
def build_detector(cfg, train_cfg=None, test_cfg=None):
# 调用 build 函数，传入 cfg, registry 对象，
# 把 train_cfg 和 test_cfg 作为默认字典传入
return build(cfg, DETECTORS, dict(train_cfg=train_cfg, test_cfg=test_cfg))
- 1
- 2
- 3
- 4
- 5
- 6
- 7
- 8
- 9
- 10
- 11
- 12
- 13
- 14
- 15
build_from_cfg 在 mmcv/utils/registery.py 中。其中参数 cfg 字典中的 type 键所对应的值表示需要创建的对象的类型。build_from_cfg 会自动在 Registry 注册的类中找到需要创建的类，并传入默认参数实例化。
def build_from_cfg(cfg, registry, default_args=None): """Build a module from config dict. Args: cfg (dict): Config dict. It should at least contain the key "type". registry (:obj:`Registry`): The registry to search the type from. default_args (dict, optional): Default initialization arguments. Returns: object: The constructed object. """ if not isinstance(cfg, dict): raise TypeError(f'cfg must be a dict, but got {type(cfg)}') if 'type' not in cfg: raise KeyError( f'the cfg dict must contain the key "type", but got {cfg}') if not isinstance(registry, Registry): raise TypeError('registry must be an mmcv.Registry object, ' f'but got {type(registry)}') if not (isinstance(default_args, dict) or default_args is None): raise TypeError('default_args must be a dict or None, ' f'but got {type(default_args)}') args = cfg.copy() # 获取 type 对应的值 obj_type = args.pop('type') if is_str(obj_type): # 获取需要创建的对象 obj_cls = registry.get(obj_type) if obj_cls is None: raise KeyError( f'{obj_type} is not in the {registry.name} registry') elif inspect.isclass(obj_type): obj_cls = obj_type else: raise TypeError( f'type must be a str or valid type, but got {type(obj_type)}') # 如果 default_args 不是 None，传入默认值再实例化。 if default_args is not None: for name, value in default_args.items(): args.setdefault(name, value) return obj_cls(**args)
- 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
那么什么是 registry？
registry 就是注册类，将一个字符串和类关联起来。如果索引字符串就会获得类。Registry 是注册所需要的类，可以用它来注册类。我们可以使用如下的方式来注册类。
backbones = Registry('backbone')
@backbones.register_module()
class ResNet:
pass
backbones = Registry('backbone')
@backbones.register_module(name='mnet')
class MobileNet:
pass
backbones = Registry('backbone')
class ResNet:
pass
backbones.register_module(ResNet)
- 1
- 2
- 3
- 4
- 5
- 6
- 7
- 8
- 9
- 10
- 11
- 12
- 13
- 14
下面是 Registry 类的代码，它的内部维护了一个已经注册的类的字典 ——_module_dict。每当注册一个类就在字典里添加一个字符串（默认为类名）与类的映射。register_module 方法，利用装饰器将类名和类添加到 _module_dict 中。对于注册的模块可以通过 build_from_cfg 来构建。
import inspect import warnings from functools import partial from .misc import is_str class Registry: """A registry to map strings to classes. Args: name (str): Registry name. """ def __init__(self, name): self._name = name # 已经注册的类的字典 self._module_dict = dict() def __len__(self): return len(self._module_dict) def __contains__(self, key): return self.get(key) is not None def __repr__(self): format_str = self.__class__.__name__ + \ f'(name={self._name}, ' \ f'items={self._module_dict})' return format_str @property def name(self): return self._name @property def module_dict(self): return self._module_dict def get(self, key): """Get the registry record. Args: key (str): The class name in string format. Returns: class: The corresponding class. """ return self._module_dict.get(key, None) def _register_module(self, module_class, module_name=None, force=False): if not inspect.isclass(module_class): raise TypeError('module must be a class, ' f'but got {type(module_class)}') if module_name is None: module_name = module_class.__name__ if not force and module_name in self._module_dict: raise KeyError(f'{module_name} is already registered ' f'in {self.name}') self._module_dict[module_name] = module_class def deprecated_register_module(self, cls=None, force=False): warnings.warn( 'The old API of register_module(module, force=False) ' 'is deprecated and will be removed, please use the new API ' 'register_module(name=None, force=False, module=None) instead.') if cls is None: return partial(self.deprecated_register_module, force=force) self._register_module(cls, force=force) return cls def register_module(self, name=None, force=False, module=None): """Register a module. A record will be added to `self._module_dict`, whose key is the class name or the specified name, and value is the class itself. It can be used as a decorator or a normal function. Example: >>> backbones = Registry('backbone') >>> @backbones.register_module() >>> class ResNet: >>> pass >>> backbones = Registry('backbone') >>> @backbones.register_module(name='mnet') >>> class MobileNet: >>> pass >>> backbones = Registry('backbone') >>> class ResNet: >>> pass >>> backbones.register_module(ResNet) Args: name (str | None): The module name to be registered. If not specified, the class name will be used. force (bool, optional): Whether to override an existing class with the same name. Default: False. module (type): Module class to be registered. """ if not isinstance(force, bool): raise TypeError(f'force must be a boolean, but got {type(force)}') # NOTE: This is a walkaround to be compatible with the old api, # while it may introduce unexpected bugs. if isinstance(name, type): return self.deprecated_register_module(name, force=force) # use it as a normal method: x.register_module(module=SomeClass) if module is not None: self._register_module( module_class=module, module_name=name, force=force) return module # raise the error ahead of time if not (name is None or isinstance(name, str)): raise TypeError(f'name must be a str, but got {type(name)}') # use it as a decorator: @x.register_module() def _register(cls): self._register_module( module_class=cls, module_name=name, force=force) return cls return _register
- 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
- 99
- 100
- 101
- 102
- 103
- 104
- 105
- 106
- 107
- 108
- 109
- 110
- 111
- 112
- 113
- 114
- 115
- 116
- 117
- 118
- 119
- 120
- 121
- 122
- 123
- 124
- 125
- 126
三、build_dataset（mmdet/datasets/builder）
build_dataset 也类似，通过调用 build_from_cfg 创建。
def build_dataset(cfg, default_args=None): from .dataset_wrappers import (ConcatDataset, RepeatDataset, ClassBalancedDataset) if isinstance(cfg, (list, tuple)): dataset = ConcatDataset([build_dataset(c, default_args) for c in cfg]) elif cfg['type'] == 'RepeatDataset': dataset = RepeatDataset( build_dataset(cfg['dataset'], default_args), cfg['times']) elif cfg['type'] == 'ClassBalancedDataset': dataset = ClassBalancedDataset( build_dataset(cfg['dataset'], default_args), cfg['oversample_thr']) elif isinstance(cfg.get('ann_file'), (list, tuple)): dataset = _concat_dataset(cfg, default_args) else: dataset = build_from_cfg(cfg, DATASETS, default_args) return dataset
- 1
- 2
- 3
- 4
- 5
- 6
- 7
- 8
- 9
- 10
- 11
- 12
- 13
- 14
- 15
- 16
- 17
四、train_detector（mmdet/apis/train.py）
train_detector 的主要流程为：
（一）构建 data loaders：
data_loaders = [
build_dataloader(
ds,
cfg.data.samples_per_gpu,
cfg.data.workers_per_gpu,
# cfg.gpus will be ignored if distributed
len(cfg.gpu_ids),
dist=distributed,
seed=cfg.seed) for ds in dataset
]
- 1
- 2
- 3
- 4
- 5
- 6
- 7
- 8
- 9
- 10
（二）构建分布式处理对象：
model = MMDistributedDataParallel(
model.cuda(),
device_ids=[torch.cuda.current_device()],
broadcast_buffers=False,
find_unused_parameters=find_unused_parameters)
- 1
- 2
- 3
- 4
- 5
（三）构建优化器：
optimizer = build_optimizer(model, cfg.optimizer)
- 1
（四）创建 EpochBasedRunner 并进行训练：
runner = EpochBasedRunner(
model,
optimizer=optimizer,
work_dir=cfg.work_dir,
logger=logger,
meta=meta)
- 1
- 2
- 3
- 4
- 5
- 6
我们来看一下源码：
def train_detector(model, dataset, cfg, distributed=False, validate=False, timestamp=None, meta=None): # 获取 logger logger = get_root_logger(cfg.log_level) # ==================== 构建 data loaders ==================== dataset = dataset if isinstance(dataset, (list, tuple)) else [dataset] # 获得 samples_per_gpu if 'imgs_per_gpu' in cfg.data: logger.warning('"imgs_per_gpu" is deprecated in MMDet V2.0. ' 'Please use "samples_per_gpu" instead') if 'samples_per_gpu' in cfg.data: logger.warning( f'Got "imgs_per_gpu"={cfg.data.imgs_per_gpu} and ' f'"samples_per_gpu"={cfg.data.samples_per_gpu}, "imgs_per_gpu"' f'={cfg.data.imgs_per_gpu} is used in this experiments') else: logger.warning( 'Automatically set "samples_per_gpu"="imgs_per_gpu"=' f'{cfg.data.imgs_per_gpu} in this experiments') cfg.data.samples_per_gpu = cfg.data.imgs_per_gpu data_loaders = [ build_dataloader( ds, cfg.data.samples_per_gpu, cfg.data.workers_per_gpu, # cfg.gpus will be ignored if distributed len(cfg.gpu_ids), dist=distributed, seed=cfg.seed) for ds in dataset ] # ==================== 构建分布式处理对象 ===================== # 如果是多卡会进入此 if 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) # ====================== 构建优化器 ========================== optimizer = build_optimizer(model, cfg.optimizer) # ============= 创建 EpochBasedRunner 并进行训练 ============== 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 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)) if distributed: runner.register_hook(DistSamplerSeedHook()) # register eval hooks if validate: val_dataset = build_dataset(cfg.data.val, dict(test_mode=True)) val_dataloader = build_dataloader( val_dataset, samples_per_gpu=1, workers_per_gpu=cfg.data.workers_per_gpu, dist=distributed, shuffle=False) eval_cfg = cfg.get('evaluation', {}) 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
- 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
- 99
- 100
- 101
- 102
- 103
- 104
- 105
在本篇文章（一）中，主要讲解了，train.py 中的主要流程，train.py 中的重要的函数以及函数的具体实现。但是 train_detector 只讲了流程，并没有拆开详细讲解。在下一小结中我们会详细讲解 train_detector 的每一步究竟做了什么。
参考：
https://zhuanlan.zhihu.com/p/163747610
声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/IT小白/article/detail/337878
推荐阅读
相关标签
