import os import tensorflow as tf import numpy as np import time import json import pickle import logging import pathlib import sys from typing import Dict, Any, Tuple, Optional, List from omegaconf import OmegaConf # For accurate PER calculation try: import editdistance except ImportError: print("Warning: editdistance not available, falling back to approximation") editdistance = None from rnn_model_tf import ( TripleGRUDecoder, CTCLoss, create_tpu_strategy, build_model_for_tpu, configure_mixed_precision ) from dataset_tf import ( BrainToTextDatasetTF, DataAugmentationTF, train_test_split_indices, create_input_fn ) class BrainToTextDecoderTrainerTF: """ TensorFlow/Keras trainer for brain-to-text phoneme decoder optimized for TPU v5e-8 This trainer implements the same training logic as the PyTorch version but uses TensorFlow operations optimized for TPU hardware. """ def __init__(self, args: Dict[str, Any]): """ Initialize the TensorFlow trainer Args: args: Configuration dictionary containing all training parameters """ self.args = args self.logger = None # Initialize TPU strategy self.strategy = create_tpu_strategy() if self.strategy is None: raise RuntimeError("Failed to create TPU strategy - strategy is None") print(f"Training on {self.strategy.num_replicas_in_sync} TPU cores") print(f"Strategy type: {type(self.strategy).__name__}") print("💡 Using tf.data.AUTOTUNE for optimal data pipeline performance") print("📝 Ensure create_input_fn uses AUTOTUNE for .map() and .prefetch() operations") print("⚠️ For best TPU performance, ensure create_input_fn uses padded_batch with fixed shapes") print(" and drop_remainder=True to avoid dynamic shape warnings") # Configure mixed precision for TPU v5e-8 if args.get('use_amp', True): configure_mixed_precision() self.mixed_precision = True else: self.mixed_precision = False # Initialize tracking variables self.best_val_per = float('inf') self.best_val_loss = float('inf') # Setup directories if args['mode'] == 'train': os.makedirs(self.args['output_dir'], exist_ok=True) if (args.get('save_best_checkpoint', True) or args.get('save_all_val_steps', False) or args.get('save_final_model', False)): os.makedirs(self.args['checkpoint_dir'], exist_ok=True) # Setup logging self._setup_logging() # Set random seeds if self.args['seed'] != -1: tf.random.set_seed(self.args['seed']) np.random.seed(self.args['seed']) # Initialize datasets self._initialize_datasets() # Build model within strategy scope with self.strategy.scope(): self.model = self._build_model() self.optimizer = self._create_optimizer() print("🔧 Initializing optimizer for TPU training...") print(f"Optimizer type: {type(self.optimizer).__name__}") # ========================= SOLUTION ========================= # Explicitly build optimizer within strategy scope before training. # This forces creation of all slot variables (e.g., AdamW momentum) # avoiding lazy initialization inside @tf.function which loses context. # Note: Model must be built first for .build() to work. # The _log_model_info method builds the model via forward pass. # Ensure model is built (will be called later in _log_model_info anyway) if not self.model.built: dummy_batch_size = 2 dummy_time_steps = 100 dummy_features = tf.zeros((dummy_batch_size, dummy_time_steps, self.args['model']['n_input_features'])) dummy_day_idx = tf.zeros((dummy_batch_size,), dtype=tf.int32) _ = self.model(dummy_features, dummy_day_idx, training=False) print("🔧 Building optimizer with model variables...") self.optimizer.build(self.model.trainable_variables) print("✅ Optimizer built successfully") # ============================================================ print("✅ Optimizer ready for TPU training") self.lr_scheduler = self._create_lr_scheduler() self.ctc_loss = CTCLoss(blank_index=0, reduction='none') # Create unified checkpoint management self.ckpt = tf.train.Checkpoint( optimizer=self.optimizer, model=self.model ) self.ckpt_manager = tf.train.CheckpointManager( self.ckpt, directory=self.args['checkpoint_dir'], max_to_keep=5 # Keep only the 5 most recent checkpoints ) # Try to restore from latest checkpoint if self.ckpt_manager.latest_checkpoint: print(f"🔄 Restoring from {self.ckpt_manager.latest_checkpoint}") if self.logger: self.logger.info(f"Restoring from {self.ckpt_manager.latest_checkpoint}") # expect_partial() avoids failures due to model structure changes self.ckpt.restore(self.ckpt_manager.latest_checkpoint).expect_partial() print("✅ Checkpoint restored successfully") else: print("🆕 Initializing training from scratch") if self.logger: self.logger.info("Initializing training from scratch") # Log model information self._log_model_info() # Adversarial training configuration adv_cfg = self.args.get('adversarial', {}) self.adv_enabled = adv_cfg.get('enabled', False) self.adv_grl_lambda = float(adv_cfg.get('grl_lambda', 0.5)) self.adv_noisy_loss_weight = float(adv_cfg.get('noisy_loss_weight', 0.2)) self.adv_noise_l2_weight = float(adv_cfg.get('noise_l2_weight', 0.0)) self.adv_warmup_steps = int(adv_cfg.get('warmup_steps', 0)) # Manual weight decay handling - disable since AdamW handles it self.manual_weight_decay = False if self.args.get('weight_decay', 0.0) > 0: print(f"🔧 Weight decay configured in AdamW: {self.args.get('weight_decay', 0.0)}") else: print("💡 No weight decay configured") if self.adv_enabled: if self.logger: self.logger.info(f"Adversarial training ENABLED | grl_lambda={self.adv_grl_lambda}, " f"noisy_loss_weight={self.adv_noisy_loss_weight}, " f"noise_l2_weight={self.adv_noise_l2_weight}, " f"warmup_steps={self.adv_warmup_steps}") else: print(f"Adversarial training ENABLED | grl_lambda={self.adv_grl_lambda}, " f"noisy_loss_weight={self.adv_noisy_loss_weight}, " f"noise_l2_weight={self.adv_noise_l2_weight}, " f"warmup_steps={self.adv_warmup_steps}") def _setup_logging(self): """Setup logging configuration""" self.logger = logging.getLogger(__name__) for handler in self.logger.handlers[:]: self.logger.removeHandler(handler) self.logger.setLevel(logging.INFO) formatter = logging.Formatter(fmt='%(asctime)s: %(message)s') if self.args['mode'] == 'train': fh = logging.FileHandler(str(pathlib.Path(self.args['output_dir'], 'training_log'))) fh.setFormatter(formatter) self.logger.addHandler(fh) sh = logging.StreamHandler(sys.stdout) sh.setFormatter(formatter) self.logger.addHandler(sh) self.logger.info(f'Using TPU strategy with {self.strategy.num_replicas_in_sync} replicas') if self.mixed_precision: self.logger.info('Mixed precision (bfloat16) enabled for TPU training') def _configure_cpu_optimization(self): """Configure CPU utilization to make use of 224 cores for data pipeline""" import multiprocessing # Get available CPU cores available_cores = multiprocessing.cpu_count() print(f"💻 Available CPU cores: {available_cores}") # Optimize for data pipeline parallelism # For 224 cores, use more threads for better data loading performance if available_cores >= 200: # High core count system inter_op_threads = min(64, available_cores // 3) # More aggressive for 224 cores intra_op_threads = min(32, available_cores // 6) else: # Use ~1/4 of cores for inter-op (between operations) # Use ~1/8 of cores for intra-op (within operations) inter_op_threads = min(32, available_cores // 4) intra_op_threads = min(16, available_cores // 8) tf.config.threading.set_inter_op_parallelism_threads(inter_op_threads) tf.config.threading.set_intra_op_parallelism_threads(intra_op_threads) print(f"🔧 CPU optimization configured:") print(f" Inter-op parallelism: {inter_op_threads} threads") print(f" Intra-op parallelism: {intra_op_threads} threads") print(f" This will accelerate data loading and preprocessing while TPU handles training") def _get_tpu_status(self) -> str: """Get current TPU status and HBM utilization info""" try: # Get TPU devices tpu_devices = tf.config.list_logical_devices('TPU') if not tpu_devices: return "TPU: No devices" # Get strategy info num_replicas = self.strategy.num_replicas_in_sync if hasattr(self.strategy, 'num_replicas_in_sync') else 1 # Get TPU memory info using the working /device:TPU:X format try: # Check all TPU devices for memory usage active_cores = 0 total_current_mb = 0 max_peak_mb = 0 for device in tpu_devices: try: memory_info = tf.config.experimental.get_memory_info(device.name) if memory_info and 'current' in memory_info: current_mb = memory_info['current'] // (1024 * 1024) peak_mb = memory_info.get('peak', memory_info['current']) // (1024 * 1024) if current_mb > 1: # >1MB considered active active_cores += 1 total_current_mb += current_mb max_peak_mb = max(max_peak_mb, peak_mb) except: continue if active_cores > 0: if active_cores == 1: hbm_info = f"HBM:{total_current_mb}MB(peak:{max_peak_mb}MB)" else: hbm_info = f"HBM:{total_current_mb}MB/{active_cores}cores(peak:{max_peak_mb}MB)" else: hbm_info = "HBM:idle" except Exception: # Fallback: simple TPU activity check try: with tf.device('/TPU:0'): _ = tf.constant(1.0) hbm_info = "HBM:active" except Exception: hbm_info = "HBM:inactive" return (f"TPU: {len(tpu_devices)}dev {num_replicas}cores " f"{hbm_info}") except Exception as e: return f"TPU: status_error({str(e)[:20]})" def _get_detailed_tpu_status(self) -> str: """Get detailed TPU status for training start""" try: # Get TPU devices tpu_devices = tf.config.list_logical_devices('TPU') if not tpu_devices: return "❌ No TPU devices detected" # Get strategy info num_replicas = self.strategy.num_replicas_in_sync if hasattr(self.strategy, 'num_replicas_in_sync') else 1 strategy_type = type(self.strategy).__name__ # Get TPU HBM memory info using working device format try: active_cores = 0 total_current_gb = 0 max_peak_gb = 0 memory_details = [] for i, device in enumerate(tpu_devices): try: memory_info = tf.config.experimental.get_memory_info(device.name) if memory_info and 'current' in memory_info: current_gb = memory_info['current'] // (1024 * 1024 * 1024) peak_gb = memory_info.get('peak', memory_info['current']) // (1024 * 1024 * 1024) if current_gb > 0 or memory_info['current'] > 1024*1024: # >1MB active_cores += 1 total_current_gb += current_gb max_peak_gb = max(max_peak_gb, peak_gb) if current_gb > 0: memory_details.append(f"Core{i}:{current_gb}GB") except: continue if active_cores > 0: # Based on your test: TPU:0 peaked at 14.5GB, suggesting ~16GB per core estimated_per_core = 16 # Conservative estimate estimated_total_gb = estimated_per_core * len(tpu_devices) hbm_usage = f"HBM: {total_current_gb}GB/{estimated_total_gb}GB (peak: {max_peak_gb}GB) active:{active_cores}cores" else: hbm_usage = "HBM: 0GB/256GB (idle)" except Exception: hbm_usage = "HBM: unavailable" # Simple TPU test try: with tf.device('/TPU:0'): test_result = tf.constant([1.0, 2.0]) _ = tf.reduce_sum(test_result) tpu_test = "✅ responsive" except Exception as e: tpu_test = f"❌ test_failed({str(e)[:15]})" return (f"TPU Devices: {len(tpu_devices)} | " f"Strategy: {strategy_type} | " f"Cores: {num_replicas} | " f"{hbm_usage} | " f"Test: {tpu_test}") except Exception as e: return f"❌ TPU status check failed: {str(e)[:50]}" def _initialize_datasets(self): """Initialize training and validation datasets""" # Create file paths train_file_paths = [ os.path.join(self.args["dataset"]["dataset_dir"], s, 'data_train.hdf5') for s in self.args['dataset']['sessions'] ] val_file_paths = [ os.path.join(self.args["dataset"]["dataset_dir"], s, 'data_val.hdf5') for s in self.args['dataset']['sessions'] ] # Validate no duplicates if len(set(train_file_paths)) != len(train_file_paths): raise ValueError("Duplicate sessions in train dataset") if len(set(val_file_paths)) != len(val_file_paths): raise ValueError("Duplicate sessions in val dataset") # Split trials train_trials, _ = train_test_split_indices( file_paths=train_file_paths, test_percentage=0, seed=self.args['dataset']['seed'], bad_trials_dict=self.args['dataset'].get('bad_trials_dict') ) _, val_trials = train_test_split_indices( file_paths=val_file_paths, test_percentage=1, seed=self.args['dataset']['seed'], bad_trials_dict=self.args['dataset'].get('bad_trials_dict') ) # Save trial splits with open(os.path.join(self.args['output_dir'], 'train_val_trials.json'), 'w') as f: json.dump({'train': train_trials, 'val': val_trials}, f) # Create TensorFlow datasets with aggressive data preloading for TPU optimization # Monitor memory usage during data preloading import psutil initial_memory_mb = psutil.Process().memory_info().rss / 1024 / 1024 print("🔄 Initializing training dataset with GPU-style memory management...") init_start_time = time.time() self.train_dataset_tf = BrainToTextDatasetTF( trial_indices=train_trials, n_batches=self.args['num_training_batches'], split='train', batch_size=self.args['dataset']['batch_size'], days_per_batch=self.args['dataset']['days_per_batch'], random_seed=self.args['dataset']['seed'], must_include_days=self.args['dataset'].get('must_include_days'), feature_subset=self.args['dataset'].get('feature_subset'), cache_data=True, # 启用智能缓存（像GPU版本一样） preload_all_data=False # 🚨 采用GPU版本策略：按需加载，避免内存溢出 ) # Log training dataset initialization performance train_init_time = time.time() - init_start_time train_memory_mb = psutil.Process().memory_info().rss / 1024 / 1024 train_memory_used = train_memory_mb - initial_memory_mb print(f"✅ Training dataset initialized in {train_init_time:.2f}s, using {train_memory_used:.1f} MB RAM") print("🔄 Initializing validation dataset with GPU-style memory management...") val_init_start_time = time.time() self.val_dataset_tf = BrainToTextDatasetTF( trial_indices=val_trials, n_batches=None, # Use all validation data split='test', batch_size=self.args['dataset']['batch_size'], days_per_batch=1, # One day per validation batch random_seed=self.args['dataset']['seed'], feature_subset=self.args['dataset'].get('feature_subset'), cache_data=True, # 启用智能缓存（像GPU版本一样） preload_all_data=False # 🚨 采用GPU版本策略：按需加载，避免内存溢出 ) # Log validation dataset initialization performance val_init_time = time.time() - val_init_start_time final_memory_mb = psutil.Process().memory_info().rss / 1024 / 1024 total_memory_used = final_memory_mb - initial_memory_mb val_memory_used = final_memory_mb - train_memory_mb print(f"✅ Validation dataset initialized in {val_init_time:.2f}s, using {val_memory_used:.1f} MB RAM") print(f"📊 Total data cache: {total_memory_used:.1f} MB RAM used for all datasets") self.logger.info("Successfully initialized TensorFlow datasets") def _build_model(self) -> TripleGRUDecoder: """Build the TripleGRUDecoder model""" model = TripleGRUDecoder( neural_dim=self.args['model']['n_input_features'], n_units=self.args['model']['n_units'], n_days=len(self.args['dataset']['sessions']), n_classes=self.args['dataset']['n_classes'], rnn_dropout=self.args['model']['rnn_dropout'], input_dropout=self.args['model']['input_network']['input_layer_dropout'], patch_size=self.args['model']['patch_size'], patch_stride=self.args['model']['patch_stride'] ) return model def _create_optimizer(self) -> tf.keras.optimizers.Optimizer: """Create AdamW optimizer""" print(f"Creating optimizer with strategy: {type(self.strategy).__name__}") print("Using AdamW optimizer for TPU training") optimizer = tf.keras.optimizers.AdamW( learning_rate=self.args['lr_max'], beta_1=self.args['beta0'], beta_2=self.args['beta1'], epsilon=self.args['epsilon'], weight_decay=self.args.get('weight_decay', 0.0) ) print("✅ Using AdamW optimizer") return optimizer def _create_lr_scheduler(self): """Create learning rate scheduler""" if self.args['lr_scheduler_type'] == 'cosine': return self._create_cosine_scheduler() elif self.args['lr_scheduler_type'] == 'linear': return tf.keras.optimizers.schedules.PolynomialDecay( initial_learning_rate=self.args['lr_max'], decay_steps=self.args['lr_decay_steps'], end_learning_rate=self.args['lr_min'], power=1.0 # Linear decay ) else: raise ValueError(f"Unknown scheduler type: {self.args['lr_scheduler_type']}") def _create_cosine_scheduler(self): """Create cosine learning rate scheduler""" return tf.keras.optimizers.schedules.CosineDecayRestarts( initial_learning_rate=self.args['lr_max'], first_decay_steps=self.args['lr_decay_steps'], t_mul=1.0, m_mul=1.0, alpha=self.args['lr_min'] / self.args['lr_max'] ) def _log_model_info(self): """Log model architecture and parameter information""" if self.logger: self.logger.info("Initialized TripleGRUDecoder model") else: print("Initialized TripleGRUDecoder model") # Build the model by calling it once with dummy data dummy_batch_size = 2 dummy_time_steps = 100 dummy_features = tf.zeros((dummy_batch_size, dummy_time_steps, self.args['model']['n_input_features'])) dummy_day_idx = tf.zeros((dummy_batch_size,), dtype=tf.int32) # Call the model to build it _ = self.model(dummy_features, dummy_day_idx, training=False) # Count parameters total_params = sum([tf.size(w).numpy() for w in self.model.trainable_weights]) if self.logger: self.logger.info(f"Model has {total_params:,} trainable parameters") else: print(f"Model has {total_params:,} trainable parameters") @tf.function def _train_step(self, batch, step): """Single training step with gradient tape""" features = batch['input_features'] labels = batch['seq_class_ids'] n_time_steps = batch['n_time_steps'] phone_seq_lens = batch['phone_seq_lens'] day_indices = batch['day_indices'] with tf.GradientTape() as tape: # Apply data transformations features, n_time_steps = DataAugmentationTF.transform_data( features, n_time_steps, self.args['dataset']['data_transforms'], training=True ) # Calculate adjusted lengths for CTC adjusted_lens = tf.cast( (tf.cast(n_time_steps, tf.float32) - self.args['model']['patch_size']) / self.args['model']['patch_stride'] + 1, tf.int32 ) # Forward pass use_full = self.adv_enabled and (step >= self.adv_warmup_steps) if use_full: clean_logits, noisy_logits, noise_output = self.model( features, day_indices, None, False, 'full', grl_lambda=self.adv_grl_lambda, training=True ) else: clean_logits = self.model( features, day_indices, None, False, 'inference', training=True ) # Calculate losses if use_full: # Clean CTC loss clean_loss_input = { 'labels': labels, 'input_lengths': adjusted_lens, 'label_lengths': phone_seq_lens } clean_loss = self.ctc_loss(clean_loss_input, clean_logits) clean_loss = tf.reduce_mean(clean_loss) # Noisy CTC loss noisy_loss_input = { 'labels': labels, 'input_lengths': adjusted_lens, 'label_lengths': phone_seq_lens } noisy_loss = self.ctc_loss(noisy_loss_input, noisy_logits) noisy_loss = tf.reduce_mean(noisy_loss) # Optional noise L2 regularization noise_l2 = tf.constant(0.0, dtype=clean_loss.dtype) if self.adv_noise_l2_weight > 0.0: noise_l2 = tf.reduce_mean(tf.square(noise_output)) loss = clean_loss + self.adv_noisy_loss_weight * noisy_loss + self.adv_noise_l2_weight * noise_l2 else: loss_input = { 'labels': labels, 'input_lengths': adjusted_lens, 'label_lengths': phone_seq_lens } loss = self.ctc_loss(loss_input, clean_logits) loss = tf.reduce_mean(loss) # AdamW handles weight decay automatically - no manual L2 regularization needed # TensorFlow混合精度处理：不需要手动scaling，Keras policy自动处理 # TPU v5e-8使用bfloat16，不需要loss scaling # Calculate gradients - TensorFlow自动处理混合精度 gradients = tape.gradient(loss, self.model.trainable_variables) # For TPU compatibility, use all variables (TensorFlow will handle None gradients automatically) # This ensures consistency with slot variable initialization all_variables = self.model.trainable_variables # Replace None gradients with zeros to maintain consistency safe_gradients = [] for grad, var in zip(gradients, all_variables): if grad is not None: safe_gradients.append(grad) else: # Create zero gradient for variables without gradients safe_gradients.append(tf.zeros_like(var)) # Clip gradients if self.args['grad_norm_clip_value'] > 0: safe_gradients, grad_norm = tf.clip_by_global_norm( safe_gradients, self.args['grad_norm_clip_value'] ) else: grad_norm = tf.global_norm(safe_gradients) # Apply gradients to ALL variables (consistent with initialization) # TensorFlow optimizer will handle zero gradients correctly self.optimizer.apply_gradients(zip(safe_gradients, all_variables)) return loss, grad_norm @tf.function def _validation_step(self, batch): """Single validation step - returns data for accurate PER calculation""" features = batch['input_features'] labels = batch['seq_class_ids'] n_time_steps = batch['n_time_steps'] phone_seq_lens = batch['phone_seq_lens'] day_indices = batch['day_indices'] # Apply data transformations (no augmentation for validation) features, n_time_steps = DataAugmentationTF.transform_data( features, n_time_steps, self.args['dataset']['data_transforms'], training=False ) # Calculate adjusted lengths adjusted_lens = tf.cast( (tf.cast(n_time_steps, tf.float32) - self.args['model']['patch_size']) / self.args['model']['patch_stride'] + 1, tf.int32 ) # Forward pass (inference mode only) logits = self.model(features, day_indices, None, False, 'inference', training=False) # Calculate loss loss_input = { 'labels': labels, 'input_lengths': adjusted_lens, 'label_lengths': phone_seq_lens } loss = self.ctc_loss(loss_input, logits) loss = tf.reduce_mean(loss) # Greedy decoding for PER calculation predicted_ids = tf.argmax(logits, axis=-1, output_type=tf.int32) # Return all necessary data for accurate PER calculation on CPU return loss, predicted_ids, labels, adjusted_lens, phone_seq_lens def train(self) -> Dict[str, Any]: """Main training loop""" self.logger.info("Starting training loop...") # Log initial TPU status initial_tpu_status = self._get_detailed_tpu_status() self.logger.info(f"Initial TPU Status: {initial_tpu_status}") # Create datasets using modern distribution API def create_dist_dataset_fn(input_dataset_tf, training): """Create distributed dataset function for modern TPU strategy""" def dataset_fn(input_context): # create_input_fn returns a complete, batched tf.data.Dataset return create_input_fn( input_dataset_tf, self.args['dataset']['data_transforms'], training=training ) return self.strategy.distribute_datasets_from_function(dataset_fn) # Distribute datasets using modern API self.logger.info("🔄 Distributing training dataset across TPU cores...") dist_start_time = time.time() train_dist_dataset = create_dist_dataset_fn(self.train_dataset_tf, training=True) train_dist_time = time.time() - dist_start_time self.logger.info(f"✅ Training dataset distributed in {train_dist_time:.2f}s") self.logger.info("🔄 Distributing validation dataset across TPU cores...") val_start_time = time.time() val_dist_dataset = create_dist_dataset_fn(self.val_dataset_tf, training=False) val_dist_time = time.time() - val_start_time self.logger.info(f"✅ Validation dataset distributed in {val_dist_time:.2f}s") self.logger.info("Created distributed training and validation datasets") # Training metrics train_losses = [] val_losses = [] val_pers = [] val_results = [] val_steps_since_improvement = 0 self.logger.info("Training time count beginning...") train_start_time = time.time() # Training loop step = 0 self.logger.info("🔄 Starting training loop...") self.logger.info("📋 Note: If you see 'TPU has inputs with dynamic shapes' warnings,") self.logger.info(" consider using padded_batch with fixed shapes in create_input_fn") for batch in train_dist_dataset: if step >= self.args['num_training_batches']: self.logger.info("Reached maximum training batches, stopping training") break start_time = time.time() # Distributed training step self.logger.info("Running distributed training step...") # Ensure we're in the correct TPU strategy scope try: with self.strategy.scope(): per_replica_losses, per_replica_grad_norms = self.strategy.run( self._train_step, args=(batch, step) ) except AttributeError as e: if "merge_call" in str(e): error_msg = f"Strategy merge_call error at step {step}: {e}" print(error_msg) if self.logger: self.logger.error(error_msg) self.logger.error("This indicates the strategy is not properly initialized") raise RuntimeError(f"TPU strategy failed during training step {step}: {e}") else: raise # Reduce across replicas self.logger.info("Reducing results across replicas...") loss = self.strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_losses, axis=None) grad_norm = self.strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_grad_norms, axis=None) train_step_duration = time.time() - start_time train_losses.append(float(loss.numpy())) # Log training progress with TPU status if step % self.args['batches_per_train_log'] == 0: tpu_status = self._get_tpu_status() self.logger.info(f'Train batch {step}: ' f'loss: {float(loss.numpy()):.2f} ' f'grad norm: {float(grad_norm.numpy()):.2f} ' f'time: {train_step_duration:.3f}s ' f'| {tpu_status}') # Validation step if step % self.args['batches_per_val_step'] == 0 or step == (self.args['num_training_batches'] - 1): self.logger.info(f"Running validation after training batch: {step}") val_start_time = time.time() val_metrics = self._validate(val_dist_dataset) val_step_duration = time.time() - val_start_time tpu_status = self._get_tpu_status() self.logger.info(f'Val batch {step}: ' f'PER (avg): {val_metrics["avg_per"]:.4f} ' f'CTC Loss (avg): {val_metrics["avg_loss"]:.4f} ' f'time: {val_step_duration:.3f}s ' f'| {tpu_status}') val_pers.append(val_metrics['avg_per']) val_losses.append(val_metrics['avg_loss']) val_results.append(val_metrics) # Check for improvement new_best = False if val_metrics['avg_per'] < self.best_val_per: self.logger.info(f"New best test PER {self.best_val_per:.4f} --> {val_metrics['avg_per']:.4f}") self.best_val_per = val_metrics['avg_per'] self.best_val_loss = val_metrics['avg_loss'] new_best = True elif (val_metrics['avg_per'] == self.best_val_per and val_metrics['avg_loss'] < self.best_val_loss): self.logger.info(f"New best test loss {self.best_val_loss:.4f} --> {val_metrics['avg_loss']:.4f}") self.best_val_loss = val_metrics['avg_loss'] new_best = True if new_best: if self.args.get('save_best_checkpoint', True): self.logger.info("Checkpointing model") self._save_checkpoint(step) if self.args.get('save_val_metrics', True): with open(f'{self.args["checkpoint_dir"]}/val_metrics.pkl', 'wb') as f: pickle.dump(val_metrics, f) val_steps_since_improvement = 0 else: val_steps_since_improvement += 1 # Optional save all validation checkpoints if self.args.get('save_all_val_steps', False): self._save_checkpoint(step) # Early stopping if (self.args.get('early_stopping', False) and val_steps_since_improvement >= self.args.get('early_stopping_val_steps', 20)): self.logger.info(f'Validation PER has not improved in {self.args["early_stopping_val_steps"]} ' f'validation steps. Stopping training early at batch: {step}') break step += 1 # Training completed training_duration = time.time() - train_start_time self.logger.info(f'Best avg val PER achieved: {self.best_val_per:.5f}') self.logger.info(f'Total training time: {(training_duration / 60):.2f} minutes') # Save final model if self.args.get('save_final_model', False): last_loss = val_losses[-1] if len(val_losses) > 0 else float('inf') self._save_checkpoint(step-1) return { 'train_losses': train_losses, 'val_losses': val_losses, 'val_pers': val_pers, 'val_metrics': val_results } def _validate(self, val_dataset) -> Dict[str, Any]: """Run validation on entire validation dataset with accurate PER calculation""" total_loss = 0.0 total_edit_distance = 0 total_seq_length = 0 num_batches = 0 for batch in val_dataset: # Get predictions and labels from all TPU cores per_replica_losses, per_replica_preds, per_replica_labels, per_replica_pred_lens, per_replica_label_lens = ( self.strategy.run(self._validation_step, args=(batch,)) ) # Reduce loss across replicas batch_loss = self.strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_losses, axis=None) total_loss += float(batch_loss.numpy()) # Gather all data from TPU cores to CPU for accurate PER calculation all_preds = self.strategy.gather(per_replica_preds, axis=0) all_labels = self.strategy.gather(per_replica_labels, axis=0) all_pred_lens = self.strategy.gather(per_replica_pred_lens, axis=0) all_label_lens = self.strategy.gather(per_replica_label_lens, axis=0) # Calculate accurate edit distance on CPU batch_size = all_preds.shape[0] for i in range(batch_size): pred_len = int(all_pred_lens[i].numpy()) label_len = int(all_label_lens[i].numpy()) # Extract sequences and remove CTC blanks (assuming blank_index=0) pred_seq = all_preds[i, :pred_len].numpy() pred_seq = [p for p in pred_seq if p != 0] # Remove blanks # Remove consecutive duplicates (CTC decoding) if len(pred_seq) > 0: deduped_pred = [pred_seq[0]] for j in range(1, len(pred_seq)): if pred_seq[j] != pred_seq[j-1]: deduped_pred.append(pred_seq[j]) pred_seq = deduped_pred true_seq = all_labels[i, :label_len].numpy().tolist() # Calculate edit distance using proper library if available if editdistance is not None: edit_dist = editdistance.eval(pred_seq, true_seq) else: # Fallback to simple approximation if editdistance not available edit_dist = self._simple_edit_distance(pred_seq, true_seq) total_edit_distance += edit_dist total_seq_length += label_len num_batches += 1 avg_loss = total_loss / max(num_batches, 1) avg_per = total_edit_distance / max(total_seq_length, 1e-6) return { 'avg_loss': avg_loss, 'avg_per': avg_per, 'total_edit_distance': total_edit_distance, 'total_seq_length': total_seq_length, 'num_batches': num_batches } def _simple_edit_distance(self, seq1, seq2): """Simple edit distance implementation as fallback""" # Dynamic programming implementation of edit distance m, n = len(seq1), len(seq2) dp = [[0] * (n + 1) for _ in range(m + 1)] # Initialize base cases for i in range(m + 1): dp[i][0] = i for j in range(n + 1): dp[0][j] = j # Fill the DP table for i in range(1, m + 1): for j in range(1, n + 1): if seq1[i-1] == seq2[j-1]: dp[i][j] = dp[i-1][j-1] else: dp[i][j] = 1 + min( dp[i-1][j], # deletion dp[i][j-1], # insertion dp[i-1][j-1] # substitution ) return dp[m][n] def _save_checkpoint(self, step: int, name: str = ""): """Save checkpoint using the unified CheckpointManager""" # CheckpointManager automatically handles naming and numbering # The 'name' parameter is kept for backward compatibility but not used save_path = self.ckpt_manager.save(checkpoint_number=step) if self.logger: self.logger.info(f"Saved checkpoint for step {step}: {save_path}") else: print(f"Saved checkpoint for step {step}: {save_path}") # Save non-TensorFlow Python state separately state = { 'step': step, 'best_val_per': float(self.best_val_per), 'best_val_loss': float(self.best_val_loss) } # Associate state file with checkpoint state_path = os.path.join(self.args['checkpoint_dir'], f'state-{step}.json') with open(state_path, 'w') as f: json.dump(state, f) # Save config file (only once) config_path = os.path.join(self.args['checkpoint_dir'], 'args.yaml') if not os.path.exists(config_path): with open(config_path, 'w') as f: OmegaConf.save(config=self.args, f=f) def load_checkpoint(self, checkpoint_path: str): """Load a specific checkpoint and its associated training state""" if self.logger: self.logger.info(f"Loading checkpoint from: {checkpoint_path}") else: print(f"Loading checkpoint from: {checkpoint_path}") # Restore TensorFlow objects (model and optimizer) self.ckpt.restore(checkpoint_path).expect_partial() # Restore non-TensorFlow training state try: # Extract step number from checkpoint path (e.g., ckpt-123 -> 123) step = int(checkpoint_path.split('-')[-1]) state_path = os.path.join(os.path.dirname(checkpoint_path), f'state-{step}.json') with open(state_path, 'r') as f: state = json.load(f) self.best_val_per = state['best_val_per'] self.best_val_loss = state['best_val_loss'] if self.logger: self.logger.info(f"Restored training state from: {state_path}") else: print(f"Restored training state from: {state_path}") except (IOError, ValueError, KeyError) as e: warning_msg = (f"Could not load or parse state file for checkpoint {checkpoint_path}. " f"Starting with fresh state. Error: {e}") if self.logger: self.logger.warning(warning_msg) else: print(f"⚠️ {warning_msg}") def inference(self, features: tf.Tensor, day_indices: tf.Tensor, n_time_steps: tf.Tensor, mode: str = 'inference') -> tf.Tensor: """ Run inference on input features Args: features: Input neural features [batch_size, time_steps, features] day_indices: Day indices [batch_size] n_time_steps: Number of valid time steps [batch_size] mode: 'inference' or 'full' Returns: Phoneme logits [batch_size, time_steps, n_classes] """ # Apply data transformations (no augmentation) features, n_time_steps = DataAugmentationTF.transform_data( features, n_time_steps, self.args['dataset']['data_transforms'], training=False ) # Run model inference logits = self.model(features, day_indices, None, False, mode, training=False) return logits