Static Quantization Fails for Transformer

[nogifeet]

I am trying to quantize the pix2struct-docvqa-base model using the quantize_static() provided by onnxruntime.

I am able to run quantize_static() without any errors, but the code runs in an infinite loop and eventually kills the kernel.

Here is what I have tried..

To Export the model from HF to Onnx:

!pip install 'optimum[exporters]'
  
!optimum-cli export onnx --model="google/pix2struct-docvqa-base" \
    --device "cpu" --atol=1e-3 --framework="pt" \
    --task="visual-question-answering-with-past" \
    "./temp/"

Running Static Quantization:

import onnx
from onnxruntime.quantization import quantize_static, QuantFormat, CalibrationMethod, CalibrationDataReader
from onnxruntime.quantization.preprocess import quant_pre_process
from onnxruntime.quantization.calibrate import CalibrationDataReader
import numpy as np
import os
import time

class CalibrationDataProvider(CalibrationDataReader):
    def __init__(self):
        super(CalibrationDataProvider, self).__init__()
        self.counter = 0

    def get_next(self):
        if self.counter > 2:
            return None
        else:
            self.counter += 1
            return {
                    'flattened_patches': np.random.randn(1,2048,770).astype(np.float32),
                    'attention_mask' : np.random.randn(1,2048).astype(np.int64)
                   }

cdr = CalibrationDataProvider()

model_fp32 = './temp/encoder_model.onnx'
model_preprocess = './temp/encoder_preprocess.onnx'
model_quant = './temp/ecnoder_quantized_static.onnx'

preprocess_model = quant_pre_process(model_fp32,model_preprocess)

Up until this point, everything runs perfectly fine, but the next command just fails without any error indication as to why...

quantized_model = quantize_static(
    model_preprocess,  
    model_quant,
    calibrate_method=CalibrationMethod.Entropy,
    quant_format=QuantFormat.QOperator,
    calibration_data_reader=cdr)

I have tried to play with the different options available in quantize_static(), but none of them seem to help.

[yihonglyu]

@nogifeet Could you provide more information about the Platform, OS Version, ONNX Runtime Installation, ONNX Runtime Version or Commit ID, Architecture .. etc? I run the script with onnxruntime 1.16.1 on Linux and it is successful.

[nogifeet]

@yihonglyu yes, the crash was around my system running out of memory.
Thanks for verifying that the code is working.