quantized_clip_test.py

#!/usr/bin/env python3

import torch
import torch.nn as nn
from transformers import CLIPVisionModel, CLIPImageProcessor
from PIL import Image
import requests
from io import BytesIO
import argparse


def load_image(image_file):
    """Load image from file or URL"""
    if image_file.startswith('http://') or image_file.startswith('https://'):
        response = requests.get(image_file)
        image = Image.open(BytesIO(response.content)).convert('RGB')
    else:
        image = Image.open(image_file).convert('RGB')
    return image


class QuantizedCLIPVisionTower(nn.Module):
    """CLIP Vision Tower with INT8 quantization support"""

    def __init__(self, vision_tower_name="openai/clip-vit-large-patch14-336", device="cuda"):
        super().__init__()
        self.vision_tower_name = vision_tower_name
        self.device = device
        self.is_quantized = False

        # Load the original model
        print(f"Loading CLIP vision model: {vision_tower_name}")
        self.image_processor = CLIPImageProcessor.from_pretrained(vision_tower_name)
        self.vision_tower = CLIPVisionModel.from_pretrained(vision_tower_name)
        self.vision_tower.to(device)

        print(f"Original model memory usage: {self.get_model_size():.2f} MB")

    def get_model_size(self):
        """Calculate model size in MB"""
        param_size = 0
        for param in self.vision_tower.parameters():
            param_size += param.nelement() * param.element_size()
        buffer_size = 0
        for buffer in self.vision_tower.buffers():
            buffer_size += buffer.nelement() * buffer.element_size()
        return (param_size + buffer_size) / 1024 / 1024

    def quantize_to_int8(self):
        """Quantize the model to INT8 using PyTorch's built-in quantization"""
        print("Quantizing model to INT8...")

        # Set model to evaluation mode
        self.vision_tower.eval()

        # Prepare model for quantization
        self.vision_tower.qconfig = torch.quantization.get_default_qconfig('fbgemm')

        # Prepare the model (insert observers)
        torch.quantization.prepare(self.vision_tower, inplace=True)

        # Create some dummy data for calibration
        dummy_input = torch.randn(1, 3, 336, 336).to(self.device)

        # Calibrate the model with dummy data
        print("Calibrating model with dummy data...")
        with torch.no_grad():
            self.vision_tower(dummy_input)

        # Convert to quantized model
        torch.quantization.convert(self.vision_tower, inplace=True)

        self.is_quantized = True
        print(f"Quantized model memory usage: {self.get_model_size():.2f} MB")
        print("Model successfully quantized to INT8!")

    def quantize_to_int8_dynamic(self):
        """Apply dynamic quantization (simpler approach)"""
        print("Applying dynamic INT8 quantization...")

        # Apply dynamic quantization to linear layers
        self.vision_tower = torch.quantization.quantize_dynamic(
            self.vision_tower,
            {torch.nn.Linear},  # Quantize linear layers
            dtype=torch.qint8   # Use INT8
        )

        self.is_quantized = True
        print(f"Dynamically quantized model memory usage: {self.get_model_size():.2f} MB")
        print("Model successfully quantized with dynamic INT8!")

    def forward(self, images):
        """Forward pass through the vision tower"""
        with torch.no_grad():
            if isinstance(images, torch.Tensor):
                images = images.to(self.device)
                outputs = self.vision_tower(images, output_hidden_states=True)
                # Get features from the last hidden state, excluding CLS token
                image_features = outputs.hidden_states[-1][:, 1:]  # Remove CLS token
                return image_features
            else:
                raise ValueError("Expected torch.Tensor input")

    def process_image(self, image):
        """Process PIL image to tensor"""
        inputs = self.image_processor(image, return_tensors="pt")
        return inputs['pixel_values']


def main():
    parser = argparse.ArgumentParser(description='Test quantized CLIP vision encoder')
    parser.add_argument('--image-file', type=str,
                      default='https://llava-vl.github.io/static/images/view.jpg',
                      help='Path to image file or URL')
    parser.add_argument('--vision-tower', type=str,
                      default='openai/clip-vit-large-patch14-336',
                      help='CLIP vision tower model name')
    parser.add_argument('--device', type=str, default='cuda' if torch.cuda.is_available() else 'cpu',
                      help='Device to run on')
    parser.add_argument('--quantization-type', type=str, choices=['static', 'dynamic'],
                      default='dynamic',
                      help='Type of quantization to use')

    args = parser.parse_args()

    print("=" * 50)
    print("CLIP Vision Encoder INT8 Quantization Test")
    print("=" * 50)

    # Initialize the quantized CLIP tower
    clip_tower = QuantizedCLIPVisionTower(args.vision_tower, args.device)

    # Load test image
    print(f"Loading image: {args.image_file}")
    image = load_image(args.image_file)
    print(f"Image size: {image.size}")

    # Process image
    image_tensor = clip_tower.process_image(image)
    print(f"Image tensor shape: {image_tensor.shape}")

    # Test original model
    print("\n" + "="*30)
    print("Testing Original Model")
    print("="*30)

    start_time = torch.cuda.Event(enable_timing=True)
    end_time = torch.cuda.Event(enable_timing=True)

    if args.device == 'cuda':
        start_time.record()

    original_features = clip_tower.forward(image_tensor)

    if args.device == 'cuda':
        end_time.record()
        torch.cuda.synchronize()
        original_time = start_time.elapsed_time(end_time)
        print(f"Original model inference time: {original_time:.2f} ms")

    print(f"Original features shape: {original_features.shape}")
    print(f"Original features mean: {original_features.mean().item():.6f}")
    print(f"Original features std: {original_features.std().item():.6f}")

    # Quantize the model
    print("\n" + "="*30)
    print("Quantizing Model")
    print("="*30)

    if args.quantization_type == 'dynamic':
        clip_tower.quantize_to_int8_dynamic()
    else:
        clip_tower.quantize_to_int8()

    # Test quantized model
    print("\n" + "="*30)
    print("Testing Quantized Model")
    print("="*30)

    if args.device == 'cuda':
        start_time.record()

    quantized_features = clip_tower.forward(image_tensor)

    if args.device == 'cuda':
        end_time.record()
        torch.cuda.synchronize()
        quantized_time = start_time.elapsed_time(end_time)
        print(f"Quantized model inference time: {quantized_time:.2f} ms")
        print(f"Speedup: {original_time/quantized_time:.2f}x")

    print(f"Quantized features shape: {quantized_features.shape}")
    print(f"Quantized features mean: {quantized_features.mean().item():.6f}")
    print(f"Quantized features std: {quantized_features.std().item():.6f}")

    # Compare outputs
    print("\n" + "="*30)
    print("Comparing Outputs")
    print("="*30)

    # Calculate similarity metrics
    mse = torch.nn.functional.mse_loss(original_features.float(), quantized_features.float())
    cosine_sim = torch.nn.functional.cosine_similarity(
        original_features.flatten().float(),
        quantized_features.flatten().float(),
        dim=0
    )

    print(f"Mean Squared Error: {mse.item():.6f}")
    print(f"Cosine Similarity: {cosine_sim.item():.6f}")

    # Calculate relative error
    rel_error = torch.abs(original_features - quantized_features) / (torch.abs(original_features) + 1e-8)
    print(f"Mean Relative Error: {rel_error.mean().item():.6f}")
    print(f"Max Relative Error: {rel_error.max().item():.6f}")

    print("\n" + "="*50)
    print("Quantization Test Completed!")
    print("="*50)


if __name__ == "__main__":
    main()