> ## Content Index
> Fetch the complete content index at: https://yinguobing.com/blog/llms.txt
> Use this file to discover other available public pages before exploring further.

# 在TensorFlow模型训练任务中使用YAML配置文件
- URL: https://yinguobing.com/blog/use-yaml-for-tensorflow-training-job-configuration/
- Published: 2021-06-27T03:03:59.000Z
- Updated: 2026-09-08T02:06:37.000Z
- Description: 训练任务工程配置的首选方式。
- Author: 尹国冰
- Tags: TensorFlow

封面图片： [Mitchell Leach](https://unsplash.com/@mitchellleach?utm%5Fsource=unsplash&utm%5Fmedium=referral&utm%5Fcontent=creditCopyText)

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

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

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

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

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

### YAML

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

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

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

```yaml
name: "resnet"
```

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

```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数量。

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

模板如下：

```yaml
# This is the configuration file for TensorFlow model training job.

# This will be used as the directory name for checkpoint and logging.
name: "resnet50_v2"

# Defines the model architecture parameters.
model:
  input_shape: [112, 112, 3]
  output_size: 1024

# The training hyper parameters.
training:
  train_file_template: "/data/train_{:05d}.tfrecord"
  train_file_number: 24
  batch_size: 128
  epochs: 30
  initial_epoch: 0

# Where to log the training process used by tensorboard later.
# dir: the log directory.
# frequency: how ofen to log and save the checkpoint, in steps.
logging:
  dir: "/home/robin/resnet/logs"
  frequency: 1000

# How to save the model weights to disk during training.
# dir: a string of path to checkpoint the current model.
# restore_dir: points to a previous checkpoint directory used to restore the
#   model weights. Use `!!null` to skip restoring.
checkpoint:
  dir: "/home/robin/resnet/checkpoints"
  restore_dir: !!null

# The multiworker distribute training requires OS enviroument setup.
distribute:
  tf_config:
    cluster:
      worker:
        - "192.168.1.12:50085"
        - "192.168.1.13:50085"
    task: { type: "worker", index: 0 }

# Setup the GPUs
gpu:
  num_visible_devices: 2

```

TensorFlow训练任务配置模板

将以上内容复制并保存为YAML文件即可。你还可以视情况使用git来跟踪这个文件，方便将来回溯训练过程。