Press "Enter" to skip to content

使用TOML配置轻松管理深度学习模型

您可能永远不需要为您的train.py使用那些冗长的CLI参数

由Scott Graham在Unsplash上的照片

由于所有模块都需要大量的参数和设置,因此管理深度学习模型可能很困难。训练模块可能需要像batch_sizenum_epochs这样的参数,或者是学习率调度程序的参数。同样,数据预处理模块可能需要train_test_split或图像增强的参数。

管理或将这些参数引入管道的一个天真的方法是在运行脚本时将它们用作CLI参数。命令行参数可能很难输入,并且在单个文件中管理所有参数可能不可能。TOML文件提供了一种更清洁的方式来管理配置,脚本可以以Python dict的形式加载配置的必要部分,而无需编写读取/解析命令行参数的样板代码。

在本文中,我们将探讨在配置文件中使用TOML以及如何在训练/部署脚本中高效地使用它们。

什么是TOML文件?

TOML是Tom’s Obvious Minimal Language的缩写,它是专门为配置文件设计的文件格式。 TOML文件的概念与能够以树形层次结构存储键值对的YAML / YML文件非常相似。 TOML比YAML的优点是其可读性,这在存在多个嵌套级别时变得重要。

图1.用TOML(左)和YAML(右)编写的相同模型配置。 TOML允许我们在相同缩进级别上编写键值对,而不考虑层次结构。

个人而言,除了提高可读性之外,我找不到更喜欢TOML而不是YAML的实际理由。使用YAML是完全没有问题的,这里有一个用于解析YAML的Python包。

为什么我们需要TOML中的配置?

使用TOML存储ML模型的模型/数据/部署配置有两个优点:

在单个文件中管理所有配置:使用TOML文件,我们可以创建多个设置组,这些设置组是不同模块所需的。例如,在图1中,与模型训练过程相关的设置嵌套在[train]属性下,类似地,用于部署模型所需的porthost存储在deploy下。我们不需要在train.pydeploy.py之间跳转以更改其参数,而是可以从单个TOML配置文件全局化所有设置。

如果我们在虚拟机上训练模型,而代码编辑器或IDE不可用于编辑文件,则这可能非常有帮助。单个配置文件易于使用大多数VM上可用的vimnano进行编辑。

我们如何从TOML中读取配置?

要从TOML文件中读取配置,可以使用两个Python包tomlmunchtoml将帮助我们读取TOML文件,并将文件的内容返回为Python dictmunch将转换dict的内容,以启用元素的属性样式访问。例如,而不是编写config [ "training"] [ "num_epochs"] ,我们可以只编写config.training.num_epochs,这提高了可读性。

考虑以下文件结构,

- config.py- train.py- project_config.toml

project_config.toml 包含了我们的 ML 项目的配置信息,例如,

[data]vocab_size = 5589seq_length = 10test_split = 0.3data_path = "dataset/"data_tensors_path = "data_tensors/"[model]embedding_dim = 256num_blocks = 5num_heads_in_block = 3[train]num_epochs = 10batch_size = 32learning_rate = 0.001checkpoint_path = "auto"

config.py 文件中,我们创建了一个函数,它使用 tomlmunch 返回了这个配置的 munchified 版本,

$> pip install toml munch

import tomlimport munchdef load_global_config( filepath : str = "project_config.toml" ):    return munch.munchify( toml.load( filepath ) )def save_global_config( new_config , filepath : str = "project_config.toml" ):    with open( filepath , "w" ) as file:        toml.dump( new_config , file )

现在,在我们的任何项目文件中,例如 train.pypredict.py,我们都可以加载这个配置,

from config import load_global_configconfig = load_global_config()batch_size = config.train.batch_sizelr = config.train.learning_rateif config.train.checkpoint_path == "auto":    # Make a directory with name as current timestamp    pass

print( toml.load( filepath ) ) 的输出结果是,

{'data': {'data_path': 'dataset/',          'data_tensors_path': 'data_tensors/',          'seq_length': 10,          'test_split': 0.3,          'vocab_size': 5589}, 'model': {'embedding_dim': 256, 'num_blocks': 5, 'num_heads_in_block': 3}, 'train': {'batch_size': 32,           'checkpoint_path': 'auto',           'learning_rate': 0.001,           'num_epochs': 10}}

如果您正在使用 W&B Tracking 或 MLFlow 等 MLOps 工具,将配置文件作为 dict 进行维护可能会很有帮助,因为我们可以直接将其作为参数传递。

结束

希望您在下一个 ML 项目中考虑使用 TOML 配置!这是一种管理全局或本地训练 / 部署或推理脚本设置的干净方式。

与编写冗长的 CLI 参数不同,脚本可以直接从 TOML 文件加载配置。如果我们希望使用不同的超参数训练两个版本的模型,我们只需要更改 config.py 中的 TOML 文件。我最近在我的项目中开始使用 TOML 文件,实验变得更快了。MLOps 工具还可以管理模型的版本以及它们的配置,但上述方法的简单性是独特的,并且需要最少的更改现有项目。

希望您喜欢阅读。祝您有一个愉快的一天!

Leave a Reply

Your email address will not be published. Required fields are marked *