SigLIP によるゼロショット画像分類#

import requests 
from io import BytesIO 
from PIL import Image 
from requests.packages.urllib3.exceptions import InsecureRequestWarning 

requests.packages.urllib3.disable_warnings(InsecureRequestWarning) 

def check_text_data(data): 
    """ 
    Check if the given data is text-based.
    """ 
    if isinstance(data, str): 
        return True 
    if isinstance(data, list): 
        return all(isinstance(x, str) for x in data) 
    return False 

def get_pil_from_url(url): 
    """ 
    Downloads and converts an image from a URL to a PIL Image object.
    """ 
    response = requests.get(url, verify=False, timeout=20) 
    image = Image.open(BytesIO(response.content)) 
    return image.convert("RGB") 

def collate_fn(example, image_column="image_url", text_column="caption"): 
    """ 
    Preprocesses an example by loading and transforming image and text data. 
    Checks if the text data in the example is valid by calling the `check_text_data` function. 
    Downloads the image specified by the URL in the image_column by calling the `get_pil_from_url` function. 
    If there is any error during the download process, returns None. 
    Returns the preprocessed inputs with transformed image and text data.
    """ 
    assert len(example) == 1 
    example = example[0] 

    if not check_text_data(example[text_column]): 
        raise ValueError("Text data is not valid") 

    url = example[image_column] 
    try: 
        image = get_pil_from_url(url) 
        h, w = image.size 
        if h == 1 or w == 1: 
            return None 
    except Exception: 
        return None 

    inputs = processor( 
        text=example[text_column], 
        images=[image], 
        return_tensors="pt", 
        padding="max_length", 
    ) 
    if inputs["input_ids"].shape[1] > model.config.text_config.max_position_embeddings: 
        return None 
    return inputs

import torch 
from datasets import load_dataset 
from tqdm.notebook import tqdm 

def prepare_calibration_data(dataloader, init_steps): 
    """ 
    This function prepares calibration data from a dataloader for a specified number of initialization steps. 
    It iterates over the dataloader, fetching batches and storing the relevant data.
    """ 
    data = [] 
    print(f"Fetching {init_steps} for the initialization...") 
    counter = 0 
    for batch in tqdm(dataloader): 
        if counter == init_steps: 
            break 
        if batch: 
            counter += 1 
            with torch.no_grad(): 
                data.append( 
                    { 
                        "pixel_values": batch["pixel_values"].to("cpu"), 
                        "input_ids": batch["input_ids"].to("cpu"), 
                    } 
                ) 
    return data 

def prepare_dataset(opt_init_steps=300, max_train_samples=1000): 
    """ 
    Prepares a vision-text dataset for quantization.
    """ 
    dataset = load_dataset("google-research-datasets/conceptual_captions", streaming=True, trust_remote_code=True) 
    train_dataset = dataset["train"].shuffle(seed=42, buffer_size=max_train_samples) 
    dataloader = torch.utils.data.DataLoader(train_dataset, collate_fn=collate_fn, batch_size=1) 
    calibration_data = prepare_calibration_data(dataloader, opt_init_steps) 
    return calibration_data

calibration_data = prepare_dataset()

import nncf 
import logging 

nncf.set_log_level(logging.ERROR) 

if len(calibration_data) == 0: 
    raise RuntimeError("Calibration dataset is empty.Please check internet connection and try to download images manually.") 

calibration_dataset = nncf.Dataset(calibration_data) 
quantized_ov_model = nncf.quantize( 
    model=ov_model, 
    calibration_dataset=calibration_dataset, 
    model_type=nncf.ModelType.TRANSFORMER, 
)

from scipy.special import softmax 

input_labels = [ 
    "cat", 
    "dog", 
    "wolf", 
    "tiger", 
    "man", 
    "horse", 
    "frog", 
    "tree", 
    "house", 
    "computer", 
] 
text_descriptions = [f"This is a photo of a {label}" for label in input_labels] 

inputs = processor(text=text_descriptions, images=[image], return_tensors="pt", padding="max_length") 
compiled_int8_ov_model = ov.compile_model(quantized_ov_model, device.value) 

logits_per_image_out = compiled_int8_ov_model.output(0) 
ov_logits_per_image = compiled_int8_ov_model(dict(inputs))[logits_per_image_out] 
probs = softmax(ov_logits_per_image, axis=1) 
visualize_result(image, input_labels, probs[0])

from pathlib import Path 

fp16_model_path = "siglip-base-patch16-224.xml" 
ov.save_model(ov_model, fp16_model_path) 

int8_model_path = "siglip-base-patch16-224_int8.xml" 
ov.save_model(quantized_ov_model, int8_model_path) 

fp16_ir_model_size = Path(fp16_model_path).with_suffix(".bin").stat().st_size / 1024 / 1024 
quantized_model_size = Path(int8_model_path).with_suffix(".bin").stat().st_size / 1024 / 1024 
print(f"FP16 IR model size: {fp16_ir_model_size:.2f} MB") 
print(f"INT8 model size: {quantized_model_size:.2f} MB") 
print(f"Model compression rate: {fp16_ir_model_size / quantized_model_size:.3f}")

import time 

def calculate_inference_time(model_path, calibration_data): 
    model = ov.compile_model(model_path, device.value) 
    output_layer = model.output(0) 
    inference_time = [] 
    for batch in calibration_data: 
        start = time.perf_counter() 
        _ = model(batch)[output_layer] 
        end = time.perf_counter() 
        delta = end - start 
        inference_time.append(delta) 
    return np.median(inference_time)

fp16_latency = calculate_inference_time(fp16_model_path, calibration_data) 
int8_latency = calculate_inference_time(int8_model_path, calibration_data) 
print(f"Performance speed up: {fp16_latency / int8_latency:.3f}")