#!/usr/bin/env python3 """ 分阶段TTA-E参数搜索 先粗搜索找到有希望的区域，再精细搜索 """ import os import sys import argparse import json import numpy as np from itertools import product import subprocess import time def parse_arguments(): parser = argparse.ArgumentParser(description='分阶段TTA-E参数搜索') # 基础参数 parser.add_argument('--base_script', type=str, default='evaluate_model.py') parser.add_argument('--data_dir', type=str, default='../data/hdf5_data_final') parser.add_argument('--eval_type', type=str, default='val') parser.add_argument('--gpu_number', type=int, default=0) # 搜索阶段控制 parser.add_argument('--stage', type=str, default='coarse', choices=['coarse', 'fine', 'both'], help='搜索阶段：coarse=粗搜索，fine=精细搜索，both=两阶段') parser.add_argument('--coarse_results', type=str, default='coarse_results.json', help='粗搜索结果文件（用于精细搜索阶段）') parser.add_argument('--final_results', type=str, default='final_results.json', help='最终结果文件') # 粗搜索参数（步长0.2） parser.add_argument('--coarse_gru_weights', type=str, default='0.2,0.4,0.6,0.8,1.0') parser.add_argument('--coarse_tta_weights', type=str, default='0.0,0.5,1.0') # 精细搜索参数（步长0.1，在最佳配置周围） parser.add_argument('--fine_range', type=float, default=0.3, help='精细搜索范围（围绕最佳配置的±范围）') parser.add_argument('--fine_step', type=float, default=0.1, help='精细搜索步长') # 筛选控制 parser.add_argument('--top_k', type=int, default=5, help='选择前K个最佳配置进行精细搜索') return parser.parse_args() def generate_coarse_search_space(args): """生成粗搜索空间""" gru_weights = [float(x.strip()) for x in args.coarse_gru_weights.split(',')] tta_weights = [float(x.strip()) for x in args.coarse_tta_weights.split(',')] search_space = [] for gru_w in gru_weights: for noise_w in tta_weights: for scale_w in tta_weights: for shift_w in tta_weights: for smooth_w in tta_weights: search_space.append((gru_w, 1.0, noise_w, scale_w, shift_w, smooth_w)) return search_space def generate_fine_search_space(best_configs, args): """基于最佳配置生成精细搜索空间""" fine_search_space = [] for config in best_configs: gru_w = config['gru_weight'] tta_w = config['tta_weights'] # 在每个参数周围生成精细搜索点 gru_range = np.arange( max(0.1, gru_w - args.fine_range), min(1.0, gru_w + args.fine_range) + args.fine_step, args.fine_step ) for param_name in ['noise', 'scale', 'shift', 'smooth']: base_val = tta_w[param_name] param_range = np.arange( max(0.0, base_val - args.fine_range), min(1.0, base_val + args.fine_range) + args.fine_step, args.fine_step ) # 围绕当前最佳配置生成邻域 for gru_fine in gru_range: for noise_fine in param_range if param_name == 'noise' else [tta_w['noise']]: for scale_fine in param_range if param_name == 'scale' else [tta_w['scale']]: for shift_fine in param_range if param_name == 'shift' else [tta_w['shift']]: for smooth_fine in param_range if param_name == 'smooth' else [tta_w['smooth']]: config_tuple = ( round(gru_fine, 1), 1.0, round(noise_fine, 1), round(scale_fine, 1), round(shift_fine, 1), round(smooth_fine, 1) ) if config_tuple not in fine_search_space: fine_search_space.append(config_tuple) return fine_search_space def run_evaluation(config, args): """运行单个配置的评估""" gru_w, orig_w, noise_w, scale_w, shift_w, smooth_w = config tta_weights_str = f"{orig_w},{noise_w},{scale_w},{shift_w},{smooth_w}" cmd = [ 'python', args.base_script, '--gru_weight', str(gru_w), '--tta_weights', tta_weights_str, '--data_dir', args.data_dir, '--eval_type', args.eval_type, '--gpu_number', str(args.gpu_number) ] try: result = subprocess.run(cmd, capture_output=True, text=True, timeout=1800) # 30分钟超时 # 解析PER结果 per = None for line in result.stdout.split('\n'): if 'Aggregate Phoneme Error Rate (PER):' in line: per_str = line.split(':')[-1].strip().replace('%', '') per = float(per_str) break if per is None: print(f"⚠️ 无法解析PER结果: {config}") per = float('inf') return { 'config': config, 'gru_weight': gru_w, 'tta_weights': { 'original': orig_w, 'noise': noise_w, 'scale': scale_w, 'shift': shift_w, 'smooth': smooth_w }, 'per': per, 'success': result.returncode == 0, 'stdout': result.stdout[:1000], # 只保存前1000字符 } except subprocess.TimeoutExpired: return {'config': config, 'per': float('inf'), 'error': 'Timeout'} except Exception as e: return {'config': config, 'per': float('inf'), 'error': str(e)} def run_coarse_search(args): """运行粗搜索""" print("🔍 第一阶段：粗搜索") print("=" * 50) search_space = generate_coarse_search_space(args) total_configs = len(search_space) print(f"粗搜索空间: {total_configs} 个配置") print(f"GRU权重: {args.coarse_gru_weights}") print(f"TTA权重: {args.coarse_tta_weights}") print() results = [] best_per = float('inf') for i, config in enumerate(search_space): print(f"进度: {i+1}/{total_configs} ({100*(i+1)/total_configs:.1f}%)") print(f"配置: GRU={config[0]:.1f}, TTA=({config[2]},{config[3]},{config[4]},{config[5]})") result = run_evaluation(config, args) results.append(result) if result['per'] < best_per: best_per = result['per'] print(f"🎯 新最佳PER: {best_per:.3f}%") else: print(f" PER: {result['per']:.3f}%") # 保存粗搜索结果 coarse_results = { 'results': results, 'stage': 'coarse', 'timestamp': time.strftime("%Y-%m-%d %H:%M:%S"), 'args': vars(args) } with open(args.coarse_results, 'w') as f: json.dump(coarse_results, f, indent=2) # 选择最佳配置 valid_results = [r for r in results if r['per'] != float('inf')] best_configs = sorted(valid_results, key=lambda x: x['per'])[:args.top_k] print(f"\n粗搜索完成！选择前{args.top_k}个配置进行精细搜索:") for i, config in enumerate(best_configs): print(f"{i+1}. PER={config['per']:.3f}% | GRU={config['gru_weight']:.1f} | {config['tta_weights']}") return best_configs def run_fine_search(best_configs, args): """运行精细搜索""" print(f"\n🔬 第二阶段：精细搜索") print("=" * 50) fine_search_space = generate_fine_search_space(best_configs, args) total_configs = len(fine_search_space) print(f"精细搜索空间: {total_configs} 个配置") print(f"搜索范围: ±{args.fine_range}") print(f"搜索步长: {args.fine_step}") print() results = [] best_per = float('inf') for i, config in enumerate(fine_search_space): print(f"进度: {i+1}/{total_configs} ({100*(i+1)/total_configs:.1f}%)") result = run_evaluation(config, args) results.append(result) if result['per'] < best_per: best_per = result['per'] print(f"🎯 新最佳PER: {best_per:.3f}%") print(f" 配置: GRU={result['gru_weight']:.1f} | {result['tta_weights']}") if i % 10 == 0: # 每10个配置显示一次进度 print(f" 当前PER: {result['per']:.3f}%") return results def main(): args = parse_arguments() print("🚀 分阶段TTA-E参数搜索") print("=" * 60) if args.stage in ['coarse', 'both']: # 运行粗搜索 best_configs = run_coarse_search(args) if args.stage == 'coarse': print(f"\n✅ 粗搜索完成，结果保存到: {args.coarse_results}") return else: # 从文件加载粗搜索结果 print(f"📁 加载粗搜索结果: {args.coarse_results}") with open(args.coarse_results, 'r') as f: coarse_data = json.load(f) valid_results = [r for r in coarse_data['results'] if r['per'] != float('inf')] best_configs = sorted(valid_results, key=lambda x: x['per'])[:args.top_k] if args.stage in ['fine', 'both']: # 运行精细搜索 fine_results = run_fine_search(best_configs, args) # 合并所有结果 all_results = fine_results if args.stage == 'both': all_results.extend([r for r in coarse_data['results'] if 'results' in locals()]) # 找到最终最佳配置 valid_results = [r for r in all_results if r['per'] != float('inf')] final_best = min(valid_results, key=lambda x: x['per']) # 保存最终结果 final_results = { 'best_config': final_best, 'all_fine_results': fine_results, 'stage': args.stage, 'timestamp': time.strftime("%Y-%m-%d %H:%M:%S"), 'args': vars(args) } with open(args.final_results, 'w') as f: json.dump(final_results, f, indent=2) print(f"\n🏆 最终最佳配置:") print(f"PER: {final_best['per']:.3f}%") print(f"GRU权重: {final_best['gru_weight']:.1f}") print(f"TTA权重: {final_best['tta_weights']}") print(f"结果保存到: {args.final_results}") # 显示top-10 sorted_results = sorted(valid_results, key=lambda x: x['per'])[:10] print(f"\n📊 Top-10配置:") for i, result in enumerate(sorted_results): tw = result['tta_weights'] print(f"{i+1:2d}. PER={result['per']:6.3f}% | GRU={result['gru_weight']:.1f} | " f"TTA=({tw['noise']:.1f},{tw['scale']:.1f},{tw['shift']:.1f},{tw['smooth']:.1f})") if __name__ == "__main__": main()