#!/usr/bin/env python3 """ TTA-E参数搜索脚本 先运行所有基础配置获取预测结果，然后搜索最优的参数组合 避免重复模型推理，提高搜索效率 """ import os import sys import torch import numpy as np import pandas as pd from omegaconf import OmegaConf import time from tqdm import tqdm import editdistance import argparse import itertools import json from concurrent.futures import ProcessPoolExecutor import pickle # Add parent directories to path to import models sys.path.append(os.path.join(os.path.dirname(__file__), '..', 'model_training')) sys.path.append(os.path.join(os.path.dirname(__file__), '..', 'model_training_lstm')) from model_training.rnn_model import GRUDecoder from model_training_lstm.rnn_model import LSTMDecoder from model_training.evaluate_model_helpers import * def parse_arguments(): parser = argparse.ArgumentParser(description='TTA-E Parameter Search for optimal configuration.') parser.add_argument('--gru_model_path', type=str, default='/root/autodl-tmp/nejm-brain-to-text/data/t15_pretrained_rnn_baseline', help='Path to the pretrained GRU model directory.') parser.add_argument('--lstm_model_path', type=str, default='/root/autodl-tmp/nejm-brain-to-text/model_training_lstm/trained_models/baseline_rnn', help='Path to the pretrained LSTM model directory.') parser.add_argument('--data_dir', type=str, default='../data/hdf5_data_final', help='Path to the dataset directory.') parser.add_argument('--csv_path', type=str, default='../data/t15_copyTaskData_description.csv', help='Path to the CSV file with metadata.') parser.add_argument('--gpu_number', type=int, default=0, help='GPU number to use for model inference.') parser.add_argument('--eval_type', type=str, default='val', choices=['val', 'test'], help='Evaluation type.') # 搜索空间参数 parser.add_argument('--gru_weights', type=str, default='0.4,0.5,0.6,0.7,0.8,1.0', help='Comma-separated GRU weights to search.') parser.add_argument('--tta_noise_weights', type=str, default='0.0,0.5,1.0', help='Comma-separated noise weights to search.') parser.add_argument('--tta_scale_weights', type=str, default='0.0,0.5,1.0', help='Comma-separated scale weights to search.') parser.add_argument('--tta_shift_weights', type=str, default='0.0,0.5,1.0', help='Comma-separated shift weights to search.') parser.add_argument('--tta_smooth_weights', type=str, default='0.0,0.5,1.0', help='Comma-separated smooth weights to search.') # TTA固定参数 parser.add_argument('--tta_noise_std', type=float, default=0.01, help='Standard deviation for TTA noise augmentation.') parser.add_argument('--tta_smooth_range', type=float, default=0.5, help='Range for TTA smoothing kernel variation.') parser.add_argument('--tta_scale_range', type=float, default=0.05, help='Range for TTA amplitude scaling.') parser.add_argument('--tta_cut_max', type=int, default=3, help='Maximum timesteps for TTA shift.') # 输出控制 parser.add_argument('--cache_file', type=str, default='tta_predictions_cache.pkl', help='File to cache model predictions.') parser.add_argument('--results_file', type=str, default='parameter_search_results.json', help='File to save search results.') parser.add_argument('--force_recache', action='store_true', help='Force re-computation of predictions cache.') return parser.parse_args() def generate_all_base_configs(): """生成所有需要运行的基础配置（每种增强单独运行）""" base_configs = [ {'name': 'original', 'tta_weights': {'original': 1.0, 'noise': 0.0, 'scale': 0.0, 'shift': 0.0, 'smooth': 0.0}}, {'name': 'noise', 'tta_weights': {'original': 0.0, 'noise': 1.0, 'scale': 0.0, 'shift': 0.0, 'smooth': 0.0}}, {'name': 'scale', 'tta_weights': {'original': 0.0, 'noise': 0.0, 'scale': 1.0, 'shift': 0.0, 'smooth': 0.0}}, {'name': 'shift', 'tta_weights': {'original': 0.0, 'noise': 0.0, 'scale': 0.0, 'shift': 1.0, 'smooth': 0.0}}, {'name': 'smooth', 'tta_weights': {'original': 0.0, 'noise': 0.0, 'scale': 0.0, 'shift': 0.0, 'smooth': 1.0}}, ] return base_configs def run_single_tta_prediction(x, input_layer, gru_model, lstm_model, gru_model_args, lstm_model_args, device, aug_type, tta_noise_std, tta_smooth_range, tta_scale_range, tta_cut_max): """运行单个TTA增强的预测""" x_augmented = x.clone() # Get default smoothing parameters default_smooth_std = gru_model_args['dataset']['data_transforms']['smooth_kernel_std'] default_smooth_size = gru_model_args['dataset']['data_transforms']['smooth_kernel_size'] if aug_type == 'original': pass elif aug_type == 'noise': noise_scale = tta_noise_std * (0.5 + 0.5 * np.random.rand()) noise = torch.randn_like(x_augmented) * noise_scale x_augmented = x_augmented + noise elif aug_type == 'scale': scale_factor = 1.0 + (torch.rand(1).item() - 0.5) * 2 * tta_scale_range x_augmented = x_augmented * scale_factor elif aug_type == 'shift' and tta_cut_max > 0: shift_amount = np.random.randint(1, min(tta_cut_max + 1, x_augmented.shape[1] // 8)) x_augmented = torch.cat([x_augmented[:, shift_amount:, :], x_augmented[:, :shift_amount, :]], dim=1) elif aug_type == 'smooth': smooth_variation = (torch.rand(1).item() - 0.5) * 2 * tta_smooth_range varied_smooth_std = max(0.3, default_smooth_std + smooth_variation) # Use autocast for efficiency with torch.autocast(device_type="cuda", enabled=gru_model_args['use_amp'], dtype=torch.bfloat16): # Apply Gaussian smoothing if aug_type == 'smooth': x_smoothed = gauss_smooth( inputs=x_augmented, device=device, smooth_kernel_std=varied_smooth_std, smooth_kernel_size=default_smooth_size, padding='valid', ) else: x_smoothed = gauss_smooth( inputs=x_augmented, device=device, smooth_kernel_std=default_smooth_std, smooth_kernel_size=default_smooth_size, padding='valid', ) with torch.no_grad(): # Get GRU logits and convert to probabilities gru_logits, _ = gru_model( x=x_smoothed, day_idx=torch.tensor([input_layer], device=device), states=None, return_state=True, ) gru_probs = torch.softmax(gru_logits, dim=-1) # Get LSTM logits and convert to probabilities lstm_logits, _ = lstm_model( x=x_smoothed, day_idx=torch.tensor([input_layer], device=device), states=None, return_state=True, ) lstm_probs = torch.softmax(lstm_logits, dim=-1) return gru_probs.float().cpu().numpy(), lstm_probs.float().cpu().numpy() def cache_model_predictions(args): """缓存所有基础模型预测结果""" print("=== 第一阶段：缓存模型预测结果 ===") # 设置设备 if torch.cuda.is_available() and args.gpu_number >= 0: device = torch.device(f'cuda:{args.gpu_number}') else: device = torch.device('cpu') print(f'Using device: {device}') # 加载模型 print("Loading models...") gru_model_args = OmegaConf.load(os.path.join(args.gru_model_path, 'checkpoint/args.yaml')) lstm_model_args = OmegaConf.load(os.path.join(args.lstm_model_path, 'checkpoint/args.yaml')) # Define GRU model gru_model = GRUDecoder( neural_dim=gru_model_args['model']['n_input_features'], n_units=gru_model_args['model']['n_units'], n_days=len(gru_model_args['dataset']['sessions']), n_classes=gru_model_args['dataset']['n_classes'], rnn_dropout=gru_model_args['model']['rnn_dropout'], input_dropout=gru_model_args['model']['input_network']['input_layer_dropout'], n_layers=gru_model_args['model']['n_layers'], patch_size=gru_model_args['model']['patch_size'], patch_stride=gru_model_args['model']['patch_stride'], ) # Load GRU model weights gru_checkpoint = torch.load(os.path.join(args.gru_model_path, 'checkpoint/best_checkpoint'), weights_only=False, map_location=device) for key in list(gru_checkpoint['model_state_dict'].keys()): gru_checkpoint['model_state_dict'][key.replace("module.", "")] = gru_checkpoint['model_state_dict'].pop(key) gru_checkpoint['model_state_dict'][key.replace("_orig_mod.", "")] = gru_checkpoint['model_state_dict'].pop(key) gru_model.load_state_dict(gru_checkpoint['model_state_dict']) # Define LSTM model lstm_model = LSTMDecoder( neural_dim=lstm_model_args['model']['n_input_features'], n_units=lstm_model_args['model']['n_units'], n_days=len(lstm_model_args['dataset']['sessions']), n_classes=lstm_model_args['dataset']['n_classes'], rnn_dropout=lstm_model_args['model']['rnn_dropout'], input_dropout=lstm_model_args['model']['input_network']['input_layer_dropout'], n_layers=lstm_model_args['model']['n_layers'], patch_size=lstm_model_args['model']['patch_size'], patch_stride=lstm_model_args['model']['patch_stride'], ) # Load LSTM model weights lstm_checkpoint = torch.load(os.path.join(args.lstm_model_path, 'checkpoint/best_checkpoint'), weights_only=False, map_location=device) for key in list(lstm_checkpoint['model_state_dict'].keys()): lstm_checkpoint['model_state_dict'][key.replace("module.", "")] = lstm_checkpoint['model_state_dict'].pop(key) lstm_checkpoint['model_state_dict'][key.replace("_orig_mod.", "")] = lstm_checkpoint['model_state_dict'].pop(key) lstm_model.load_state_dict(lstm_checkpoint['model_state_dict']) gru_model.to(device) lstm_model.to(device) gru_model.eval() lstm_model.eval() # 加载数据 print("Loading data...") b2txt_csv_df = pd.read_csv(args.csv_path) test_data = {} total_trials = 0 for session in gru_model_args['dataset']['sessions']: files = [f for f in os.listdir(os.path.join(args.data_dir, session)) if f.endswith('.hdf5')] if f'data_{args.eval_type}.hdf5' in files: eval_file = os.path.join(args.data_dir, session, f'data_{args.eval_type}.hdf5') data = load_h5py_file(eval_file, b2txt_csv_df) test_data[session] = data total_trials += len(test_data[session]["neural_features"]) print(f'Loaded {len(test_data[session]["neural_features"])} {args.eval_type} trials for session {session}.') print(f'Total trials: {total_trials}') # 生成所有基础预测 base_configs = generate_all_base_configs() predictions_cache = {} for config in base_configs: config_name = config['name'] print(f"\nGenerating predictions for: {config_name}") predictions_cache[config_name] = {} with tqdm(total=total_trials, desc=f'Processing {config_name}', unit='trial') as pbar: for session, data in test_data.items(): input_layer = gru_model_args['dataset']['sessions'].index(session) session_predictions = [] for trial in range(len(data['neural_features'])): neural_input = data['neural_features'][trial] neural_input = np.expand_dims(neural_input, axis=0) neural_input = torch.tensor(neural_input, device=device, dtype=torch.bfloat16) gru_probs, lstm_probs = run_single_tta_prediction( neural_input, input_layer, gru_model, lstm_model, gru_model_args, lstm_model_args, device, config_name, args.tta_noise_std, args.tta_smooth_range, args.tta_scale_range, args.tta_cut_max ) session_predictions.append({ 'gru_probs': gru_probs, 'lstm_probs': lstm_probs, 'trial_info': { 'session': session, 'block_num': data['block_num'][trial], 'trial_num': data['trial_num'][trial], 'seq_class_ids': data['seq_class_ids'][trial] if args.eval_type == 'val' else None, 'seq_len': data['seq_len'][trial] if args.eval_type == 'val' else None, 'sentence_label': data['sentence_label'][trial] if args.eval_type == 'val' else None, } }) pbar.update(1) predictions_cache[config_name][session] = session_predictions # 保存缓存 print(f"\nSaving predictions cache to {args.cache_file}...") with open(args.cache_file, 'wb') as f: pickle.dump(predictions_cache, f) print("✓ Predictions cache saved successfully!") return predictions_cache def ensemble_and_evaluate(predictions_cache, gru_weight, tta_weights, eval_type='val'): """基于缓存的预测结果进行集成和评估""" lstm_weight = 1.0 - gru_weight epsilon = 1e-8 total_trials = 0 total_edit_distance = 0 total_true_length = 0 # 检查哪些增强被启用 enabled_augmentations = [aug_type for aug_type, weight in tta_weights.items() if weight > 0] if len(enabled_augmentations) == 0: return float('inf') # 无效配置 # 归一化TTA权重 total_tta_weight = sum(weight for weight in tta_weights.values() if weight > 0) normalized_tta_weights = {k: v/total_tta_weight for k, v in tta_weights.items() if v > 0} for session in predictions_cache['original'].keys(): session_predictions = predictions_cache['original'][session] for trial_idx in range(len(session_predictions)): trial_info = session_predictions[trial_idx]['trial_info'] if eval_type == 'val' and trial_info['seq_class_ids'] is None: continue # 收集所有启用增强的概率 weighted_gru_probs = None weighted_lstm_probs = None for aug_type in enabled_augmentations: weight = normalized_tta_weights[aug_type] gru_probs = predictions_cache[aug_type][session][trial_idx]['gru_probs'] lstm_probs = predictions_cache[aug_type][session][trial_idx]['lstm_probs'] gru_probs = torch.tensor(gru_probs) lstm_probs = torch.tensor(lstm_probs) if weighted_gru_probs is None: weighted_gru_probs = weight * gru_probs weighted_lstm_probs = weight * lstm_probs else: # 处理序列长度不同的情况 min_len = min(weighted_gru_probs.shape[1], gru_probs.shape[1]) weighted_gru_probs = weighted_gru_probs[:, :min_len, :] + weight * gru_probs[:, :min_len, :] weighted_lstm_probs = weighted_lstm_probs[:, :min_len, :] + weight * lstm_probs[:, :min_len, :] # 集成GRU和LSTM（几何平均） weighted_gru_probs = weighted_gru_probs + epsilon weighted_lstm_probs = weighted_lstm_probs + epsilon log_ensemble_probs = (gru_weight * torch.log(weighted_gru_probs) + lstm_weight * torch.log(weighted_lstm_probs)) ensemble_probs = torch.exp(log_ensemble_probs) ensemble_probs = ensemble_probs / ensemble_probs.sum(dim=-1, keepdim=True) # 解码预测序列 pred_seq = torch.argmax(ensemble_probs[0], dim=-1).numpy() pred_seq = [int(p) for p in pred_seq if p != 0] pred_seq = [pred_seq[i] for i in range(len(pred_seq)) if i == 0 or pred_seq[i] != pred_seq[i-1]] pred_phonemes = [LOGIT_TO_PHONEME[p] for p in pred_seq] if eval_type == 'val': # 计算PER true_seq = trial_info['seq_class_ids'][0:trial_info['seq_len']] true_phonemes = [LOGIT_TO_PHONEME[p] for p in true_seq] ed = editdistance.eval(true_phonemes, pred_phonemes) total_edit_distance += ed total_true_length += len(true_phonemes) total_trials += 1 if eval_type == 'val' and total_true_length > 0: per = 100 * total_edit_distance / total_true_length return per else: return 0.0 # test模式返回0 def search_optimal_parameters(predictions_cache, args): """搜索最优参数组合""" print("\n=== 第二阶段：参数搜索 ===") # 解析搜索空间 gru_weights = [float(x) for x in args.gru_weights.split(',')] noise_weights = [float(x) for x in args.tta_noise_weights.split(',')] scale_weights = [float(x) for x in args.tta_scale_weights.split(',')] shift_weights = [float(x) for x in args.tta_shift_weights.split(',')] smooth_weights = [float(x) for x in args.tta_smooth_weights.split(',')] print(f"Search space:") print(f" GRU weights: {gru_weights}") print(f" Noise weights: {noise_weights}") print(f" Scale weights: {scale_weights}") print(f" Shift weights: {shift_weights}") print(f" Smooth weights: {smooth_weights}") # 生成所有参数组合 all_combinations = list(itertools.product(gru_weights, noise_weights, scale_weights, shift_weights, smooth_weights)) total_combinations = len(all_combinations) print(f"Total combinations to evaluate: {total_combinations}") best_per = float('inf') best_config = None results = [] with tqdm(total=total_combinations, desc='Parameter search', unit='config') as pbar: for gru_w, noise_w, scale_w, shift_w, smooth_w in all_combinations: tta_weights = { 'original': 1.0, # 总是包含原始数据 'noise': noise_w, 'scale': scale_w, 'shift': shift_w, 'smooth': smooth_w } per = ensemble_and_evaluate(predictions_cache, gru_w, tta_weights, args.eval_type) config = { 'gru_weight': gru_w, 'lstm_weight': 1.0 - gru_w, 'tta_weights': tta_weights, 'per': per } results.append(config) if per < best_per: best_per = per best_config = config print(f"\n🎯 New best PER: {per:.3f}%") print(f" GRU weight: {gru_w:.1f}") print(f" TTA weights: {tta_weights}") pbar.update(1) return results, best_config def main(): args = parse_arguments() print("TTA-E Parameter Search") print("=" * 50) # 第一阶段：缓存预测结果 if args.force_recache or not os.path.exists(args.cache_file): predictions_cache = cache_model_predictions(args) else: print(f"Loading existing predictions cache from {args.cache_file}...") with open(args.cache_file, 'rb') as f: predictions_cache = pickle.load(f) print("✓ Cache loaded successfully!") # 第二阶段：参数搜索 results, best_config = search_optimal_parameters(predictions_cache, args) # 保存结果 print(f"\n=== 搜索完成 ===") print(f"Best configuration:") print(f" PER: {best_config['per']:.3f}%") print(f" GRU weight: {best_config['gru_weight']:.1f}") print(f" LSTM weight: {best_config['lstm_weight']:.1f}") print(f" TTA weights: {best_config['tta_weights']}") # 保存所有结果 search_results = { 'best_config': best_config, 'all_results': results, 'search_args': vars(args), 'timestamp': time.strftime("%Y-%m-%d %H:%M:%S") } with open(args.results_file, 'w') as f: json.dump(search_results, f, indent=2) print(f"\n✓ Results saved to {args.results_file}") # 显示前10个最佳配置 sorted_results = sorted(results, key=lambda x: x['per']) print(f"\nTop 10 configurations:") for i, config in enumerate(sorted_results[:10]): print(f"{i+1:2d}. PER={config['per']:6.3f}% | GRU={config['gru_weight']:.1f} | TTA={config['tta_weights']}") if __name__ == "__main__": main()