speed_comparison.py

import torch
import time
import numpy as np
import os
import json
import onnxruntime as ort
from PIL import Image
import torchvision.transforms as transforms
from models.bisenet import BiSeNet

def prepare_image_for_test():
    """Prepare a test image for inference"""
    # Create a dummy image or load from assets
    image = Image.new('RGB', (256, 256), color='red')
    
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
    ])
    
    image_tensor = transform(image)
    image_batch = image_tensor.unsqueeze(0)
    
    return image_batch

def test_pytorch_model(model_name):
    """Test PyTorch model speed"""
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    # Load model
    model = BiSeNet(num_classes=19, backbone_name=model_name)
    checkpoint_path = f'./weights/{model_name}.pt'
    
    if not os.path.exists(checkpoint_path):
        return None
    
    model.load_state_dict(torch.load(checkpoint_path, map_location='cpu'))
    model.to(device)
    model.eval()
    
    # Calculate model size
    temp_path = f"temp_{model_name}.pt"
    torch.save(model.state_dict(), temp_path)
    size_mb = os.path.getsize(temp_path) / (1024 * 1024)
    os.remove(temp_path)
    
    # Prepare input
    dummy_input = prepare_image_for_test().to(device)
    
    # Warmup
    with torch.no_grad():
        for _ in range(10):
            _ = model(dummy_input)
    
    # Test speed
    times = []
    with torch.no_grad():
        for _ in range(100):
            start_time = time.time()
            _ = model(dummy_input)
            end_time = time.time()
            times.append((end_time - start_time) * 1000)  # ms
    
    mean_time = np.mean(times)
    std_time = np.std(times)
    fps = 1000 / mean_time
    
    return {
        'speed_ms': round(mean_time, 2),
        'std_ms': round(std_time, 2),
        'fps': round(fps, 1),
        'size_mb': round(size_mb, 1)
    }

def test_onnx_model(model_name):
    """Test ONNX model speed"""
    onnx_path = f'./weights/{model_name}.onnx'
    
    if not os.path.exists(onnx_path):
        return None
    
    # Load ONNX model
    providers = ['CPUExecutionProvider']
    session = ort.InferenceSession(onnx_path, providers=providers)
    input_name = session.get_inputs()[0].name
    
    # Calculate model size
    size_mb = os.path.getsize(onnx_path) / (1024 * 1024)
    
    # Prepare input
    dummy_input = prepare_image_for_test().numpy()
    
    # Warmup
    for _ in range(10):
        _ = session.run(None, {input_name: dummy_input})
    
    # Test speed
    times = []
    for _ in range(100):
        start_time = time.time()
        _ = session.run(None, {input_name: dummy_input})
        end_time = time.time()
        times.append((end_time - start_time) * 1000)  # ms
    
    mean_time = np.mean(times)
    std_time = np.std(times)
    fps = 1000 / mean_time
    
    return {
        'speed_ms': round(mean_time, 2),
        'std_ms': round(std_time, 2),
        'fps': round(fps, 1),
        'size_mb': round(size_mb, 1)
    }

def compare_models():
    """Compare PyTorch vs ONNX performance"""
    
    models = ['resnet18', 'resnet34', 'efficientnet_b0', 'efficientnet_b1', 'efficientnet_b2']
    results = []
    
    print("🚀 COMPARING PYTORCH vs ONNX PERFORMANCE")
    print("="*80)
    
    for model_name in models:
        print(f"\n📊 Testing {model_name}...")
        
        # Test PyTorch
        pytorch_result = test_pytorch_model(model_name)
        if pytorch_result:
            print(f"   PyTorch: {pytorch_result['speed_ms']}ms ± {pytorch_result['std_ms']}ms")
        else:
            print(f"   PyTorch: ❌ Model not found")
        
        # Test ONNX
        onnx_result = test_onnx_model(model_name)
        if onnx_result:
            print(f"   ONNX:    {onnx_result['speed_ms']}ms ± {onnx_result['std_ms']}ms")
        else:
            print(f"   ONNX:    ❌ Model not found")
        
        # Calculate speedup
        if pytorch_result and onnx_result:
            speedup = pytorch_result['speed_ms'] / onnx_result['speed_ms']
            print(f"   Speedup: {speedup:.2f}x faster with ONNX")
        
        results.append({
            'model': model_name,
            'pytorch': pytorch_result,
            'onnx': onnx_result
        })
    
    # Print detailed comparison table
    print("\n" + "="*100)
    print("DETAILED COMPARISON TABLE")
    print("="*100)
    print(f"{'Model':<15} {'PyTorch(ms)':<15} {'ONNX(ms)':<12} {'Speedup':<10} {'PyTorch(MB)':<12} {'ONNX(MB)':<10}")
    print("-"*100)
    
    for result in results:
        model = result['model']
        pytorch = result['pytorch']
        onnx = result['onnx']
        
        if pytorch and onnx:
            speedup = pytorch['speed_ms'] / onnx['speed_ms']
            print(f"{model:<15} {pytorch['speed_ms']:<15} {onnx['speed_ms']:<12} "
                  f"{speedup:.2f}x{'':<6} {pytorch['size_mb']:<12} {onnx['size_mb']:<10}")
        else:
            print(f"{model:<15} {'N/A':<15} {'N/A':<12} {'N/A':<10} {'N/A':<12} {'N/A':<10}")
    
    print("="*100)
    
    # Summary statistics
    print("\n📈 SUMMARY STATISTICS:")
    
    valid_results = [r for r in results if r['pytorch'] and r['onnx']]
    if valid_results:
        speedups = [r['pytorch']['speed_ms'] / r['onnx']['speed_ms'] for r in valid_results]
        avg_speedup = np.mean(speedups)
        max_speedup = max(speedups)
        min_speedup = min(speedups)
        
        print(f"   Average speedup: {avg_speedup:.2f}x")
        print(f"   Max speedup: {max_speedup:.2f}x")
        print(f"   Min speedup: {min_speedup:.2f}x")
        
        # Best performing models
        fastest_pytorch = min(valid_results, key=lambda x: x['pytorch']['speed_ms'])
        fastest_onnx = min(valid_results, key=lambda x: x['onnx']['speed_ms'])
        
        print(f"\n🏆 BEST PERFORMERS:")
        print(f"   Fastest PyTorch: {fastest_pytorch['model']} ({fastest_pytorch['pytorch']['speed_ms']}ms)")
        print(f"   Fastest ONNX: {fastest_onnx['model']} ({fastest_onnx['onnx']['speed_ms']}ms)")
    
    # Save results
    with open('speed_comparison_results.json', 'w') as f:
        json.dump(results, f, indent=2)
    
    print(f"\n💾 Results saved to: speed_comparison_results.json")
    
    # Recommendations
    print("\n💡 RECOMMENDATIONS:")
    print("   • Use ONNX for production deployment (faster inference)")
    print("   • Use PyTorch for development and training")
    print("   • ONNX models are typically 1.5-3x faster than PyTorch")
    print("   • Consider model size vs speed trade-off for your use case")

if __name__ == "__main__":
    compare_models()