聊聊炼丹效率

炼丹多年，辗转在不同地方待过，发现还是有相当部分的小伙伴在手动敲命令开所有的实验。高效一点的操作是写一个 bash script 然后用 for loop 把实验跑完，但似乎每次跑实验效率还是比较低，一轮下来也相当累人。

总结下来有以下问题：

1. 效率。在迭代的早期通常会用小模型加小数据。单个模型往往只需占用一张卡，在常见的 8 卡服务器中便留着其他卡在那干等。
2. 认知负载。训一个模型涉及到非常冗长的 pipeline 以及众多超参：数据处理，模型结构，模型参数，训练参数，测试参数。这些可调的节点通常分布在代码、数据或者命令行的参数里面。这让检查、排错和调整极其费劲。往往开完一组实验一天的精力就见底了，更不用说随时出现的错误会让这组实验的结果白白报废。
3. 可用性。怎么在文件系统里面区分这些不同的实验？怎么样高效地区分并分析这些实验的结果？怎么样把在一个 project 里面开发的工具快速迁移到其他的 project？
4. 鲁棒性。如果机器突然宕机，哪些实验需要重新跑一遍？

笔者深受这些痛点折磨，所以一直也在寻找解决方案。现在迭代的工作流感觉还算满意，欢迎评论指教：

1. 把所有的模型、流程改动都映射到命令行参数。模型结构变化用 if/else 或者 switch/case 引出来。这符合把 operation 变成 code 的主流趋势，也能够无缝对接到大多数主流代码框架。
2. 把不关心的默认参数放到一个“命令模板”里面，然后将感兴趣的参数变成变量，将感兴趣的的取值组合写到一个文件里面。把 baseline 取值也写到文件里，方便对比。
3. 把文件当做任务池，用一组 worker (分配了 gpu 资源) 并发地去拉任务，跑任务。每个任务的中间结果写到以超参组合命名的文件夹中，这比时间戳更可读，也有足够区分不同实验，还可以查重防止重复跑实验。用 tensorboard 跟踪训练进程。
4. 写一个评价标准和感兴趣指标的 parser 对所有实验的中间文件进行处理，再把结果拉到 jupyter notebook 或者 excel 表里做可视化和分析。

为了降低认知负担，1 做了第一层简化：把分散在各个地方的调节点全部抽象成命令行参数，这样只需在开发的时候保证每个调节点都正常 work 就行了，查错和决策就到了参数选择的层面。2 则是做了第二层简化：将无关的参数和实验相关的参数区分开，这样就避免了不小心改到别的参数而出错 -- 而写出这些无关的参数又迫使自己思考和过一遍所有可能的影响因子，查漏补缺。3 是实现层面的，任务池 + 并发的模式可以最大化硬件利用率，不需要惦记实验有没有跑完。处理完这几点后，基本上工作流就变成了白天读 paper、沟通、debug、分析结果，晚上回家前列好要跑的实验跑起来，到家后看一下实验是否正常运行，然后就可以倒头睡觉等第二天出结果了。

基本原则讲完之后也贴一个我的 python 实现。工具路径在这里，希望大家实验跑的比谁都快：

https://github.com/simtony/runner

欢迎使用，star，二次开发，提 issue。非常相似的工具有微软的 NNI（https://github.com/microsoft/nni），但是作为跑实验的工具来说太重了，而且做了太多抽象，很难对某个实验的结果做直观的分析。另外这个回答（https://www.zhihu.com/question/384519338/answer/2152639948）提到的工具虽然实现了并发实验的功能，但感觉不太够用。如果有更好的方案可以评论区分享一下。

举个栗子

假设现在我们开发了一个新的 normalization 层叫 “newnorm”，baseline 是 batchnorm。每一个实验涉及 train、checkpoint average 和 test 三个流程。现在希望看不同的 normalization 以及不同的 momentum 参数对结果的影响，对应的配置文件如下：

---
template:
  train: >
    python train.py data-bin/{data}
      --seed 1
      --criterion label_smoothed_cross_entropy
      --arch transformer_iwslt_de_en --share-all-embeddings
      --optimizer adam --adam-betas '(0.9,0.98)' --clip-norm 0.0
      --dropout 0.3 --lr-scheduler inverse_sqrt --warmup-updates 8000
      --lr 0.0015 --min-lr 1e-09
      --label-smoothing 0.1 --weight-decay 0.0001
      --max-tokens 4096 
      --save-dir {_output}
      --tensorboard-logdir {_output}
      --no-save-optimizer-state
      --update-freq 1 --log-format simple --log-interval 50
      --ddp-backend no_c10d
      --keep-last-epochs 5 --early-stop 5
      --normalization {norm} [moment]

  avg: >
    python scripts/average_checkpoints.py --inputs {_output}
      --num-epoch-checkpoints 5 --output {_output}/averaged_model.pt

  test: >
    python generate.py data-bin/{data}
        --max-tokens 4096 --beam 5 --lenpen 1.0 --remove-bpe
        --path {_output}/averaged_model.pt --gen-subset test

default:
  data: iwslt14
  norm: batch
  moment: 0.1

resource: [ 0, 1, 2, 3 ]

---
norm: [ new, batch ]
moment: [ 0.1, 0.05 ]

第一个 yaml doc 作为实验的 specification。template 下面指定了 train, checkpoint average 和 test 的模板命令，其中需要调的参数用 {param} 作为占位符。工具还定义了一些默认的参数，比如这个实验对应的路径 {_output}。指定了要调的超参后，default 里面指定了这些超参的 baseline 值，最后在 resource 里指定了 4 个 worker，每个 worker 对应一个 GPU。

从第二个 yaml doc 开始指定要格点搜的超参。默认会把所有超参组合跑一遍。这里有 4 个任务。同步代码和配置文件到服务器后，直接 run 并发地跑这4个任务:

$ run
Orphan params: set()
Tasks: 4, commands to run: 12
START   gpu: 0, train: 1/ 4, output/Norm_new-Moment_0.1
START   gpu: 1, train: 2/ 4, output/Norm_new-Moment_0.05
START   gpu: 2, train: 3/ 4, output/Norm_batch-Moment_0.1
START   gpu: 3, train: 4/ 4, output/Norm_power-Moment_0.05
START   gpu: 2, avg  : 3/ 4, output/Norm_batch-Moment_0.1
FAIL    gpu: 2, avg  : 3/ 4, output/Norm_batch-Moment_0.1
...

每个输出文件夹里面会写入相应的文件

$ ls output/Norm_batch-Moment_0.1
checkpoint51.pt
checkpoint52.pt
averaged_model.pt
log.train.20220316.030151
log.avg.20220316.030151
log.test.20220316.030151
param
stat

其中 log.* 是每个任务本来会打到命令行里面的 log。param 是每个任务对应的一些参数设定，方便 debug，stat 则是任务状态，分为success 和 fail。这可以用来帮助工具判断是否需要重跑，也可以后期debug。跑实验的过程可以开 tensorboard 跟踪结果，一旦不对劲马上 kill。

实验跑完之后可以开一个 jupyter notebook 写实验结果分析的 parser。在这个例子里面只需要从 log.test 里面读出 BLEU 就好了。写完之后可以调用 Examiner 对所有结果做分析：

from runner.examine import Examiner, latest_log

# define a metric parser for each directory (experiment)
def add_bleu(output_dir, experiment, caches):
    # Each parser follows the same signature
    # It can read/write to a global cache dict `caches`, 
    # and read/write each experiment: 
    # collections.namedtuple("Experiment", ["cache", "metric", "param"])
    latest_test_log = latest_log("test", output_dir)
    bleu = parse_bleu(latest_test_log) #  a user-defined log parser
    experiment.metric["test_bleu"] = bleu
    
examiner = Examiner()  # container for parsed results
# register parser for each directory (experiment)
examiner.add(add_bleu)
# run all parsers for directories matched by regex 
examiner.exam(output="output", regex=".*batch.*")
# print the tsv table with all (different) params and metrics of each experiment
examiner.table()

为什么这样写

格点搜索可以适配大多数调参的场景。首先随机暴力格点搜比较有效的调参方式，特别是当计算资源比较充足的时候。其次做对比实验的时候也会用到参数的格点组合。最后如果不想格点搜，可以手动把想跑的超参组合各自写到配置文件里。

每一个实验都是由一系列顺序执行的命令组成的，比如上述例子的 train - checkpoint average - test。所以相比简单的命令，打包后的顺序执行的命令是更好的任务池的单元。

超参配置模式先后有两个版本。一开始是直接定义一个 config 类并对其操作。但是这样会跟当前 project 深度耦合，换一个代码库就得改很多地方，还会出错，并发部分也不好迁移到其他任务。最后将任务抽象成了一组命令，把修改超参转化成修改任务命令，然后借用了 python 调用 bash 的接口进行并发跑任务。这个方案完美匹配各大主流框架。

并发部分前后迭代了三个版本。第一版的 multiprocessing 最简单，但是对主进程 Ctrl + C 后经常出现 orphan process，还需要查 pid 手动去 kill。第二版的 thread 虽然没有 orphan process 的问题，但是和 multiprocessing 一样需要对全局共享的队列和 io 加锁，也很麻烦。最后收敛到了 asyncio 的 coroutine。后续加 cursor 的用户界面也好写一点。

一些不怎么高级的进阶功能

单个 worker 需要多 GPU 的话可以在 resource 里面用引号框起来：resource: ["0,1", "2,3"] 。

日常需要一组实验在多个机器跑。不同机器卡数不同，需要跑的任务也不同。我先是用了 pycharm 的 deployment -> server group 的配置，让每次 Ctrl + S都会把本地代码 push 到所有服务器上。在上述工具方面做了几个改动：在命令行工具 run 中增加了 -t 和 -r。其中 -t 可以指定跑 yaml 文件中对应_title 参数的任务。-r 指定 gpu index，这样在不同机器通过命令行参数修改资源和 worker 数量。

经常会出现跑一个 train 和多个 test 的情况。为了避免每次 test 都得从头跑一次 train，加了 -c 命令来选择要跑的 command。同时在 yaml 文件里面也加了 _cmd 字段方便按每组实验配置。

一个参数打包很多超参的情况也非常常见。典型的如 Transformer 的 pre/post layernorm 需要同时改 encoder 和 decoder 的 normalization 方式。在切换数据集的时候也是如此，不同数据往往意味着一整套超参的改变。所以在 yaml 的第一个文档中加了 alias 字段，用来将某一个参数的取值映射到一组参数的取值。

有的时候快速试一些改进的时候会懒得把它引到命令行参数里面。这个时候为了将当前结果和已有结果区分开，可以在参数选择中引入 template 里面不用的参数。这些参数只会起到改输出路径名的作用。

公众号后台回复“CVPR 2022”获取论文打包合集下载～