Model Quantization Deep Dive

import torch
import time
import logging
from typing import Dict, Any
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel

class QuantizedModelService:
    def __init__(self, model_path: str, precision: str = "int8"):
        self.precision = precision
        self.model = self._load_quantized_model(model_path)
        self.warmup_model()
        
        # Performance metrics
        self.inference_times = []
        self.memory_usage = self._get_memory_usage()
        
    def _load_quantized_model(self, model_path: str):
        """Load quantized model with error handling"""
        try:
            if self.precision == "int8":
                model = torch.jit.load(model_path)
                model.eval()
                return model
            else:
                raise ValueError(f"Unsupported precision: {self.precision}")
        except Exception as e:
            logging.error(f"Failed to load quantized model: {e}")
            raise
    
    def warmup_model(self, num_warmup: int = 5):
        """Warmup model for consistent performance"""
        dummy_input = torch.randn(1, 3, 224, 224)  # Adjust for your input size
        
        with torch.no_grad():
            for _ in range(num_warmup):
                _ = self.model(dummy_input)
        
        logging.info(f"Model warmed up with {num_warmup} iterations")
    
    def predict(self, input_data: torch.Tensor) -> Dict[str, Any]:
        """Run inference with performance monitoring"""
        start_time = time.time()
        
        try:
            with torch.no_grad():
                output = self.model(input_data)
                
            inference_time = (time.time() - start_time) * 1000  # ms
            self.inference_times.append(inference_time)
            
            # Keep only last 1000 measurements
            if len(self.inference_times) > 1000:
                self.inference_times = self.inference_times[-1000:]
            
            return {
                "prediction": output.tolist(),
                "inference_time_ms": inference_time,
                "avg_inference_time_ms": sum(self.inference_times) / len(self.inference_times),
                "model_precision": self.precision
            }
            
        except Exception as e:
            logging.error(f"Inference failed: {e}")
            raise HTTPException(status_code=500, detail=f"Inference failed: {str(e)}")
    
    def _get_memory_usage(self) -> float:
        """Get current memory usage in MB"""
        if torch.cuda.is_available():
            return torch.cuda.memory_allocated() / 1024 / 1024
        else:
            import psutil
            return psutil.Process().memory_info().rss / 1024 / 1024
    
    def get_metrics(self) -> Dict[str, Any]:
        """Get service performance metrics"""
        return {
            "avg_inference_time_ms": sum(self.inference_times) / len(self.inference_times) if self.inference_times else 0,
            "total_predictions": len(self.inference_times),
            "memory_usage_mb": self._get_memory_usage(),
            "precision": self.precision,
            "throughput_qps": 1000 / (sum(self.inference_times) / len(self.inference_times)) if self.inference_times else 0
        }

# FastAPI application
app = FastAPI(title="Quantized Model Service")
quantized_service = QuantizedModelService("quantized_model.pt", precision="int8")

class PredictionRequest(BaseModel):
    data: list

@app.post("/predict")
async def predict(request: PredictionRequest):
    input_tensor = torch.tensor(request.data).float()
    return quantized_service.predict(input_tensor)

@app.get("/metrics")
async def get_metrics():
    return quantized_service.get_metrics()

@app.get("/health")
async def health_check():
    return {"status": "healthy", "precision": quantized_service.precision}