Pytorch实验代码的亿些小细节

你是否有过这样的经历:炼了一大堆的丹,但过了一周回来看结果,忘记了每个模型对应的配置;改了模型中的一个组件,跑起来一个新的训练,这时候测试旧模型却发现结果跟原来不一样了;把所有的训练测试代码写在一个文件里,加入各种if else,最后一个文件上千行,一个周末没看,回来改一个逻辑要找半天……其实这些情况除了深度学习相关的开发,在别的软件开发中也是很常见的,为了解决这些问题,软件行业的开发者形成了很多套路,比如设计模式,提高代码复用性,或者各种最佳实践,比如谷歌、阿里都有一套Java开发最佳实践,各种框架比如客户端的Android,后端的spring,也有各种最佳实践,让开发者的代码更加简洁,更专注于核心的业务实现。在炼丹领域,从2016年至今,各大训练框架互相竞争,互相学习,学术界基于这些框架产出了很多论文,质量越高的论文,往往代码写得也越有条理,最新的论文代码也渐渐形成了一些固定的范式。笔者在学校和工业界阅读过许多论文的开源代码,也基于别人的代码做过不少的改进,有些代码也在工业界落地,自己的炼丹代码也渐渐形成了一些风格,今天就讲讲自己炼丹代码中的一些能让实验更有条理的小习惯。
先上代码,欢迎star:https://github.com/ahangchen/torch_base

代码结构

1
torch_base
2
├── checkpoints # 存放模型的地方
3
├── data # 定义各种用于训练测试的dataset
4
├── eval.py # 测试代码
5
├── loss.py # 定义各种花里胡哨的loss
6
├── metrics.py # 定义各种约定俗成的评估指标
7
├── model # 定义各种实验中的模型
8
├── options.py # 定义各种实验参数,以命令行形式传入
9
├── README.md # 介绍一下自己的repo
10
├── scripts # 各种训练,测试脚本
11
├── train.py # 训练代码
12
└── utils # 各种工具代码
Copied!
checkpoints比较简单,每次训练的模型各自放在一个目录里,scripts目录可以放每次训练或测试用的命令脚本,README.md往往是这个repo的门面,可以放一些介绍性的内容;其他都是代码目录,下面会逐一讲解。

options

首先要介绍的是options.py这个文件,因为这里定义了各种实验参数,其他模块多多少少都会与它有关,受它控制;通常我们需要把各种参数通过某种方式传给程序,比如命令行参数,或者yaml配置文件,我比较习惯用命令行参数,配合pycharm的configuration使用,或者写在scripts目录的脚本里边,都很方便清晰。命令行传参用到了argparse这个lib,这里lib的详细介绍可以看官网教程,这里只挑重点来讲一下:
1
def parse_common_args(parser):
2
parser.add_argument('--model_type', type=str, default='base_model', help='used in model_entry.py')
3
parser.add_argument('--data_type', type=str, default='base_dataset', help='used in data_entry.py')
4
parser.add_argument('--save_prefix', type=str, default='pref', help='some comment for model or test result dir')
5
parser.add_argument('--load_model_path', type=str, default='checkpoints/base_model_pref/0.pth',
6
help='model path for pretrain or test')
7
parser.add_argument('--load_not_strict', action='store_true', help='allow to load only common state dicts')
8
parser.add_argument('--val_list', type=str, default='/data/dataset1/list/base/val.txt',
9
help='val list in train, test list path in test')
10
parser.add_argument('--gpus', nargs='+', type=int)
11
return parser
12
13
def parse_train_args(parser):
14
parser = parse_common_args(parser)
15
...
16
return parser
17
18
def parse_test_args(parser):
19
parser = parse_common_args(parser)
20
...
21
return parser
Copied!
我会在外面初始化一个parser,先用parse_common_args添加训练测试共用的一些参数,在parse_train_args和parse_test_args中调用这个公共的函数,这样可以避免有些参数在训练时写了,测试时忘了写,一跑就报错。parse_train_args解析训练相关的参数,parse_test_args解析测试相关的参数;具体参数和用途如下:
  • parse_common_args
    • 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这个环境变量来控制
  • parse_train_args
    • lrmomentum, beta, weight-decay: optmizer相关参数,在train.py中初始化optimizer
    • model_dir:模型的存储目录,留空,不用传入,会在get_train_model_dir函数中确定这个字段的值,创建对应的目录,填充到args中,方便其他模块获得模型路径
    • train_list:训练集list路径
    • batch_size:训练时的batch size,有人可能会问,为啥测试时不用设置batch size?主要是出于测试时的可视化需求,往往测试需要一张一张forward,所以我习惯将测试batch size为1
    • epochs:模型训练epoch数
    • parse_test_args
      • save_viz:控制是否保存可视化结果的开关
      • result_dir:可视化结果和测试结果的存储目录,留空,不用传入,会在get_test_result_dir中自动生成,自动创建目录,这个目录通常位于模型路径下,形如checkpoints/model_name/checkpoint_num/val_info_save_prefix
使用时,调用prepare_train_args,就会创建一个包含所有公共参数和训练参数的parser,然后创建一个模型目录,并调用save_args函数保存所有参数,返回对应的args。保存参数这一步十分重要,能够避免模型训练完成之后,脚本或命令找不到,忘记自己训练的模型配置这种尴尬局面。
测试时也类似,调用prepare_test_args,创建parser,创建目录,保存参数,并返回对应的args。

data

接下来是data package,在这里,可以为每种数据集定义一个dataset,最好是每个dataset各自形成一个文件,比如list_dataset.py, mem_list_dataset.py,如果多个dataset都写到一个文件里,随着实验进行,各种修修补补下来,代码就会很长,很难查阅。
这里我们还有一个data_entry.py,可以根据命令行参数,以字典的形式,快捷地选择要构造的dataset,如果你有更多的dataset,可以继续扩展这个字典,字典访问是O(1)的,也可以避免一堆if-else的判断。有了dataset,再用pytorch的dataloader接口包一下,可以支持shuffle,多线程加载数据,非常方便。
通常我们还会在data package里放一个augment.py,可以把数据扩增操作都放进去,因为往往多个dataset都需要调用相同的augmentor,所以最好独立出来,在dataset文件中分别调用。

model

这里放的就是各种花里胡哨的模型啦,也是炼丹工作最主要的部分。建议每个模型创建一个package,比如basebetterbest,甚至sota,现代的神经网络结构有一些常用的小组件,比如conv-bn-relu这样的结构,我习惯把它们都放在一个单独的文件submodules.py中,可以在各种任务中复用。
与data_entry类似,我们有一个model_entry.py,在select_model函数中也是通过字典实现参数名和模型的对应,在equip_multi_gpu函数中,可以方便的实现单机多卡,至于多机多卡,我自己用得不多,因为我大多是训练面向无人机上的模型,参数量和计算量要求很小,我们的单机服务器足够train绝大多数模型了,如果是为了更大的batch size加速训练,不如在另一台机器上多跑一组别的实验,总体效率更高。如果大家想看这方面教程,可以留言,我可以补一下对应的代码。

utils

存放各种可复用的util函数或者类,比如一些通用的可视化代码放到viz.py,一些pytorch魔改函数放到torch_utils.py,还有基于tensorboard的存图存曲线的logger.py,这里主要介绍一下这个日志组件:

Recoder

一个数据统计工具,在循环里record每次迭代的数据(比如各种评价指标metrics),在每个epoch训练完成后,调用summary,得到之前统计的指标各自的均值。这个工具在训练时嵌入到Logger中使用,在测试时由于不需要调用tensorboard,所以直接被eval.py调用。

Logger

将tensorboard的SummaryWritter包了一层,包含一个recorder,还有一个SummaryWritter;在训练或验证的每个step以name-value的形式record一下对应的曲线数据,name最好用train/xxxval/xxx这种形式,这样训练和测试的曲线会显示在两个图中,在每个epoch的最后一个step在每次训练或验证的epoch循环结束时,调用一次save_curves保存曲线,调用一次save_checkpoint保存模型参数;这些操作都在下面的train.py中体现。

train.py

终于来到核心的训练代码环节,这里我整了一个trainer,将训练中固定的操作封装成一些函数,需要按实际情况修改的操作封装成另外的函数,这样有新任务来了,只需要修改这些函数就行。现在依次介绍这些函数:
  • __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和反向传播;

eval.py

最后介绍的是测试代码,我把测试的过程包成了一个Evaluator,和trainer也比较类似:
  • __init__:构造函数,初始化命令行参数args,加载模型model并切换到eval模式,初始化测试集的data_loader,设置一个recorder用于统计各种评估指标;
  • eval:测试核心代码,遍历整个测试集,执行forward,得到输入,输出,真值,调用compute_metrics,调用recorder做记录,根据viz_freq,决定这个step是否调用viz_per_epoch可视化并保存结果(与训练不同,往往测试集可视化的内容是要向领导/导师/甲方汇报的,不能存到tensorboard里),循环结束时,调用recorder得到所有的评估指标,并将所有metrics写到result.txt里,避免测试窗口一关就找不到测试结果了。

总结

至此,torch_base这个工程就基本介绍完了,主要还是实践中遇到的各种大坑小坑,逼着自己给工程加上了亿点点小细节,如果基于这个工程去开发新的任务,可以省去一些的脚手架开发工作,专注于model&&data&&metric&&loss&&viz相关的一些内容,让炼丹bring up更快,效率更高,对我来说还是挺有用的,不知道有没有给你一些启发?如果有什么建议也欢迎在issue或者知乎评论区告诉我,感谢你的阅读~
最近更新 3mo ago