294 lines
11 KiB
Python
294 lines
11 KiB
Python
#!/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() |