如何在PyTorch中将模型配置日志保存为configs.yml文件

前置依赖 #

首先安装yaml处理库PyYAML：
pip install pyyaml

实现步骤 #

• 第一步：将所有模型配置、训练超参数、数据集配置等内容整理为Python字典格式，和yaml的层级键值对结构完全匹配
• 第二步：使用PyYAML的dump方法将字典序列化后写入.yml后缀的文件即可

代码示例 #

import yaml
# 1. 定义配置字典，可根据自己的需求增减字段
config = {
    "model_config": {
        "backbone": "ResNet50",
        "num_classes": 8,
        "pretrained": True,
        "dropout": 0.3
    },
    "train_config": {
        "batch_size": 64,
        "epoch": 100,
        "lr": 0.001,
        "optimizer": "AdamW",
        "loss_fn": "CrossEntropyLoss"
    },
    "data_config": {
        "dataset_path": "./dataset",
        "img_size": 224,
        "augment": True
    }
}

# 2. 写入config.yml文件
with open("config.yml", "w", encoding="utf-8") as f:
    yaml.dump(
        config,
        f,
        sort_keys=False, # 保留配置的定义顺序，不按字母自动排序
        default_flow_style=False, # 输出带层级缩进的块格式，符合常规配置文件阅读习惯
        allow_unicode=True # 避免中文内容乱码
    )

可选拓展 #

如果你需要把已实例化的模型参数也写入配置，可在配置字典中加入模型参数字段：

from your_model_file import MyModel

# 实例化模型
model = MyModel(num_classes=8, dropout=0.3)
# 追加模型初始化参数到配置，过滤不可序列化的对象
config["model_config"]["init_params"] = {k:v for k,v in model.__dict__.items() if isinstance(v, (int, float, str, bool, list, dict))}
# 重新写入文件即可
with open("config.yml", "w", encoding="utf-8") as f:
    yaml.dump(config, f, sort_keys=False, default_flow_style=False, allow_unicode=True)

注意事项 #

• 配置字典中不要直接传入PyTorch张量、函数对象、类实例这类不能被yaml直接序列化的内容，需要提前转为字符串、数值等基础Python类型，否则会触发序列化报错
• 后续需要读取配置时，执行config = yaml.safe_load(open("config.yml", "r", encoding="utf-8"))即可将配置读回字典格式，直接用于初始化模型、训练流程

内容的提问来源于stack exchange，提问作者Muhammad Salman