aMUSEd と OpenVINO による軽量な画像生成#

class TextEncoderWrapper(torch.nn.Module): 
    def __init__(self, text_encoder): 
        super().__init__() 
        self.text_encoder = text_encoder 

    def forward(self, input_ids=None, return_dict=None, output_hidden_states=None): 
        outputs = self.text_encoder( 
            input_ids=input_ids, 
            return_dict=return_dict, 
            output_hidden_states=output_hidden_states, 
        ) 

        return outputs.text_embeds, outputs.last_hidden_state, outputs.hidden_states 

input_ids = pipe.tokenizer( 
    prompt, 
    return_tensors="pt", 
    padding="max_length", 
    truncation=True, 
    max_length=pipe.tokenizer.model_max_length, 
) 

input_example = { 
    "input_ids": input_ids.input_ids, 
    "return_dict": torch.tensor(True), 
    "output_hidden_states": torch.tensor(True), 
} 

convert(TextEncoderWrapper(pipe.text_encoder), TEXT_ENCODER_OV_PATH, input_example)

class TransformerWrapper(torch.nn.Module): 
    def __init__(self, transformer): 
        super().__init__() 
        self.transformer = transformer 

    def forward( 
        self, latents=None, 
        micro_conds=None, 
        pooled_text_emb=None, 
        encoder_hidden_states=None, 
    ): 
        return self.transformer( 
            latents, 
            micro_conds=micro_conds, 
            pooled_text_emb=pooled_text_emb, 
            encoder_hidden_states=encoder_hidden_states, 
        ) 

shape = (1, 16, 16) 
latents = torch.full(shape, pipe.scheduler.config.mask_token_id, dtype=torch.long) 
latents = torch.cat([latents] * 2) 

example_input = { 
    "latents": latents, 
    "micro_conds": torch.rand([2, 5], dtype=torch.float32), 
    "pooled_text_emb": torch.rand([2, 768], dtype=torch.float32), 
    "encoder_hidden_states": torch.rand([2, 77, 768], dtype=torch.float32), 
} 

pipe.transformer.eval() 
w_transformer = TransformerWrapper(pipe.transformer) 
convert(w_transformer, TRANSFORMER_OV_PATH, example_input)

def get_latents(): 
    shape = (1, 16, 16) 
    latents = torch.full(shape, pipe.scheduler.config.mask_token_id, dtype=torch.long) 
    model_input = torch.cat([latents] * 2) 

    model_output = pipe.transformer( 
        model_input, 
        micro_conds=torch.rand([2, 5], dtype=torch.float32), 
        pooled_text_emb=torch.rand([2, 768], dtype=torch.float32), 
        encoder_hidden_states=torch.rand([2, 77, 768], dtype=torch.float32), 
    ) 
    guidance_scale = 10.0 
    uncond_logits, cond_logits = model_output.chunk(2) 
    model_output = uncond_logits + guidance_scale * (cond_logits - uncond_logits) 

    latents = pipe.scheduler.step( 
        model_output=model_output, 
        timestep=torch.tensor(0), 
        sample=latents, 
    ).prev_sample 

    return latents 

class VQVAEWrapper(torch.nn.Module): 
    def __init__(self, vqvae): 
        super().__init__() 
        self.vqvae = vqvae 

    def forward(self, latents=None, force_not_quantize=True, shape=None): 
    outputs = self.vqvae.decode( 
        latents, 
        force_not_quantize=force_not_quantize, 
        shape=shape.tolist(), 
        ) 

        return outputs 

latents = get_latents() 
example_vqvae_input = { 
    "latents": latents, 
    "force_not_quantize": torch.tensor(True), 
    "shape": torch.tensor((1, 16, 16, 64)), 
} 

convert(VQVAEWrapper(pipe.vqvae), VQVAE_OV_PATH, example_vqvae_input)

%%skip not $to_quantize.value 

import datasets 
from tqdm.auto import tqdm 
from typing import Any, Dict, List 
import pickle 
import numpy as np 

def disable_progress_bar(pipeline, disable=True): 
    if not hasattr(pipeline, "_progress_bar_config"): 
        pipeline._progress_bar_config = {'disable': disable} 
    else: 
        pipeline._progress_bar_config['disable'] = disable 

class CompiledModelDecorator(ov.CompiledModel): 
    def __init__(self, compiled_model: ov.CompiledModel, data_cache: List[Any] = None, keep_prob: float = 0.5): 
        super().__init__(compiled_model) 
        self.data_cache = data_cache if data_cache is not None else [] 
        self.keep_prob = keep_prob 

    def __call__(self, *args, **kwargs): 
        if np.random.rand() <= self.keep_prob: 
            self.data_cache.append(*args) 
        return super().__call__(*args, **kwargs) 

def collect_calibration_data(ov_transformer_model, calibration_dataset_size: int) -> List[Dict]: 
    calibration_dataset_filepath = Path(f"calibration_data/{calibration_dataset_size}.pkl") 
    if not calibration_dataset_filepath.exists(): 
        calibration_data = [] 
        pipe.transformer.transformer = CompiledModelDecorator(ov_transformer_model, calibration_data, keep_prob=1.0) 
        disable_progress_bar(pipe) 

        dataset = datasets.load_dataset("google-research-datasets/conceptual_captions", split="train", trust_remote_code=True).shuffle(seed=42) 

        # データ収集のための推論を実行 
        pbar = tqdm(total=calibration_dataset_size) 
        for batch in dataset: 
            prompt = batch["caption"] 
            if len(prompt) > pipe.tokenizer.model_max_length: 
                continue 
            pipe(prompt, generator=torch.Generator('cpu').manual_seed(0)) 
            pbar.update(len(calibration_data) - pbar.n) 
            if pbar.n >= calibration_dataset_size: 
                break 

        pipe.transformer.transformer = ov_transformer_model 
        disable_progress_bar(pipe, disable=False) 

        calibration_dataset_filepath.parent.mkdir(exist_ok=True, parents=True) 
        with open(calibration_dataset_filepath, 'wb') as f: 
            pickle.dump(calibration_data, f) 
    with open(calibration_dataset_filepath, 'rb') as f: 
        calibration_data = pickle.load(f) 
    return calibration_data

%%skip not $to_quantize.value 

from nncf.quantization.advanced_parameters import AdvancedSmoothQuantParameters 
from nncf.quantization.range_estimator import RangeEstimatorParameters, StatisticsCollectorParameters, 
StatisticsType, \ 
    AggregatorType 
import nncf 

CALIBRATION_DATASET_SIZE = 12 * 25 

if not QUANTIZED_TRANSFORMER_OV_PATH.exists(): 
    calibration_data = collect_calibration_data(ov_transformer, CALIBRATION_DATASET_SIZE) 
    quantized_model = nncf.quantize( 
        core.read_model(TRANSFORMER_OV_PATH), 
        nncf.Dataset(calibration_data), 
        model_type=nncf.ModelType.TRANSFORMER, 
        subset_size=len(calibration_data), 
        # 畳み込みを無視することで、推論速度を大幅に低下させることなく生成の品質を向上させる 
        ignored_scope=nncf.IgnoredScope(types=["Convolution"]), 
        # 0.85 という値は、以下で計算されたインセプション・スコアに基づくグリッド検索を使用して取得 
        advanced_parameters=nncf.AdvancedQuantizationParameters( 
            smooth_quant_alphas=AdvancedSmoothQuantParameters(matmul=0.85), 
            # 活性化統計収集中に、外れ値の 1 % を無視することで量子化品質が向上     
            activations_range_estimator_params=RangeEstimatorParameters( 
                min=StatisticsCollectorParameters(statistics_type=StatisticsType.MIN, 
                    aggregator_type=AggregatorType.MEAN_NO_OUTLIERS, 
                    quantile_outlier_prob=0.01), 
                max=StatisticsCollectorParameters(statistics_type=StatisticsType.MAX, 
                    aggregator_type=AggregatorType.MEAN_NO_OUTLIERS, 
                    quantile_outlier_prob=0.01) 
            )
         )
     ) 
    ov.save_model(quantized_model, QUANTIZED_TRANSFORMER_OV_PATH)

%%skip not $to_quantize.value 

original_ov_transformer_model = pipe.transformer.transformer 
pipe.transformer.transformer = core.compile_model(QUANTIZED_TRANSFORMER_OV_PATH, device.value) 

image = pipe(prompt, generator=torch.Generator('cpu').manual_seed(8)).images[0] 
image.save('text2image_256_quantized.png') 

pipe.transformer.transformer = original_ov_transformer_model 

display(image)

%%skip not $to_quantize.value 

from torchmetrics.image.inception import InceptionScore 
from torchvision import transforms as transforms 
from itertools import islice 
import time 

VALIDATION_DATASET_SIZE = 100 

def compute_inception_score(ov_transformer_model_path, validation_set_size, batch_size=100): 
    original_ov_transformer_model = pipe.transformer.transformer 
    pipe.transformer.transformer = core.compile_model(ov_transformer_model_path, device.value) 

    disable_progress_bar(pipe) 
    dataset = datasets.load_dataset("google-research-datasets/conceptual_captions", "unlabeled", split="validation", trust_remote_code=True).shuffle(seed=42) 
    dataset = islice(dataset, validation_set_size) 

    inception_score = InceptionScore(normalize=True, splits=1) 

    images = [] 
    infer_times = [] 
    for batch in tqdm(dataset, total=validation_set_size, desc="Computing Inception Score"): 
        prompt = batch["caption"] 
        if len(prompt) > pipe.tokenizer.model_max_length: 
            continue 
        start_time = time.perf_counter() 
        image = pipe(prompt, generator=torch.Generator('cpu').manual_seed(0)).images[0] 
        infer_times.append(time.perf_counter() - start_time) 
        image = transforms.ToTensor()(image) 
        images.append(image) 

    mean_perf_time = sum(infer_times) / len(infer_times) 

    while len(images) > 0: 
        images_batch = torch.stack(images[-batch_size:]) 
        images = images[:-batch_size] 
        inception_score.update(images_batch) 
    kl_mean, kl_std = inception_score.compute() 

    pipe.transformer.transformer = original_ov_transformer_model 
    disable_progress_bar(pipe, disable=False) 

    return kl_mean, mean_perf_time 

original_inception_score, original_time = compute_inception_score(TRANSFORMER_OV_PATH, VALIDATION_DATASET_SIZE) 
print(f"Original pipeline Inception Score: {original_inception_score}") 
quantized_inception_score, quantized_time = compute_inception_score(QUANTIZED_TRANSFORMER_OV_PATH, VALIDATION_DATASET_SIZE) 
print(f"Quantized pipeline Inception Score: {quantized_inception_score}") 
print(f"Quantization speed-up: {original_time / quantized_time:.2f}x")