Sign in Subscribe
TensorFlow

在TensorFlow模型训练任务中使用YAML配置文件

训练任务工程配置的首选方式。

在TensorFlow模型训练任务中使用YAML配置文件

封面图片： Mitchell Leach

深度学习模型的训练任务涉及到大量的配置参数。除了常见的超参数以外还有一些辅助任务参数，例如checkpoint与log的写入目录、GPU数量、分布式训练角色等。在实际训练过程中这些参数会按需调整，为了减少工作量，同时避免无意中修改了源代码，我通常采用命令参数的方式提供修改方法。例如在arcface项目中，可以通过如下命令设定batch size与epoch：

python3 train.py --epochs=2 --batch_size=192 --softmax=True

这种设定训练参数的方式使用起来简单方便，但是也有不足之处。

  1. 任务配置丢失。一旦任务由于各种原因（例如掉电）中断，有可能面临训练参数不可知的局面。
  2. 不利于任务版本控制。对于有必要跟踪训练进展的状况，需要记录每个任务启动时的命令。不方便不优雅。
  3. 配置繁琐。对于某些任务来说，涉及到的配置参数数量庞大，手动修改多有不便。

不仅如此，面对一个正在运行中的训练任务，若是想知晓它的训练参数就必须在终端输出中查找它的启动命令，这需要翻阅大量的日志输出。这种情况下，使用任务配置文件是一个更好的方案。

YAML

如果你在网络上搜索配置文件的话可以找到的方式多如牛毛。我最终选择了YAML的方案是因为

  • 配置文件格式清晰明了，使用缩进来格式化，没有大堆的括号干扰阅读。
  • 可以被Python等多种常用语言解析。
  • 语法很像Python。

例如新建一个cfg.yml文件，填入下面这一行配置：

name: "resnet"

然后使用Python读取它并解析：

import yaml
with open("cfg.yml", r) as f:
    cfg = yaml.load(f, Loader=yaml.FullLoader)
print(cfg)
---------------
{'name': 'resnet'}

配置文件在经过Python解析之后返回一个字典，可以很方便的在代码中读取调用。

配置文件模板

真实的训练任务中配置参数数量庞大且类别多样。为了方便使用，我在这里给出一份配置文件模板。该模板按照人物属性对配置参数做了分组。一级分组包括：

  • name 模型名称，可以用来区分不同的训练任务模型。
  • model 模型相关参数，例如输入的形状，输出形状等等。
  • training 训练参数，主要包括batch size, epochs等超参。
  • logging 与训练日志相关，例如存储位置，频率等。
  • checkpoint 模型文件检查点，通常用来存储训练过程中的权重参数。
  • distribute 分布式训练相关参数。
  • gpu 使用的GPU数量。

这只是一个模板，尽可能的覆盖常用的配置选项。实际使用时多半要根据自己的情况修改。

模板如下：

Read next