道为某个特定的机器学习项目服务,最典型的就是机器学习相关的比赛。在机器学习比赛中,通常会存在很多棘手的问题,例如实验管理、复现,代码混乱等...而其中的大多数问题,究其根本都是开发人员的懈怠导致的。道尝试通过约定或规则,来尽量避免这些问题。
道的设计或开发遵循以下一些原则:
- 使用工具
道会尽可能得使用工具来辅助或规范用户代码,例如,类型检查 mypy,代码风格 black 等。另外,道本身还会提供一些工具,例如运行训练的命令 tao run
- 提供短小精悍的模块而难以理解的高级抽象
道只提供简单的函数和工具,项目的主要逻辑仍然由用户掌握。道并不是一个几行配置或几行代码就能让程序运行的工具,这类「把问题变简单」的框架有很多,比如 mmcv 或 pytorch-lightning,但实际上工具并不能降低问题的复杂性,越简单通用的工具,灵活性则越差,这些框架一般也提供插件机制来解决灵活性的问题,但开发插件所带来的额外的学习成本也不容小觑。道也依赖框架,我们认为 pytorch-ignite 是个不错的选择,他几乎就是一个「插件系统」。
- 分布式优先
在做任何模块的开发,都将优先考虑 distributed 的可能性,而 tao run
内部本身也使用 torchrun
来运行代码。
- 代码/实验管理
不同于一般的软件,机器学习的训练一个很大的特点是需要不断修改代码或超参数来重复运行。如何记录每一次运行的完整输入(参数)和代码是实验是否能复现的关键(当然 seed 也很重要)。用户需要自己确认哪些是参数,然后通过命令行传入,参数通常还可以通过 tao tune
来做自动搜索。
使用 tao run
来进行单次训练,使用 tao tune
来自动化调参。
tao run 的执行流程
- 将某个提交(默认为当前提交)的代码复制(clone)到
$run_dir / $hash
下,并将工作目录更改为该目录。 - 设置以
TAO
开头的环境变量。 - 调用
torchrun
来运行训练脚本。 - 若不指定
--name
参数,则大部分日志或文件,将产生在$log_dir / $datetime_$hash
目录中,而日志(通过logging
模块或warnings
模块的输出)将全部放置在$log_dir / log.txt
中。若指定--name
,则大部分日志或文件会放置在$log_dir / $name
下。
上文中的 $run_dir
表示配置文件中的 run_dir,$log_dir
表示配置文件中的 log_dir,$datetime
表示运行时间,$hash
表示某个 git commit 的哈希值。
tao tune 的执行流程
- 将某个提交(默认为当前提交)的代码复制(clone)到
$run_dir / $hash
下,并将工作目录更改为该目录。 - 设置以
TAO
开头的环境变量,特别的,会设置TAO_TUNE
的变量。 - 调用
torchrun
来运行训练脚本,并重复运行多次(每次运行会获得不同的参数组合)。 - 若不指定
--name
参数,则大部分日志或文件,将产生在$log_dir / tune_$datetime_$hash / $trial_id
目录中,而日志(通过logging
模块或warnings
模块的输出)将全部放置在$log_dir / log.txt
中。若指定--name
,则大部分日志或文件会放置在$log_dir / $name / $trial_id
下。
其中 $trial_id
为 optuna
中的 trial_id。