quantize_ptq.py

"""
PTQ Quantization Script.

Post-Training Quantization for edge deployment.
"""

import argparse
from pathlib import Path

import torch
from loguru import logger

from src.data.data_loader import create_dataloaders
from src.models.autoencoder import VibrationAutoencoder
from src.quantization.ptq import PostTrainingQuantizer
from src.utils.config_loader import load_config


def parse_args():
    """Parse command line arguments."""
    parser = argparse.ArgumentParser(description="Quantize model using PTQ")
    
    parser.add_argument(
        "--checkpoint",
        type=str,
        required=True,
        help="Path to FP32 model checkpoint",
    )
    
    parser.add_argument(
        "--config",
        type=str,
        default="config/config.yaml",
        help="Path to configuration file",
    )
    
    parser.add_argument(
        "--output",
        type=str,
        default="checkpoints/quantized_ptq.ckpt",
        help="Output path for quantized model",
    )
    
    parser.add_argument(
        "--calibration-batches",
        type=int,
        default=100,
        help="Number of batches for calibration",
    )
    
    parser.add_argument(
        "--backend",
        type=str,
        default="qnnpack",
        choices=["qnnpack", "fbgemm", "x86"],
        help="Quantization backend",
    )
    
    return parser.parse_args()


def main():
    """Main quantization function."""
    args = parse_args()
    
    # Load configuration
    logger.info(f"Loading configuration from {args.config}")
    config = load_config(args.config)
    
    # Load model
    logger.info(f"Loading model from {args.checkpoint}")
    checkpoint = torch.load(args.checkpoint, map_location="cpu")
    
    if "model" in checkpoint:
        model = checkpoint["model"]
    else:
        model = VibrationAutoencoder()
        model.load_state_dict(checkpoint["model_state_dict"])
    
    model.eval()
    model.cpu()
    
    # Load calibration data
    logger.info("Loading calibration data...")
    train_loader, _, _, _ = create_dataloaders(config)
    
    # Create quantizer
    quantizer = PostTrainingQuantizer(
        backend=args.backend,
        calibration_batches=args.calibration_batches,
    )
    
    # Prepare model
    logger.info("Preparing model for quantization...")
    prepared_model = quantizer.prepare_model(model)
    
    # Calibrate
    logger.info("Calibrating model...")
    calibrated_model = quantizer.calibrate(
        prepared_model,
        train_loader,
        device="cpu",
    )
    
    # Convert
    logger.info("Converting to INT8...")
    quantized_model = quantizer.convert(calibrated_model)
    
    # Get model sizes
    fp32_size = Path(args.checkpoint).stat().st_size / (1024 * 1024)
    
    # Save quantized model
    output_path = Path(args.output)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    
    torch.save({
        "model": quantized_model,
        "model_state_dict": quantized_model.state_dict(),
        "config": config,
        "quantization": {
            "method": "ptq",
            "backend": args.backend,
            "calibration_batches": args.calibration_batches,
        },
    }, output_path)
    
    int8_size = output_path.stat().st_size / (1024 * 1024)
    
    logger.info(f"\n{'='*50}")
    logger.info("QUANTIZATION COMPLETE")
    logger.info(f"{'='*50}")
    logger.info(f"FP32 Model Size: {fp32_size:.2f} MB")
    logger.info(f"INT8 Model Size: {int8_size:.2f} MB")
    logger.info(f"Compression Ratio: {fp32_size / int8_size:.2f}x")
    logger.info(f"Quantized model saved to: {output_path}")
    logger.info(f"{'='*50}\n")


if __name__ == "__main__":
    main()