""" Wakeword detection trainer implementation. This module provides a complete trainer for wakeword detection models using RepCNN architecture with comprehensive training, validation, and deployment capabilities. """ import os import time import logging from pathlib import Path from typing import Dict, List, Tuple, Optional, Any import numpy as np import torch import torch.nn as nn import torch.optim as optim import torch.nn.functional as F from torch.utils.data import DataLoader from ..base import BaseTrainer, TrainerConfig, TrainingState from ..metadata import ModelMetadata, MetadataManager, ModelType from ..model_formats import ModelFormatManager from ..data_pipeline import AudioProcessingConfig from ..utils import TrainerLogger, ProgressMonitor, ValidationMetrics from ..validation import ModelValidator, ValidationConfig from .models import RepCNN, ImprovedRepCNN, LightweightRepCNN, create_repcnn_model from .data import WakewordDataPreprocessor, create_wakeword_dataloaders class FocalLoss(nn.Module): """ Focal Loss for addressing class imbalance in wakeword detection. Particularly useful when dealing with imbalanced datasets where negative samples significantly outnumber positive wakeword samples. """ def __init__(self, alpha: float = 1.0, gamma: float = 2.0, class_weights: Optional[torch.Tensor] = None): """ Initialize Focal Loss. Args: alpha: Weighting factor for rare class gamma: Focusing parameter class_weights: Optional class weights """ super().__init__() self.alpha = alpha self.gamma = gamma self.class_weights = class_weights def forward(self, inputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor: """Compute focal loss.""" ce_loss = F.cross_entropy(inputs, targets, weight=self.class_weights, reduction='none') pt = torch.exp(-ce_loss) focal_loss = self.alpha * (1 - pt) ** self.gamma * ce_loss return focal_loss.mean() class WakewordTrainer(BaseTrainer): """ Specialized trainer for wakeword detection models. Provides end-to-end training pipeline for wakeword detection including data preprocessing, model training, validation, and deployment preparation. """ def __init__(self, config: TrainerConfig): """Initialize wakeword trainer.""" super().__init__(config) # Wakeword-specific configuration self.wakeword_classes = config.custom_params.get('wakeword_classes', { 0: "custom", 1: "system_command", 2: "negative" }) self.model_type = config.custom_params.get('model_type', 'standard') self.detection_threshold = config.custom_params.get('detection_threshold', 0.5) self.use_focal_loss = config.custom_params.get('use_focal_loss', True) self.class_weights = config.custom_params.get('class_weights', None) # Audio processing configuration self.audio_config = AudioProcessingConfig( sample_rate=config.sample_rate, target_length=config.audio_length, n_mels=config.n_mels, n_fft=config.n_fft, hop_length=config.hop_length, win_length=config.win_length ) # Initialize specialized components self.data_preprocessor = WakewordDataPreprocessor(self.audio_config, self.logger) self.metadata_manager = MetadataManager(self.logger) self.format_manager = ModelFormatManager(self.logger) # Model validation self.validator = None self.logger.info(f"Initialized WakewordTrainer with classes: {self.wakeword_classes}") def prepare_data(self) -> Tuple[DataLoader, DataLoader, DataLoader]: """ Prepare wakeword detection data. Returns: Tuple of (train_loader, val_loader, test_loader) """ self.logger.info("Preparing wakeword detection data...") # Check if data is already preprocessed processed_data_dir = Path(self.config.data_dir) / "processed" if not processed_data_dir.exists() or not any(processed_data_dir.iterdir()): # Preprocess raw data raw_data_dir = Path(self.config.data_dir) / "raw" if not raw_data_dir.exists(): raise FileNotFoundError(f"Raw data directory not found: {raw_data_dir}") self.logger.info("Preprocessing raw wakeword data...") processed_path = self.data_preprocessor.preprocess_wakeword_data( str(raw_data_dir), str(processed_data_dir), self.wakeword_classes ) self.logger.info(f"Data preprocessing completed: {processed_path}") # Create data loaders augmentation_config = { 'noise_factor': self.config.noise_factor, 'speed_factor': self.config.speed_factor, 'pitch_factor': self.config.pitch_factor, 'volume_factor': self.config.volume_factor, 'time_shift_factor': 0.1 } if self.config.use_augmentation else None train_loader, val_loader, test_loader = create_wakeword_dataloaders( str(processed_data_dir), self.wakeword_classes, self.audio_config, batch_size=self.config.batch_size, num_workers=self.config.num_workers, augmentation_config=augmentation_config, balance_training=True ) self.logger.info(f"Created data loaders - Train: {len(train_loader.dataset)}, " f"Val: {len(val_loader.dataset)}, Test: {len(test_loader.dataset)}") return train_loader, val_loader, test_loader def build_model(self) -> nn.Module: """ Build RepCNN model for wakeword detection. Returns: RepCNN model """ self.logger.info(f"Building {self.model_type} RepCNN model...") # Calculate input dimensions time_frames = int(self.audio_config.target_length * self.audio_config.sample_rate // self.audio_config.hop_length) + 1 model_kwargs = { 'num_classes': len(self.wakeword_classes), 'input_channels': 1, 'num_mels': self.audio_config.n_mels, 'time_frames': time_frames, 'width_multiplier': self.config.custom_params.get('width_multiplier', 1.0), 'use_se': self.config.custom_params.get('use_se', False), 'dropout_rate': self.config.custom_params.get('dropout_rate', 0.2) } # Add model-specific parameters if self.model_type == 'improved': model_kwargs['use_temporal_attention'] = self.config.custom_params.get('use_temporal_attention', True) elif self.model_type == 'lightweight': model_kwargs['width_multiplier'] = min(model_kwargs['width_multiplier'], 0.75) model = create_repcnn_model(self.model_type, **model_kwargs) # Log model information total_params = sum(p.numel() for p in model.parameters()) trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad) self.logger.info(f"Model created - Type: {self.model_type}") self.logger.info(f"Total parameters: {total_params:,}") self.logger.info(f"Trainable parameters: {trainable_params:,}") return model def create_criterion(self) -> nn.Module: """ Create loss criterion for wakeword detection. Returns: Loss function """ if self.use_focal_loss: # Use Focal Loss for imbalanced datasets class_weights = None if self.class_weights: class_weights = torch.tensor(self.class_weights, dtype=torch.float32) criterion = FocalLoss( alpha=self.config.custom_params.get('focal_alpha', 1.0), gamma=self.config.custom_params.get('focal_gamma', 2.0), class_weights=class_weights ) self.logger.info("Using Focal Loss for training") else: # Use standard Cross Entropy Loss class_weights = None if self.class_weights: class_weights = torch.tensor(self.class_weights, dtype=torch.float32) criterion = nn.CrossEntropyLoss(weight=class_weights) self.logger.info("Using Cross Entropy Loss for training") return criterion def train(self) -> Any: """ Train the wakeword detection model with comprehensive tracking. Returns: Training metrics and results """ self.logger.info("Starting wakeword detection training...") # Setup training components self.setup_training() # Initialize model validator self.validator = ModelValidator(self.model, self.config.device, self.logger) # Create metadata metadata = self.metadata_manager.create_metadata( ModelType.WAKEWORD, self.config.model_name, "Trixy ML Trainer" ) metadata.description = f"RepCNN wakeword detection model ({self.model_type})" metadata.update_from_training_config(self.config) metadata.update_from_model(self.model) metadata.update_from_dataset(self.train_loader, self.val_loader, self.test_loader) metadata.update_wakeword_info(self.wakeword_classes, { 'detection_threshold': self.detection_threshold }) # Run parent training loop training_metrics = super().train() # Update metadata with training results metadata.update_from_training_results(training_metrics) # Validate model validation_config = ValidationConfig( test_augmentations=True, robustness_tests=True, performance_profiling=True ) self.logger.info("Running comprehensive model validation...") validation_results = self.validator.validate(self.val_loader, self.criterion, validation_config) metadata.add_test_results(validation_results) # Test model test_results = self.validator.test(self.test_loader, self.criterion, validation_config) # Save model with metadata self._save_trained_model(metadata, test_results) # Generate training report self._generate_training_report(metadata, training_metrics, test_results) self.logger.info("Wakeword training completed successfully!") return { 'training_metrics': training_metrics, 'validation_results': validation_results, 'test_results': test_results, 'metadata': metadata } def _save_trained_model(self, metadata: ModelMetadata, test_results: Dict[str, Any]): """Save the trained model with comprehensive metadata.""" model_dir = Path(self.config.output_dir) / self.config.model_name model_dir.mkdir(parents=True, exist_ok=True) # Prepare model for saving self.model.eval() # Reparameterize RepCNN if applicable if hasattr(self.model, 'reparameterize'): self.logger.info("Reparameterizing model for inference...") self.model.reparameterize() # Update metadata with final model info metadata.set_file_info(str(model_dir / f"{self.config.model_name}.pth")) metadata.add_test_results(test_results) # Save model in requested format model_file = model_dir / f"{self.config.model_name}{self.config.model_format.value}" success = self.format_manager.save_model( self.model, str(model_file), metadata.to_dict(), password=self.config.password if self.config.use_password_protection else None ) if success: self.logger.info(f"Model saved successfully: {model_file}") else: self.logger.error(f"Failed to save model: {model_file}") # Save standalone metadata file metadata_file = model_dir / "metadata.json" metadata.save_to_file(str(metadata_file)) # Save training configuration config_file = model_dir / "training_config.json" with open(config_file, 'w') as f: import json json.dump(self.config.to_dict(), f, indent=2) # Save validation results if hasattr(self, 'validator') and self.validator: self.validator.save_results(str(model_dir / "validation")) def _generate_training_report(self, metadata: ModelMetadata, training_metrics: Any, test_results: Dict[str, Any]): """Generate comprehensive training report.""" report_dir = Path(self.config.output_dir) / self.config.model_name / "reports" report_dir.mkdir(parents=True, exist_ok=True) # Generate metadata report report = self.metadata_manager.create_training_report(metadata) # Add wakeword-specific information report += "\n## Wakeword Detection Specific Results\n\n" # Add detection metrics if 'standard_test' in test_results: std_test = test_results['standard_test'] if 'metrics' in std_test: metrics = std_test['metrics'] report += "### Classification Performance\n" for class_id, class_name in self.wakeword_classes.items(): if f'f1_{class_name}' in metrics: precision = metrics.get(f'precision_{class_name}', 0) recall = metrics.get(f'recall_{class_name}', 0) f1 = metrics.get(f'f1_{class_name}', 0) report += f"- **{class_name.title()}**:\n" report += f" - Precision: {precision:.4f}\n" report += f" - Recall: {recall:.4f}\n" report += f" - F1-Score: {f1:.4f}\n" # Add confusion matrix info if 'confusion_matrix' in metrics: report += "\n### Confusion Matrix\n" cm = metrics['confusion_matrix'] report += "```\n" class_names = list(self.wakeword_classes.values()) # Header report += "Actual\\Predicted " for name in class_names: report += f"{name[:8]:>8} " report += "\n" # Matrix rows for i, row in enumerate(cm): if i < len(class_names): report += f"{class_names[i][:12]:>12} " for val in row: report += f"{val:>8} " report += "\n" report += "```\n\n" # Add robustness test results if 'robustness_tests' in test_results: report += "### Robustness Tests\n" robustness = test_results['robustness_tests'] if 'input_corruptions' in robustness: report += "#### Input Corruption Robustness\n" for corruption, results in robustness['input_corruptions'].items(): accuracy = results.get('accuracy', 0) report += f"- {corruption}: {accuracy:.4f} accuracy\n" report += "\n" # Add performance profiling if 'performance_profile' in test_results: profile = test_results['performance_profile'] report += "### Performance Profile\n" if 'inference_performance' in profile: perf = profile['inference_performance'] report += f"- Inference time: {perf['mean_inference_time']*1000:.2f} ms\n" report += f"- Throughput: {perf['throughput_samples_per_second']:.1f} samples/sec\n" if 'model_size' in profile: size = profile['model_size'] report += f"- Model size: {size['total_size_mb']:.2f} MB\n" report += "\n" # Add deployment recommendations report += "## Deployment Recommendations\n\n" report += "### Detection Threshold\n" report += f"- Recommended threshold: {self.detection_threshold}\n" report += "- Monitor false positive/negative rates in production\n" report += "- Consider threshold adjustment based on use case requirements\n\n" report += "### Hardware Requirements\n" if 'performance_profile' in test_results: profile = test_results['performance_profile'] inference_time = profile.get('inference_performance', {}).get('mean_inference_time', 0) * 1000 if inference_time < 50: report += "- Suitable for real-time processing on most devices\n" elif inference_time < 100: report += "- Suitable for real-time processing on modern devices\n" else: report += "- May require optimization for real-time processing\n" report += "\n### Model Optimization\n" if self.model_type == 'lightweight': report += "- Already optimized for edge deployment\n" else: report += "- Consider using lightweight variant for edge deployment\n" if hasattr(self.model, 'reparameterize'): report += "- Model has been reparameterized for efficient inference\n" # Save report report_file = report_dir / "training_report.md" with open(report_file, 'w', encoding='utf-8') as f: f.write(report) self.logger.info(f"Training report saved: {report_file}") def evaluate_detection_performance(self, test_loader: DataLoader, thresholds: List[float] = None) -> Dict[str, Any]: """ Evaluate wakeword detection performance at different thresholds. Args: test_loader: Test data loader thresholds: List of detection thresholds to evaluate Returns: Detection performance results """ if thresholds is None: thresholds = [0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9] self.logger.info("Evaluating detection performance...") self.model.eval() all_predictions = [] all_probabilities = [] all_targets = [] with torch.no_grad(): for data, targets in test_loader: data = data.to(self.config.device) targets = targets.to(self.config.device) outputs = self.model(data) probabilities = F.softmax(outputs, dim=1) all_probabilities.append(probabilities.cpu().numpy()) all_targets.append(targets.cpu().numpy()) all_probabilities = np.concatenate(all_probabilities) all_targets = np.concatenate(all_targets) # Evaluate at different thresholds results = {} for threshold in thresholds: # Get predictions based on threshold max_probs = np.max(all_probabilities, axis=1) predictions = np.argmax(all_probabilities, axis=1) # Apply threshold (predict negative if below threshold) negative_class_id = max([id for id, name in self.wakeword_classes.items() if name in ['negative', 'background']], default=2) low_confidence_mask = max_probs < threshold predictions[low_confidence_mask] = negative_class_id # Calculate metrics accuracy = np.mean(predictions == all_targets) # Calculate per-class metrics class_metrics = {} for class_id, class_name in self.wakeword_classes.items(): mask = all_targets == class_id if np.sum(mask) > 0: class_acc = np.mean(predictions[mask] == all_targets[mask]) class_metrics[class_name] = { 'accuracy': class_acc, 'samples': np.sum(mask) } results[f"threshold_{threshold}"] = { 'threshold': threshold, 'overall_accuracy': accuracy, 'class_metrics': class_metrics } self.logger.info("Detection performance evaluation completed") return results def create_wakeword_trainer_config(model_name: str = "wakeword_model", model_type: str = "standard", data_dir: str = "./trainer/data/wakeword", **kwargs) -> TrainerConfig: """ Create a TrainerConfig specifically configured for wakeword detection. Args: model_name: Name of the model model_type: Type of RepCNN model ('standard', 'improved', 'lightweight') data_dir: Directory containing wakeword data **kwargs: Additional configuration parameters Returns: TrainerConfig instance for wakeword detection """ # Default wakeword configuration config_dict = { 'trainer_name': 'wakeword_trainer', 'model_name': model_name, 'data_dir': data_dir, 'output_dir': './models/wakeword', # Training parameters optimized for wakeword detection 'batch_size': 64, 'learning_rate': 0.001, 'num_epochs': 100, 'min_epochs': 20, # Allow at least 20 epochs for wakeword model to learn patterns 'early_stopping_patience': 15, # Audio parameters for wakeword detection 'sample_rate': 16000, 'audio_length': 1.5, # Typical wakeword length 'n_mels': 40, 'n_fft': 512, 'hop_length': 160, 'win_length': 400, # Augmentation for robustness 'use_augmentation': True, 'noise_factor': 0.1, 'speed_factor': 0.1, 'pitch_factor': 0.05, 'volume_factor': 0.2, # Model-specific parameters 'custom_params': { 'model_type': model_type, 'wakeword_classes': {0: "custom", 1: "system_command", 2: "negative"}, 'detection_threshold': 0.5, 'use_focal_loss': True, 'focal_alpha': 1.0, 'focal_gamma': 2.0, 'width_multiplier': 1.0, 'use_se': False, 'dropout_rate': 0.2 } } # Separate custom parameters from standard TrainerConfig parameters custom_param_keys = { 'use_focal_loss', 'detection_threshold', 'focal_alpha', 'focal_gamma', 'width_multiplier', 'use_se', 'dropout_rate', 'wakeword_classes', 'use_temporal_attention' } # Extract custom parameters from kwargs custom_params = config_dict['custom_params'].copy() for key, value in kwargs.items(): if key in custom_param_keys: custom_params[key] = value else: config_dict[key] = value # Handle model_type specially if 'model_type' in kwargs: custom_params['model_type'] = kwargs['model_type'] # Update custom_params config_dict['custom_params'] = custom_params return TrainerConfig(**config_dict)