argparse
这个lib,这里lib的详细介绍可以看官网教程,这里只挑重点来讲一下:model_type
: 模型的名字,配合model目录和model_entry.py使用;data_type
:数据集的名字,配合data目录和data_entry.py使用;save_prefix
:训练时:实验的名字,可以备注自己改了那些重要组件,具体的参数,会用于创建保存模型的目录;测试时:测试的名字,可以备注测试时做了哪些配置,会用于创建保存测试结果的目录;load_model_path
:模型加载路径,训练时,作为预训练模型路径,测试时,作为待测模型路径,有的人喜欢传入一个模型名字,再传入一个epoch,但其实没啥必要,就算要循环测多个目录,我们也可以写shell生成对应的load_model_path,而且通常只需要测最后一个epoch的模型;load_not_strict
:我写了一个load_match_dict
函数(utils/torch_utils.py),允许加载的模型和当前模型的参数不完全匹配,可多可少,如果打开这个选项,就会调用此函数,这样我们就可以修改模型的某个组件,然后用之前的模型来做预训练啦!如果关闭,就会用torch原本的加载逻辑,要求比较严格的参数匹配;val_list
: 训练时可以传入验证集list,测试时可以传入测试集list;gpus
:可以配置训练或测试时使用的显卡编号,在多卡训练时需要用到,测试时也可以指定显卡编号,绕开其他正在用的显卡,当然你也可以在命令行里export CUDA_VISIBLE_DEVICES这个环境变量来控制lr
,momentum
, beta
, weight-decay
: optmizer相关参数,在train.py中初始化optimizermodel_dir
:模型的存储目录,留空,不用传入,会在get_train_model_dir
函数中确定这个字段的值,创建对应的目录,填充到args中,方便其他模块获得模型路径train_list
:训练集list路径batch_size
:训练时的batch size,有人可能会问,为啥测试时不用设置batch size?主要是出于测试时的可视化需求,往往测试需要一张一张forward,所以我习惯将测试batch size为1epochs
:模型训练epoch数save_viz
:控制是否保存可视化结果的开关result_dir
:可视化结果和测试结果的存储目录,留空,不用传入,会在get_test_result_dir
中自动生成,自动创建目录,这个目录通常位于模型路径下,形如checkpoints/model_name/checkpoint_num/val_info_save_prefixprepare_train_args
,就会创建一个包含所有公共参数和训练参数的parser,然后创建一个模型目录,并调用save_args
函数保存所有参数,返回对应的args。保存参数这一步十分重要,能够避免模型训练完成之后,脚本或命令找不到,忘记自己训练的模型配置这种尴尬局面。prepare_test_args
,创建parser,创建目录,保存参数,并返回对应的args。select_model
函数中也是通过字典实现参数名和模型的对应,在equip_multi_gpu
函数中,可以方便的实现单机多卡,至于多机多卡,我自己用得不多,因为我大多是训练面向无人机上的模型,参数量和计算量要求很小,我们的单机服务器足够train绝大多数模型了,如果是为了更大的batch size加速训练,不如在另一台机器上多跑一组别的实验,总体效率更高。如果大家想看这方面教程,可以留言,我可以补一下对应的代码。metrics
),在每个epoch训练完成后,调用summary,得到之前统计的指标各自的均值。这个工具在训练时嵌入到Logger中使用,在测试时由于不需要调用tensorboard,所以直接被eval.py调用。train/xxx
,val/xxx
这种形式,这样训练和测试的曲线会显示在两个图中,在每个epoch的最后一个step在每次训练或验证的epoch循环结束时,调用一次save_curves保存曲线,调用一次save_checkpoint保存模型参数;这些操作都在下面的train.py中体现。__init__
:构造函数,初始化了命令行参数args
,日志工具(Logger对象)logger
,训练验证的两个dataloader,参数优化器optimizer
,以及模型本身model
,这里我们有三种方式初始化模型:1. 根据模型的构造函数初始化模型参数,2. 使用torch.load加载模型参数,这种方式要求模型参数和我们的模型定义完全匹配,3. 使用load_match_dict加载模型参数,可以找到模型参数和模型定义中,参数量和名字相同的部分进行初始化,适合只改了部分网络结构的模型初始化,作为一种局部pretrain。train
:训练入口,迭代epochs次,每次调用train_per_epoch, val_per_epoch执行训练和测试,再调用logger存储曲线和图像。train_per_epoch
:训练核心代码,将模型切换到训练模式,遍历整个train_loader,调用step进行数据拆包,不同loader返回的数据不同,拆包方式也有差异,还需要用Variable对数据再打包一下,这些操作都独立到step函数里,方便单独修改;再执行模型forward,获取结果,调用compute_metrics计算metrics(训练中也需要观察各种指标,这些指标的计算推荐放在metrics.py),计算loss(各种花里胡哨的loss请放到loss.py),反向传播,在每次迭代中都调用logger的record函数,记录metrics,在最后一个step,调用gen_imgs_to_write,将torch的数据转成图像可视化,各种可视化可以写在viz.py再调用图像的存储(曲线的存储可以放到外面,每个epoch存一次,但图像不行,除非把图传出去,比较蛋疼)。最后根据print_freq,每隔一段时间打印日志方便观察。val_per_epoch
:与训练类似,差别就是模型在eval模式下,不用计算loss和反向传播;__init__
:构造函数,初始化命令行参数args
,加载模型model
并切换到eval模式,初始化测试集的data_loader,设置一个recorder用于统计各种评估指标;eval
:测试核心代码,遍历整个测试集,执行forward,得到输入,输出,真值,调用compute_metrics,调用recorder做记录,根据viz_freq,决定这个step是否调用viz_per_epoch
可视化并保存结果(与训练不同,往往测试集可视化的内容是要向领导/导师/甲方汇报的,不能存到tensorboard里),循环结束时,调用recorder得到所有的评估指标,并将所有metrics写到result.txt
里,避免测试窗口一关就找不到测试结果了。