如何获取现有自定义YOLOv5模型的训练规格信息？

YOLOv5导出的自定义模型文件（.pt）会默认保存训练时的关键元数据，你可以直接通过模型对象的属性提取这些信息，无需额外工具。以下是实现你需求的代码示例：

import torch

# 加载本地自定义YOLOv5模型
model = torch.hub.load('./yolov5', 'custom', path='./custom_model.pt', source='local')

def get_train_specs(model):
    # 提取模型版本
    model_version = model.yaml.get('version', '未知')
    # 提取训练时使用的图像尺寸
    train_img_size = model.yaml.get('imgsz', '未知')
    # 提取训练超参数
    train_hyp = model.hyp

    # 格式化输出信息
    print(f"Model version: {model_version}")
    print(f"Trained on image size: {train_img_size}")
    print("Hyperparameters used during training:")
    for param, value in train_hyp.items():
        print(f"    {param}: {value}")

# 调用函数获取并打印训练规格
get_train_specs(model)

关键信息说明 #

• 模型版本：存储在模型内置的yaml配置中，通过model.yaml['version']直接读取
• 训练图像尺寸：对应yaml配置里的imgsz字段，记录了训练时设置的输入图像分辨率
• 训练超参数：所有训练阶段使用的超参数（如lr0、lrf、momentum等）都会被保存到model.hyp属性中，直接遍历即可获取全部参数

如果你的模型版本较旧，部分属性可能需要从模型的检查点对象中读取，比如model.ckpt.get('hyp', {})，但绝大多数YOLOv5 v5.0+版本都支持上述直接访问的方式。

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