@diottodev/runic 1.0.0 → 1.1.0

package/engineering-team/skills/senior-computer-vision/scripts/inference_optimizer.py

@@ -0,0 +1,557 @@
+#!/usr/bin/env python3
+"""
+Inference Optimizer
+
+Analyzes and benchmarks vision models, and provides optimization recommendations.
+Supports PyTorch, ONNX, and TensorRT models.
+
+Usage:
+    python inference_optimizer.py model.pt --benchmark
+    python inference_optimizer.py model.pt --export onnx --output model.onnx
+    python inference_optimizer.py model.onnx --analyze
+"""
+
+import os
+import sys
+import json
+import argparse
+import logging
+import time
+from pathlib import Path
+from typing import Dict, List, Optional, Any, Tuple
+from datetime import datetime
+import statistics
+
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+
+
+# Model format signatures
+MODEL_FORMATS = {
+    '.pt': 'pytorch',
+    '.pth': 'pytorch',
+    '.onnx': 'onnx',
+    '.engine': 'tensorrt',
+    '.trt': 'tensorrt',
+    '.xml': 'openvino',
+    '.mlpackage': 'coreml',
+    '.mlmodel': 'coreml',
+}
+
+# Optimization recommendations
+OPTIMIZATION_PATHS = {
+    ('pytorch', 'gpu'): ['onnx', 'tensorrt_fp16'],
+    ('pytorch', 'cpu'): ['onnx', 'onnxruntime'],
+    ('pytorch', 'edge'): ['onnx', 'tensorrt_int8'],
+    ('pytorch', 'mobile'): ['onnx', 'tflite'],
+    ('pytorch', 'apple'): ['coreml'],
+    ('pytorch', 'intel'): ['onnx', 'openvino'],
+    ('onnx', 'gpu'): ['tensorrt_fp16'],
+    ('onnx', 'cpu'): ['onnxruntime'],
+}
+
+
+class InferenceOptimizer:
+    """Analyzes and optimizes vision model inference."""
+
+    def __init__(self, model_path: str):
+        self.model_path = Path(model_path)
+        self.model_format = self._detect_format()
+        self.model_info = {}
+        self.benchmark_results = {}
+
+    def _detect_format(self) -> str:
+        """Detect model format from file extension."""
+        suffix = self.model_path.suffix.lower()
+        if suffix in MODEL_FORMATS:
+            return MODEL_FORMATS[suffix]
+        raise ValueError(f"Unknown model format: {suffix}")
+
+    def analyze_model(self) -> Dict[str, Any]:
+        """Analyze model structure and size."""
+        logger.info(f"Analyzing model: {self.model_path}")
+
+        analysis = {
+            'path': str(self.model_path),
+            'format': self.model_format,
+            'file_size_mb': self.model_path.stat().st_size / 1024 / 1024,
+            'parameters': None,
+            'layers': [],
+            'input_shape': None,
+            'output_shape': None,
+            'ops_count': None,
+        }
+
+        if self.model_format == 'onnx':
+            analysis.update(self._analyze_onnx())
+        elif self.model_format == 'pytorch':
+            analysis.update(self._analyze_pytorch())
+
+        self.model_info = analysis
+        return analysis
+
+    def _analyze_onnx(self) -> Dict[str, Any]:
+        """Analyze ONNX model."""
+        try:
+            import onnx
+            model = onnx.load(str(self.model_path))
+            onnx.checker.check_model(model)
+
+            # Count parameters
+            total_params = 0
+            for initializer in model.graph.initializer:
+                param_count = 1
+                for dim in initializer.dims:
+                    param_count *= dim
+                total_params += param_count
+
+            # Get input/output shapes
+            inputs = []
+            for inp in model.graph.input:
+                shape = [d.dim_value if d.dim_value else -1
+                        for d in inp.type.tensor_type.shape.dim]
+                inputs.append({'name': inp.name, 'shape': shape})
+
+            outputs = []
+            for out in model.graph.output:
+                shape = [d.dim_value if d.dim_value else -1
+                        for d in out.type.tensor_type.shape.dim]
+                outputs.append({'name': out.name, 'shape': shape})
+
+            # Count operators
+            op_counts = {}
+            for node in model.graph.node:
+                op_type = node.op_type
+                op_counts[op_type] = op_counts.get(op_type, 0) + 1
+
+            return {
+                'parameters': total_params,
+                'inputs': inputs,
+                'outputs': outputs,
+                'operator_counts': op_counts,
+                'num_nodes': len(model.graph.node),
+                'opset_version': model.opset_import[0].version if model.opset_import else None,
+            }
+
+        except ImportError:
+            logger.warning("onnx package not installed, skipping detailed analysis")
+            return {}
+        except Exception as e:
+            logger.error(f"Error analyzing ONNX model: {e}")
+            return {'error': str(e)}
+
+    def _analyze_pytorch(self) -> Dict[str, Any]:
+        """Analyze PyTorch model."""
+        try:
+            import torch
+
+            # Try to load as checkpoint
+            checkpoint = torch.load(str(self.model_path), map_location='cpu')
+
+            # Handle different checkpoint formats
+            if isinstance(checkpoint, dict):
+                if 'model' in checkpoint:
+                    state_dict = checkpoint['model']
+                elif 'state_dict' in checkpoint:
+                    state_dict = checkpoint['state_dict']
+                else:
+                    state_dict = checkpoint
+            else:
+                # Assume it's the model itself
+                if hasattr(checkpoint, 'state_dict'):
+                    state_dict = checkpoint.state_dict()
+                else:
+                    return {'error': 'Could not extract state dict'}
+
+            # Count parameters
+            total_params = 0
+            layer_info = []
+            for name, param in state_dict.items():
+                if hasattr(param, 'numel'):
+                    param_count = param.numel()
+                    total_params += param_count
+                    layer_info.append({
+                        'name': name,
+                        'shape': list(param.shape),
+                        'params': param_count,
+                        'dtype': str(param.dtype)
+                    })
+
+            return {
+                'parameters': total_params,
+                'layers': layer_info[:20],  # First 20 layers
+                'num_layers': len(layer_info),
+            }
+
+        except ImportError:
+            logger.warning("torch package not installed, skipping detailed analysis")
+            return {}
+        except Exception as e:
+            logger.error(f"Error analyzing PyTorch model: {e}")
+            return {'error': str(e)}
+
+    def benchmark(self, input_size: Tuple[int, int] = (640, 640),
+                  batch_sizes: List[int] = None,
+                  num_iterations: int = 100,
+                  warmup: int = 10) -> Dict[str, Any]:
+        """Benchmark model inference speed."""
+        if batch_sizes is None:
+            batch_sizes = [1, 4, 8, 16]
+
+        logger.info(f"Benchmarking model with input size {input_size}")
+
+        results = {
+            'input_size': input_size,
+            'num_iterations': num_iterations,
+            'warmup_iterations': warmup,
+            'batch_results': [],
+            'device': 'cpu',
+        }
+
+        try:
+            if self.model_format == 'onnx':
+                results.update(self._benchmark_onnx(input_size, batch_sizes,
+                                                    num_iterations, warmup))
+            elif self.model_format == 'pytorch':
+                results.update(self._benchmark_pytorch(input_size, batch_sizes,
+                                                       num_iterations, warmup))
+            else:
+                results['error'] = f"Benchmarking not supported for {self.model_format}"
+
+        except Exception as e:
+            results['error'] = str(e)
+            logger.error(f"Benchmark failed: {e}")
+
+        self.benchmark_results = results
+        return results
+
+    def _benchmark_onnx(self, input_size: Tuple[int, int],
+                        batch_sizes: List[int],
+                        num_iterations: int, warmup: int) -> Dict[str, Any]:
+        """Benchmark ONNX model."""
+        import numpy as np
+
+        try:
+            import onnxruntime as ort
+
+            # Try GPU first, fall back to CPU
+            providers = ['CPUExecutionProvider']
+            try:
+                if 'CUDAExecutionProvider' in ort.get_available_providers():
+                    providers = ['CUDAExecutionProvider'] + providers
+            except:
+                pass
+
+            session = ort.InferenceSession(str(self.model_path), providers=providers)
+            input_name = session.get_inputs()[0].name
+            device = 'cuda' if 'CUDA' in session.get_providers()[0] else 'cpu'
+
+            results = {'device': device, 'provider': session.get_providers()[0]}
+            batch_results = []
+
+            for batch_size in batch_sizes:
+                # Create dummy input
+                dummy = np.random.randn(batch_size, 3, *input_size).astype(np.float32)
+
+                # Warmup
+                for _ in range(warmup):
+                    session.run(None, {input_name: dummy})
+
+                # Benchmark
+                latencies = []
+                for _ in range(num_iterations):
+                    start = time.perf_counter()
+                    session.run(None, {input_name: dummy})
+                    latencies.append((time.perf_counter() - start) * 1000)
+
+                batch_result = {
+                    'batch_size': batch_size,
+                    'mean_latency_ms': statistics.mean(latencies),
+                    'std_latency_ms': statistics.stdev(latencies) if len(latencies) > 1 else 0,
+                    'min_latency_ms': min(latencies),
+                    'max_latency_ms': max(latencies),
+                    'p50_latency_ms': sorted(latencies)[len(latencies) // 2],
+                    'p95_latency_ms': sorted(latencies)[int(len(latencies) * 0.95)],
+                    'p99_latency_ms': sorted(latencies)[int(len(latencies) * 0.99)],
+                    'throughput_fps': batch_size * 1000 / statistics.mean(latencies),
+                }
+                batch_results.append(batch_result)
+
+                logger.info(f"Batch {batch_size}: {batch_result['mean_latency_ms']:.2f}ms, "
+                           f"{batch_result['throughput_fps']:.1f} FPS")
+
+            results['batch_results'] = batch_results
+            return results
+
+        except ImportError:
+            return {'error': 'onnxruntime not installed'}
+
+    def _benchmark_pytorch(self, input_size: Tuple[int, int],
+                          batch_sizes: List[int],
+                          num_iterations: int, warmup: int) -> Dict[str, Any]:
+        """Benchmark PyTorch model."""
+        try:
+            import torch
+            import numpy as np
+
+            # Load model
+            device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+            checkpoint = torch.load(str(self.model_path), map_location=device)
+
+            # Handle different checkpoint formats
+            if isinstance(checkpoint, dict) and 'model' in checkpoint:
+                model = checkpoint['model']
+            elif hasattr(checkpoint, 'forward'):
+                model = checkpoint
+            else:
+                return {'error': 'Could not load model for benchmarking'}
+
+            model.to(device)
+            model.train(False)
+
+            results = {'device': str(device)}
+            batch_results = []
+
+            with torch.no_grad():
+                for batch_size in batch_sizes:
+                    dummy = torch.randn(batch_size, 3, *input_size, device=device)
+
+                    # Warmup
+                    for _ in range(warmup):
+                        _ = model(dummy)
+                    if device.type == 'cuda':
+                        torch.cuda.synchronize()
+
+                    # Benchmark
+                    latencies = []
+                    for _ in range(num_iterations):
+                        if device.type == 'cuda':
+                            torch.cuda.synchronize()
+                        start = time.perf_counter()
+                        _ = model(dummy)
+                        if device.type == 'cuda':
+                            torch.cuda.synchronize()
+                        latencies.append((time.perf_counter() - start) * 1000)
+
+                    batch_result = {
+                        'batch_size': batch_size,
+                        'mean_latency_ms': statistics.mean(latencies),
+                        'std_latency_ms': statistics.stdev(latencies) if len(latencies) > 1 else 0,
+                        'min_latency_ms': min(latencies),
+                        'max_latency_ms': max(latencies),
+                        'throughput_fps': batch_size * 1000 / statistics.mean(latencies),
+                    }
+                    batch_results.append(batch_result)
+
+                    logger.info(f"Batch {batch_size}: {batch_result['mean_latency_ms']:.2f}ms, "
+                               f"{batch_result['throughput_fps']:.1f} FPS")
+
+            results['batch_results'] = batch_results
+            return results
+
+        except ImportError:
+            return {'error': 'torch not installed'}
+        except Exception as e:
+            return {'error': str(e)}
+
+    def get_optimization_recommendations(self, target: str = 'gpu') -> List[Dict[str, Any]]:
+        """Get optimization recommendations for target platform."""
+        recommendations = []
+
+        key = (self.model_format, target)
+        if key in OPTIMIZATION_PATHS:
+            path = OPTIMIZATION_PATHS[key]
+            for step in path:
+                rec = {
+                    'step': step,
+                    'description': self._get_step_description(step),
+                    'expected_speedup': self._get_expected_speedup(step),
+                    'command': self._get_step_command(step),
+                }
+                recommendations.append(rec)
+
+        # Add general recommendations
+        if self.model_info:
+            params = self.model_info.get('parameters', 0)
+            if params and params > 50_000_000:
+                recommendations.append({
+                    'step': 'pruning',
+                    'description': f'Model has {params/1e6:.1f}M parameters. '
+                                 'Consider structured pruning to reduce size.',
+                    'expected_speedup': '1.5-2x',
+                })
+
+            file_size = self.model_info.get('file_size_mb', 0)
+            if file_size > 100:
+                recommendations.append({
+                    'step': 'quantization',
+                    'description': f'Model size is {file_size:.1f}MB. '
+                                 'INT8 quantization can reduce by 75%.',
+                    'expected_speedup': '2-4x',
+                })
+
+        return recommendations
+
+    def _get_step_description(self, step: str) -> str:
+        """Get description for optimization step."""
+        descriptions = {
+            'onnx': 'Export to ONNX format for framework-agnostic deployment',
+            'tensorrt_fp16': 'Convert to TensorRT with FP16 precision for NVIDIA GPUs',
+            'tensorrt_int8': 'Convert to TensorRT with INT8 quantization for edge devices',
+            'onnxruntime': 'Use ONNX Runtime for optimized CPU/GPU inference',
+            'openvino': 'Convert to OpenVINO for Intel CPU/GPU optimization',
+            'coreml': 'Convert to CoreML for Apple Silicon acceleration',
+            'tflite': 'Convert to TensorFlow Lite for mobile deployment',
+        }
+        return descriptions.get(step, step)
+
+    def _get_expected_speedup(self, step: str) -> str:
+        """Get expected speedup for optimization step."""
+        speedups = {
+            'onnx': '1-1.5x',
+            'tensorrt_fp16': '2-4x',
+            'tensorrt_int8': '3-6x',
+            'onnxruntime': '1.2-2x',
+            'openvino': '1.5-3x',
+            'coreml': '2-5x (on Apple Silicon)',
+            'tflite': '1-2x',
+        }
+        return speedups.get(step, 'varies')
+
+    def _get_step_command(self, step: str) -> str:
+        """Get command for optimization step."""
+        model_name = self.model_path.stem
+        commands = {
+            'onnx': f'yolo export model={model_name}.pt format=onnx',
+            'tensorrt_fp16': f'trtexec --onnx={model_name}.onnx --saveEngine={model_name}.engine --fp16',
+            'tensorrt_int8': f'trtexec --onnx={model_name}.onnx --saveEngine={model_name}.engine --int8',
+            'onnxruntime': f'pip install onnxruntime-gpu',
+            'openvino': f'mo --input_model {model_name}.onnx --output_dir openvino/',
+            'coreml': f'yolo export model={model_name}.pt format=coreml',
+        }
+        return commands.get(step, '')
+
+    def print_summary(self):
+        """Print analysis and benchmark summary."""
+        print("\n" + "=" * 70)
+        print("MODEL ANALYSIS SUMMARY")
+        print("=" * 70)
+
+        if self.model_info:
+            print(f"Path:        {self.model_info.get('path', 'N/A')}")
+            print(f"Format:      {self.model_info.get('format', 'N/A')}")
+            print(f"File Size:   {self.model_info.get('file_size_mb', 0):.2f} MB")
+
+            params = self.model_info.get('parameters')
+            if params:
+                print(f"Parameters:  {params:,} ({params/1e6:.2f}M)")
+
+            if 'num_nodes' in self.model_info:
+                print(f"Nodes:       {self.model_info['num_nodes']}")
+
+        if self.benchmark_results and 'batch_results' in self.benchmark_results:
+            print("\n" + "-" * 70)
+            print("BENCHMARK RESULTS")
+            print("-" * 70)
+            print(f"Device:      {self.benchmark_results.get('device', 'N/A')}")
+            print(f"Input Size:  {self.benchmark_results.get('input_size', 'N/A')}")
+            print()
+            print(f"{'Batch':<8} {'Latency (ms)':<15} {'Throughput (FPS)':<18} {'P99 (ms)':<12}")
+            print("-" * 55)
+
+            for result in self.benchmark_results['batch_results']:
+                print(f"{result['batch_size']:<8} "
+                      f"{result['mean_latency_ms']:<15.2f} "
+                      f"{result['throughput_fps']:<18.1f} "
+                      f"{result.get('p99_latency_ms', 0):<12.2f}")
+
+        print("=" * 70 + "\n")
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Analyze and optimize vision model inference"
+    )
+    parser.add_argument('model_path', help='Path to model file')
+    parser.add_argument('--analyze', action='store_true',
+                       help='Analyze model structure')
+    parser.add_argument('--benchmark', action='store_true',
+                       help='Benchmark inference speed')
+    parser.add_argument('--input-size', type=int, nargs=2, default=[640, 640],
+                       metavar=('H', 'W'), help='Input image size')
+    parser.add_argument('--batch-sizes', type=int, nargs='+', default=[1, 4, 8],
+                       help='Batch sizes to benchmark')
+    parser.add_argument('--iterations', type=int, default=100,
+                       help='Number of benchmark iterations')
+    parser.add_argument('--warmup', type=int, default=10,
+                       help='Number of warmup iterations')
+    parser.add_argument('--target', choices=['gpu', 'cpu', 'edge', 'mobile', 'apple', 'intel'],
+                       default='gpu', help='Target deployment platform')
+    parser.add_argument('--recommend', action='store_true',
+                       help='Show optimization recommendations')
+    parser.add_argument('--json', action='store_true',
+                       help='Output as JSON')
+    parser.add_argument('--output', '-o', help='Output file path')
+
+    args = parser.parse_args()
+
+    if not Path(args.model_path).exists():
+        logger.error(f"Model not found: {args.model_path}")
+        sys.exit(1)
+
+    try:
+        optimizer = InferenceOptimizer(args.model_path)
+    except ValueError as e:
+        logger.error(str(e))
+        sys.exit(1)
+
+    results = {}
+
+    # Analyze model
+    if args.analyze or not (args.benchmark or args.recommend):
+        results['analysis'] = optimizer.analyze_model()
+
+    # Benchmark
+    if args.benchmark:
+        results['benchmark'] = optimizer.benchmark(
+            input_size=tuple(args.input_size),
+            batch_sizes=args.batch_sizes,
+            num_iterations=args.iterations,
+            warmup=args.warmup
+        )
+
+    # Recommendations
+    if args.recommend:
+        if not optimizer.model_info:
+            optimizer.analyze_model()
+        results['recommendations'] = optimizer.get_optimization_recommendations(args.target)
+
+    # Output
+    if args.json:
+        print(json.dumps(results, indent=2, default=str))
+    else:
+        optimizer.print_summary()
+
+        if args.recommend and 'recommendations' in results:
+            print("OPTIMIZATION RECOMMENDATIONS")
+            print("-" * 70)
+            for i, rec in enumerate(results['recommendations'], 1):
+                print(f"\n{i}. {rec['step'].upper()}")
+                print(f"   {rec['description']}")
+                print(f"   Expected speedup: {rec['expected_speedup']}")
+                if rec.get('command'):
+                    print(f"   Command: {rec['command']}")
+            print()
+
+    # Save to file
+    if args.output:
+        with open(args.output, 'w') as f:
+            json.dump(results, f, indent=2, default=str)
+        logger.info(f"Results saved to {args.output}")
+
+
+if __name__ == '__main__':
+    main()