ASesYusuf1/SESA_Audio_Separation
14
1# coding: utf-82__author__ = 'PyTorch Backend Implementation'3 4import os5import pickle6import numpy as np7import torch8import torch.nn as nn9from typing import Dict, Tuple, Optional, Any10import warnings11import hashlib12import time13 14# Suppress channels_last warnings for 3D audio tensors15warnings.filterwarnings("ignore", message=".*channels_last.*")16warnings.filterwarnings("ignore", message=".*rank 3.*")17 18 19class PyTorchBackend:20 """21 ULTRA-OPTIMIZED PyTorch backend for model inference.22 Provides various optimization techniques for maximum speed.23 """24 25 def __init__(self, device='cuda:0', optimize_mode='channels_last'):26 """27 Initialize ULTRA-OPTIMIZED PyTorch backend.28 29 Parameters:30 ----------31 device : str32 Device to use for inference (cuda:0, cpu, mps, etc.)33 optimize_mode : str34 Optimization mode: 'channels_last' (recommended), 'compile', 'jit', or 'default'35 """36 self.device = device37 self.optimize_mode = optimize_mode38 self.model = None39 self.compiled_model = None40 41 # Check device availability42 if device.startswith('cuda') and not torch.cuda.is_available():43 warnings.warn("CUDA not available, falling back to CPU")44 self.device = 'cpu'45 elif device == 'mps' and not torch.backends.mps.is_available():46 warnings.warn("MPS not available, falling back to CPU")47 self.device = 'cpu'48 49 # Apply ultra optimization settings50 self._apply_ultra_optimizations()51 52 def _apply_ultra_optimizations(self):53 """Apply ultra-speed optimizations globally."""54 if self.device.startswith('cuda'):55 # Enable all CUDA optimizations56 torch.backends.cudnn.benchmark = True57 torch.backends.cuda.matmul.allow_tf32 = True58 torch.backends.cudnn.allow_tf32 = True59 60 # Set optimal CUDA settings61 torch.backends.cudnn.deterministic = False62 torch.backends.cudnn.enabled = True63 64 # Enable cuBLAS optimizations65 os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8'66 67 # Optimize CPU inference68 if self.device == 'cpu':69 import multiprocessing70 num_threads = multiprocessing.cpu_count()71 torch.set_num_threads(num_threads)72 torch.set_num_interop_threads(num_threads)73 print(f"CPU threads set to {num_threads}")74 75 def optimize_model(76 self,77 model: nn.Module,78 example_input: Optional[torch.Tensor] = None,79 use_amp: bool = True,80 use_channels_last: bool = True81 ) -> nn.Module:82 """83 Optimize PyTorch model for inference.84 85 Parameters:86 ----------87 model : nn.Module88 PyTorch model to optimize89 example_input : Optional[torch.Tensor]90 Example input for optimization (required for some modes)91 use_amp : bool92 Use automatic mixed precision (AMP)93 use_channels_last : bool94 Use channels-last memory format95 96 Returns:97 -------98 nn.Module99 Optimized model100 """101 print(f"Optimizing model with mode: {self.optimize_mode}")102 103 self.model = model.eval().to(self.device)104 self.use_amp = use_amp105 106 # Disable gradients for all parameters (inference only)107 for param in self.model.parameters():108 param.requires_grad = False109 110 # Apply memory format optimization (default: channels_last for CUDA)111 # Note: Audio models use 3D tensors, so channels_last is applied only where beneficial112 if use_channels_last and self.device.startswith('cuda'):113 print(" Using channels-last optimization")114 # Only apply to model if it has 4D conv layers, otherwise skip silently115 try:116 with warnings.catch_warnings():117 warnings.simplefilter("ignore")118 self.model = self.model.to(memory_format=torch.channels_last)119 except Exception:120 pass # Silently skip for models that don't support channels_last121 122 # Set model to inference mode123 torch.set_grad_enabled(False)124 125 # Apply optimization based on mode126 if self.optimize_mode == 'compile':127 self.compiled_model = self._compile_model(self.model)128 elif self.optimize_mode == 'jit':129 if example_input is None:130 raise ValueError("example_input required for JIT optimization")131 self.compiled_model = self._jit_trace_model(self.model, example_input)132 elif self.optimize_mode == 'channels_last':133 self.compiled_model = self.model134 else:135 print(" Using default optimization")136 self.compiled_model = self.model137 138 # Apply fusion optimizations if possible139 try:140 if hasattr(torch.nn.utils, 'fusion'):141 self.compiled_model = torch.nn.utils.fusion.fuse_conv_bn_eval(self.compiled_model)142 print(" Conv-BN fusion applied")143 except:144 pass145 146 print("Optimization complete")147 return self.compiled_model148 149 def _compile_model(self, model: nn.Module) -> nn.Module:150 """151 Compile model using torch.compile (PyTorch 2.0+) with ULTRA optimization.152 153 Parameters:154 ----------155 model : nn.Module156 Model to compile157 158 Returns:159 -------160 nn.Module161 Compiled model162 """163 try:164 if hasattr(torch, 'compile'):165 print(" Compiling model with torch.compile")166 # Try max-autotune for best performance167 try:168 compiled = torch.compile(model, mode='max-autotune', fullgraph=True)169 print(" Using max-autotune mode")170 return compiled171 except:172 # Fallback to reduce-overhead173 compiled = torch.compile(model, mode='reduce-overhead')174 print(" Using reduce-overhead mode")175 return compiled176 else:177 print(" torch.compile not available (requires PyTorch 2.0+)")178 return model179 except Exception as e:180 print(f" Compilation failed: {e}")181 return model182 183 def _jit_trace_model(self, model: nn.Module, example_input: torch.Tensor) -> nn.Module:184 """185 Trace model using TorchScript JIT.186 187 Parameters:188 ----------189 model : nn.Module190 Model to trace191 example_input : torch.Tensor192 Example input for tracing193 194 Returns:195 -------196 nn.Module197 Traced model198 """199 try:200 print(" → Tracing model with TorchScript JIT")201 with torch.no_grad():202 traced = torch.jit.trace(model, example_input)203 traced = torch.jit.optimize_for_inference(traced)204 return traced205 except Exception as e:206 print(f" JIT tracing failed: {e}")207 return model208 209 def save_optimized_model(self, save_path: str):210 """211 Save optimized model to file.212 213 Parameters:214 ----------215 save_path : str216 Path to save the model217 """218 if self.compiled_model is None:219 raise RuntimeError("No model has been optimized yet")220 221 try:222 # Save based on optimization mode223 if self.optimize_mode == 'jit':224 torch.jit.save(self.compiled_model, save_path)225 else:226 torch.save(self.compiled_model.state_dict(), save_path)227 print(f"✓ Model saved to: {save_path}")228 except Exception as e:229 print(f"✗ Failed to save model: {e}")230 231 def load_optimized_model(self, load_path: str, model_template: nn.Module) -> nn.Module:232 """233 Load optimized model from file.234 235 Parameters:236 ----------237 load_path : str238 Path to the saved model239 model_template : nn.Module240 Model template for loading state dict241 242 Returns:243 -------244 nn.Module245 Loaded model246 """247 try:248 if self.optimize_mode == 'jit':249 self.compiled_model = torch.jit.load(load_path, map_location=self.device)250 else:251 model_template.load_state_dict(torch.load(load_path, map_location=self.device, weights_only=False))252 self.compiled_model = model_template.eval()253 254 print(f"✓ Model loaded from: {load_path}")255 return self.compiled_model256 except (pickle.UnpicklingError, RuntimeError, EOFError) as e:257 error_details = f"""258CHECKPOINT FILE CORRUPTED259 260Error: {str(e)}261 262The checkpoint file appears to be corrupted or was not downloaded correctly.263File: {load_path}264 265Common causes:266 - File is an HTML page (wrong download URL, e.g., HuggingFace /blob/ instead of /resolve/)267 - Incomplete or interrupted download268 - Network issues during download269 - File system corruption270 271Solution:272 1. Delete the corrupted checkpoint file:273 {load_path}274 2. Re-run the application - it will automatically re-download the model275 3. If the problem persists, check that your model URL uses /resolve/ not /blob/276 Example: https://huggingface.co/user/repo/resolve/main/model.ckpt277"""278 print(error_details)279 raise280 except Exception as e:281 print(f"✗ Failed to load model: {e}")282 raise283 284 def __call__(self, x: torch.Tensor) -> torch.Tensor:285 """286 Run inference with optimized model.287 288 Parameters:289 ----------290 x : torch.Tensor291 Input tensor292 293 Returns:294 -------295 torch.Tensor296 Model output297 """298 if self.compiled_model is None:299 raise RuntimeError("No model has been optimized yet")300 301 # Apply memory format if needed (only for 4D tensors - images)302 # Audio models typically use 3D tensors, so we silently skip channels_last for them303 if self.optimize_mode == 'channels_last' and x.dim() == 4:304 x = x.to(memory_format=torch.channels_last)305 306 # Run inference with AMP if enabled307 try:308 if self.use_amp and self.device.startswith('cuda'):309 with torch.cuda.amp.autocast():310 with torch.no_grad():311 return self.compiled_model(x)312 else:313 with torch.no_grad():314 return self.compiled_model(x)315 except Exception as e:316 # Fallback to non-compiled model if torch.compile fails at runtime317 # This can happen with rotary embeddings that mutate class variables318 if self.optimize_mode == 'compile' and self.model is not None:319 print(f" ⚠️ torch.compile runtime error: {type(e).__name__}")320 print(f" 🔄 Falling back to non-compiled model...")321 self.compiled_model = self.model322 self.optimize_mode = 'fallback'323 # Retry with non-compiled model324 if self.use_amp and self.device.startswith('cuda'):325 with torch.cuda.amp.autocast():326 with torch.no_grad():327 return self.compiled_model(x)328 else:329 with torch.no_grad():330 return self.compiled_model(x)331 else:332 raise333 334 335class PyTorchOptimizer:336 """337 Helper class for various PyTorch optimization techniques.338 """339 340 @staticmethod341 def enable_cudnn_benchmark():342 """Enable cuDNN benchmark mode."""343 if torch.cuda.is_available():344 torch.backends.cudnn.benchmark = True345 torch.backends.cudnn.deterministic = False346 print("cuDNN benchmark enabled")347 348 @staticmethod349 def enable_cudnn_deterministic():350 """Enable cuDNN deterministic mode for reproducible results."""351 if torch.cuda.is_available():352 torch.backends.cudnn.deterministic = True353 torch.backends.cudnn.benchmark = False354 print("✓ cuDNN deterministic mode enabled")355 356 @staticmethod357 def enable_tf32():358 """Enable TF32 for Ampere GPUs (RTX 30xx+)."""359 if torch.cuda.is_available():360 torch.backends.cuda.matmul.allow_tf32 = True361 torch.backends.cudnn.allow_tf32 = True362 # Also enable for float32 matmul precision363 torch.set_float32_matmul_precision('high') # or 'highest' for max speed364 print("TF32 enabled")365 366 @staticmethod367 def set_num_threads(num_threads: int):368 """Set number of threads for CPU inference."""369 torch.set_num_threads(num_threads)370 print(f"✓ Number of threads set to: {num_threads}")371 372 @staticmethod373 def optimize_for_inference(model: nn.Module) -> nn.Module:374 """375 Apply ULTRA optimization for inference.376 377 Parameters:378 ----------379 model : nn.Module380 Model to optimize381 382 Returns:383 -------384 nn.Module385 ULTRA-optimized model386 """387 model.eval()388 torch.set_grad_enabled(False)389 390 # Disable gradient computation for all parameters391 for param in model.parameters():392 param.requires_grad = False393 394 # Fuse operations if possible395 try:396 # Try to fuse batch norm397 model = torch.quantization.fuse_modules(model, inplace=True)398 print("Batch norm fused")399 except:400 pass401 402 try:403 # Try to fuse conv-bn if available404 if hasattr(torch.nn.utils, 'fusion'):405 model = torch.nn.utils.fusion.fuse_conv_bn_eval(model)406 print("Conv-BN fused")407 except:408 pass409 410 return model411 412 413def benchmark_pytorch_optimizations(414 model: nn.Module,415 input_shape: Tuple[int, ...],416 device: str = 'cuda:0',417 num_iterations: int = 100,418 warmup_iterations: int = 10419) -> Dict[str, float]:420 """421 Benchmark different PyTorch optimization techniques.422 423 Parameters:424 ----------425 model : nn.Module426 Model to benchmark427 input_shape : Tuple[int, ...]428 Input tensor shape429 device : str430 Device to use431 num_iterations : int432 Number of benchmark iterations433 warmup_iterations : int434 Number of warmup iterations435 436 Returns:437 -------438 Dict[str, float]439 Benchmark results with average inference times440 """441 results = {}442 dummy_input = torch.randn(*input_shape).to(device)443 444 optimization_modes = ['default', 'compile', 'channels_last']445 446 for mode in optimization_modes:447 print(f"\n{'='*60}")448 print(f"Benchmarking: {mode}")449 print('='*60)450 451 try:452 backend = PyTorchBackend(device=device, optimize_mode=mode)453 454 # Optimize model455 if mode == 'compile':456 optimized_model = backend.optimize_model(model, use_amp=True)457 else:458 optimized_model = backend.optimize_model(459 model, 460 example_input=dummy_input,461 use_amp=True,462 use_channels_last=(mode == 'channels_last')463 )464 465 # Warmup466 for _ in range(warmup_iterations):467 _ = backend(dummy_input)468 469 # Benchmark470 if device.startswith('cuda'):471 torch.cuda.synchronize()472 473 start = time.time()474 for _ in range(num_iterations):475 _ = backend(dummy_input)476 477 if device.startswith('cuda'):478 torch.cuda.synchronize()479 480 elapsed = (time.time() - start) / num_iterations481 results[mode] = elapsed * 1000 # Convert to ms482 483 print(f" Average time: {results[mode]:.2f} ms")484 485 except Exception as e:486 print(f" Failed: {e}")487 results[mode] = None488 489 return results490 491 492def create_inference_session(493 model: nn.Module,494 device: str = 'cuda:0',495 optimize_mode: str = 'default',496 enable_amp: bool = True,497 enable_tf32: bool = True,498 enable_cudnn_benchmark: bool = True499) -> PyTorchBackend:500 """501 Create an optimized inference session.502 503 Parameters:504 ----------505 model : nn.Module506 Model to use for inference507 device : str508 Device to use509 optimize_mode : str510 Optimization mode511 enable_amp : bool512 Enable automatic mixed precision513 enable_tf32 : bool514 Enable TF32 (for Ampere GPUs)515 enable_cudnn_benchmark : bool516 Enable cuDNN benchmark517 518 Returns:519 -------520 PyTorchBackend521 Configured inference session522 """523 # Apply global optimizations524 optimizer = PyTorchOptimizer()525 526 if enable_cudnn_benchmark:527 optimizer.enable_cudnn_benchmark()528 529 if enable_tf32 and device.startswith('cuda'):530 optimizer.enable_tf32()531 532 # Create backend533 backend = PyTorchBackend(device=device, optimize_mode=optimize_mode)534 backend.optimize_model(model, use_amp=enable_amp)535 536 return backend537 538 539def convert_model_to_onnx(540 model: nn.Module,541 input_shape: Tuple[int, ...],542 output_path: str,543 opset_version: int = 14544):545 """546 Convert PyTorch model to ONNX format.547 548 Parameters:549 ----------550 model : nn.Module551 Model to convert552 input_shape : Tuple[int, ...]553 Input tensor shape554 output_path : str555 Path to save ONNX model556 opset_version : int557 ONNX opset version558 """559 try:560 import onnx561 562 model.eval()563 dummy_input = torch.randn(*input_shape)564 565 print(f"Converting model to ONNX (opset {opset_version})...")566 torch.onnx.export(567 model,568 dummy_input,569 output_path,570 export_params=True,571 opset_version=opset_version,572 do_constant_folding=True,573 input_names=['input'],574 output_names=['output'],575 dynamic_axes={576 'input': {0: 'batch_size'},577 'output': {0: 'batch_size'}578 }579 )580 581 # Verify ONNX model582 onnx_model = onnx.load(output_path)583 onnx.checker.check_model(onnx_model)584 585 print(f"✓ ONNX model saved to: {output_path}")586 587 except ImportError:588 print("✗ ONNX not available. Install with: pip install onnx")589 except Exception as e:590 print(f"✗ ONNX conversion failed: {e}")591 592 593def get_model_info(model: nn.Module) -> Dict[str, Any]:594 """595 Get information about a PyTorch model.596 597 Parameters:598 ----------599 model : nn.Module600 Model to analyze601 602 Returns:603 -------604 Dict[str, Any]605 Model information606 """607 total_params = sum(p.numel() for p in model.parameters())608 trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)609 610 # Estimate model size611 param_size = sum(p.nelement() * p.element_size() for p in model.parameters())612 buffer_size = sum(b.nelement() * b.element_size() for b in model.buffers())613 size_mb = (param_size + buffer_size) / (1024 ** 2)614 615 return {616 'total_parameters': total_params,617 'trainable_parameters': trainable_params,618 'model_size_mb': size_mb,619 'device': next(model.parameters()).device,620 'dtype': next(model.parameters()).dtype621 }622