check_model.py

"""
检查已保存模型的信息
"""
import torch
import os
import yaml

def check_model_info(model_path):
    if not os.path.exists(model_path):
        print(f"Model file not found: {model_path}")
        return
    print(f"Loading model info from: {model_path}")
    checkpoint = torch.load(model_path, map_location='cpu')
    print("\n" + "="*50)
    print("模型检查点信息:")
    print("="*50)
    if 'epoch' in checkpoint:
        print(f"训练轮次: {checkpoint['epoch'] + 1}")
    if 'train_loss' in checkpoint:
        print(f"训练损失: {checkpoint['train_loss']:.4f}")
    if 'val_loss' in checkpoint:
        print(f"验证损失: {checkpoint['val_loss']:.4f}")
    if 'config' in checkpoint:
        print("\n配置信息:")
        config = checkpoint['config']
        for key, value in config.items():
            print(f"  {key}: {value}")
    print(f"\n模型参数数量:")
    if 'model_state_dict' in checkpoint:
        model_params = checkpoint['model_state_dict']
        total_params = sum(p.numel() for p in model_params.values())
        print(f"  总参数: {total_params:,}")
        print(f"\n模型层信息:")
        layer_count = {}
        for name in model_params.keys():
            layer_type = name.split('.')[0] if '.' in name else name
            layer_count[layer_type] = layer_count.get(layer_type, 0) + 1
        
        for layer, count in layer_count.items():
            print(f"  {layer}: {count} 层")
    print("\n" + "="*50)

if __name__ == "__main__":
    model_path = "saves/best_model.pt"
    check_model_info(model_path)