import torch import os import argparse from collections import defaultdict import time def load_checkpoint(ckpt_path): """加载检查点文件""" if not os.path.exists(ckpt_path): return None try: state_dict = torch.load(ckpt_path, map_location='cpu') return state_dict except Exception as e: print(f"❌ 加载检查点失败: {e}") return None def compare_parameters(state_dict1, state_dict2, threshold=1e-8): """比较两个状态字典的参数差异""" if state_dict1 is None or state_dict2 is None: return None updated_params = {} unchanged_params = {} for name, param1 in state_dict1.items(): if name in state_dict2: param2 = state_dict2[name] # 计算参数差异 diff = torch.abs(param1 - param2) max_diff = torch.max(diff).item() mean_diff = torch.mean(diff).item() if max_diff > threshold: updated_params[name] = { 'max_diff': max_diff, 'mean_diff': mean_diff, 'shape': param1.shape } else: unchanged_params[name] = { 'max_diff': max_diff, 'mean_diff': mean_diff, 'shape': param1.shape } return updated_params, unchanged_params def categorize_parameters(param_dict): """将参数按类型分类""" categories = { 'moe_related': {}, 'camera_related': {}, 'framepack_related': {}, 'attention': {}, 'other': {} } for name, info in param_dict.items(): if any(keyword in name.lower() for keyword in ['moe', 'gate', 'expert', 'processor']): categories['moe_related'][name] = info elif any(keyword in name.lower() for keyword in ['cam_encoder', 'projector', 'camera']): categories['camera_related'][name] = info elif any(keyword in name.lower() for keyword in ['clean_x_embedder', 'framepack']): categories['framepack_related'][name] = info elif any(keyword in name.lower() for keyword in ['attn', 'attention']): categories['attention'][name] = info else: categories['other'][name] = info return categories def print_category_summary(category_name, params, color_code=''): """打印某类参数的摘要""" if not params: print(f"{color_code} {category_name}: 无参数") return total_params = len(params) max_diffs = [info['max_diff'] for info in params.values()] mean_diffs = [info['mean_diff'] for info in params.values()] print(f"{color_code} {category_name} ({total_params} 个参数):") print(f" 最大差异范围: {min(max_diffs):.2e} ~ {max(max_diffs):.2e}") print(f" 平均差异范围: {min(mean_diffs):.2e} ~ {max(mean_diffs):.2e}") # 显示前5个最大变化的参数 sorted_params = sorted(params.items(), key=lambda x: x[1]['max_diff'], reverse=True) print(f" 变化最大的参数:") for i, (name, info) in enumerate(sorted_params[:100]): shape_str = 'x'.join(map(str, info['shape'])) print(f" {i+1}. {name} [{shape_str}]: max_diff={info['max_diff']:.2e}") def monitor_training(checkpoint_dir, check_interval=60): """监控训练过程中的参数更新""" print(f"🔍 开始监控训练进度...") print(f"📁 检查点目录: {checkpoint_dir}") print(f"⏰ 检查间隔: {check_interval}秒") print("=" * 80) previous_ckpt = None previous_step = -1 while True: try: # 查找最新的检查点 if not os.path.exists(checkpoint_dir): print(f"❌ 检查点目录不存在: {checkpoint_dir}") time.sleep(check_interval) continue ckpt_files = [f for f in os.listdir(checkpoint_dir) if f.startswith('step') and f.endswith('.ckpt')] if not ckpt_files: print("⏳ 未找到检查点文件，等待中...") time.sleep(check_interval) continue # 按步数排序，获取最新的 ckpt_files.sort(key=lambda x: int(x.replace('step', '').replace('.ckpt', ''))) latest_ckpt_file = ckpt_files[-1] latest_ckpt_path = os.path.join(checkpoint_dir, latest_ckpt_file) # 提取步数 current_step = int(latest_ckpt_file.replace('step', '').replace('.ckpt', '')) if current_step <= previous_step: print(f"⏳ 等待新的检查点... (当前: step{current_step})") time.sleep(check_interval) continue print(f"\n🔍 发现新检查点: {latest_ckpt_file}") # 加载当前检查点 current_state_dict = load_checkpoint(latest_ckpt_path) if current_state_dict is None: print("❌ 无法加载当前检查点") time.sleep(check_interval) continue if previous_ckpt is not None: print(f"📊 比较 step{previous_step} -> step{current_step}") # 比较参数 updated_params, unchanged_params = compare_parameters( previous_ckpt, current_state_dict, threshold=1e-8 ) if updated_params is None: print("❌ 参数比较失败") else: # 分类显示结果 updated_categories = categorize_parameters(updated_params) unchanged_categories = categorize_parameters(unchanged_params) print(f"\n✅ 已更新的参数 (总共 {len(updated_params)} 个):") print_category_summary("MoE相关", updated_categories['moe_related'], '🔥') print_category_summary("Camera相关", updated_categories['camera_related'], '📷') print_category_summary("FramePack相关", updated_categories['framepack_related'], '🎞️') print_category_summary("注意力相关", updated_categories['attention'], '👁️') print_category_summary("其他", updated_categories['other'], '📦') print(f"\n⚠️ 未更新的参数 (总共 {len(unchanged_params)} 个):") print_category_summary("MoE相关", unchanged_categories['moe_related'], '❄️') print_category_summary("Camera相关", unchanged_categories['camera_related'], '❄️') print_category_summary("FramePack相关", unchanged_categories['framepack_related'], '❄️') print_category_summary("注意力相关", unchanged_categories['attention'], '❄️') print_category_summary("其他", unchanged_categories['other'], '❄️') # 检查关键组件是否在更新 critical_keywords = ['moe', 'cam_encoder', 'projector', 'clean_x_embedder'] critical_updated = any( any(keyword in name.lower() for keyword in critical_keywords) for name in updated_params.keys() ) if critical_updated: print("\n✅ 关键组件正在更新！") else: print("\n❌ 警告：关键组件可能未在更新！") # 计算更新率 total_params = len(updated_params) + len(unchanged_params) update_rate = len(updated_params) / total_params * 100 print(f"\n📈 参数更新率: {update_rate:.1f}% ({len(updated_params)}/{total_params})") # 保存当前状态用于下次比较 previous_ckpt = current_state_dict previous_step = current_step print("=" * 80) time.sleep(check_interval) except KeyboardInterrupt: print("\n👋 监控已停止") break except Exception as e: print(f"❌ 监控过程中出错: {e}") time.sleep(check_interval) def compare_two_checkpoints(ckpt1_path, ckpt2_path): """比较两个特定的检查点""" print(f"🔍 比较两个检查点:") print(f" 检查点1: {ckpt1_path}") print(f" 检查点2: {ckpt2_path}") print("=" * 80) # 加载检查点 state_dict1 = load_checkpoint(ckpt1_path) state_dict2 = load_checkpoint(ckpt2_path) if state_dict1 is None or state_dict2 is None: print("❌ 无法加载检查点文件") return # 比较参数 updated_params, unchanged_params = compare_parameters(state_dict1, state_dict2) if updated_params is None: print("❌ 参数比较失败") return # 分类显示结果 updated_categories = categorize_parameters(updated_params) unchanged_categories = categorize_parameters(unchanged_params) print(f"\n✅ 已更新的参数 (总共 {len(updated_params)} 个):") for category_name, params in updated_categories.items(): print_category_summary(category_name.replace('_', ' ').title(), params, '🔥') print(f"\n⚠️ 未更新的参数 (总共 {len(unchanged_params)} 个):") for category_name, params in unchanged_categories.items(): print_category_summary(category_name.replace('_', ' ').title(), params, '❄️') # 计算更新率 total_params = len(updated_params) + len(unchanged_params) update_rate = len(updated_params) / total_params * 100 print(f"\n📈 参数更新率: {update_rate:.1f}% ({len(updated_params)}/{total_params})") if __name__ == '__main__': parser = argparse.ArgumentParser(description="检查模型参数更新情况") parser.add_argument("--checkpoint_dir", type=str, default="/share_zhuyixuan05/zhuyixuan05/ICLR2026/framepack_moe", help="检查点目录路径") parser.add_argument("--compare", default=True, help="比较两个特定检查点，而不是监控") parser.add_argument("--ckpt1", type=str, default="/share_zhuyixuan05/zhuyixuan05/ICLR2026/framepack_moe/step1500_origin_cam_4.ckpt") parser.add_argument("--ckpt2", type=str, default="/share_zhuyixuan05/zhuyixuan05/ICLR2026/framepack_moe/step500_origin_cam_4.ckpt") parser.add_argument("--interval", type=int, default=60, help="监控检查间隔（秒）") parser.add_argument("--threshold", type=float, default=1e-8, help="参数变化阈值") args = parser.parse_args() if args.compare: if not args.ckpt1 or not args.ckpt2: print("❌ 比较模式需要指定 --ckpt1 和 --ckpt2") else: compare_two_checkpoints(args.ckpt1, args.ckpt2) else: monitor_training(args.checkpoint_dir, args.interval)