【YOLO脚本】对模型yaml文件测试
输出模型和模型每一层的参数,计算量的脚本main_profile.py
from ultralytics import YOLO# pip install timm efficientnet_pytorch dill einops
# conda install -c conda-forge pywaveletsif __name__ == '__main__':# choose your yaml filemodel = YOLO('ultralytics/cfg/models/11/yolo11n.yaml')model.info(detailed=True)try:model.profile(imgsz=[640, 640])except Exception as e:print(e)passmodel.fuse()
检查一个目录的yaml文件test_yaml.py
测试覆盖的场景:
- ✅ 配置文件语法正确,无明显结构错误。
- ✅ 模型可加载,前向传播无报错。
- ✅ 模型优化(如层融合)可正常执行。
该测试主要验证模型配置的语法正确性和基本运行能力,是模型训练前的必要预检查,但不能替代完整的训练验证。建议在测试通过后,仍进行小规模训练(如 1-2 个 epoch)以验证训练流程的完整性
import warnings
warnings.filterwarnings('ignore')
import os, tqdm
from ultralytics import YOLOif __name__ == '__main__':error_result = []for yaml_path in tqdm.tqdm(os.listdir('ultralytics/cfg/models/11')):if 'rtdetr' not in yaml_path and 'cls' not in yaml_path and 'world' not in yaml_path:try:model = YOLO(f'ultralytics/cfg/models/11/{yaml_path}')model.info(detailed=True)model.profile([640, 640])model.fuse()except Exception as e:error_result.append(f'{yaml_path} {e}')for i in error_result:print(i)
测试单个yaml模型文件
用变量存储yaml模型文件地址。
import warnings
warnings.filterwarnings('ignore')
from ultralytics import YOLOdef test_single_model(yaml_path, detailed_info=True, run_profile=True):"""测试单个 YOLO 模型配置文件参数:yaml_path: 模型配置文件的路径detailed_info: 是否显示详细的模型信息run_profile: 是否运行性能分析"""print(f"开始测试模型: {yaml_path}")print("=" * 60)try:# 加载模型model = YOLO(yaml_path)# 显示模型信息print(f"模型信息: {model.model.__class__.__name__}")model.info(detailed=detailed_info)# 获取模型参数量params = sum(p.numel() for p in model.model.parameters() if p.requires_grad)print(f"\n总训练参数: {params:,}")# 运行性能分析 (可选)if run_profile:print("\n正在分析模型性能...")try:model.profile(imgsz=[640, 640])except Exception as profile_error:print(f"性能分析失败: {profile_error}")print("提示: 某些模型架构可能不支持性能分析")# 模型融合print("\n正在优化模型结构...")model.fuse()print("模型融合成功!")return True, "模型测试通过"except Exception as e:print(f"\n错误: {e}")return False, str(e)if __name__ == '__main__':# 通过变量指定 YAML 文件路径MODEL_YAML_PATH = 'ultralytics/cfg/models/11/yolo11n.yaml'# 配置选项SHOW_DETAILED_INFO = True # 是否显示详细的模型信息RUN_PROFILE_ANALYSIS = True # 是否运行性能分析# 测试模型success, message = test_single_model(MODEL_YAML_PATH,detailed_info=SHOW_DETAILED_INFO,run_profile=RUN_PROFILE_ANALYSIS)# 输出测试结果print("\n" + "=" * 60)if success:print(f"✓ 模型配置验证通过: {MODEL_YAML_PATH}")else:print(f"✗ 模型配置验证失败: {MODEL_YAML_PATH}")print(f"错误详情: {message}")