coco
55 lines
1.6 KiB
Python
55 lines
1.6 KiB
Python
import os
|
|
from pathlib import Path
|
|
|
|
import onnxruntime as ort
|
|
from PIL import Image
|
|
from onnxruntime.quantization import quantize_dynamic, QuantType, quant_pre_process
|
|
|
|
from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize, InterpolationMode
|
|
|
|
model = "clip-image-encoder.onnx"
|
|
model_prep = "clip-image-encoder-quant-pre.onnx"
|
|
model_quant = "clip-image-encoder-quant-int8.onnx"
|
|
|
|
|
|
def quant():
|
|
cur_path = Path(os.curdir)
|
|
|
|
quant_pre_process(model, model_prep) # preprocess for quantization
|
|
quantize_dynamic(cur_path / model_prep, cur_path / model_quant, weight_type=QuantType.QInt8,
|
|
nodes_to_exclude=['/conv1/Conv'])
|
|
|
|
|
|
def test():
|
|
ort_session = ort.InferenceSession(model_quant)
|
|
|
|
input_name = ort_session.get_inputs()[0].name
|
|
|
|
def _convert_image_to_rgb(image: Image):
|
|
return image.convert("RGB")
|
|
|
|
def _transform(n_px):
|
|
return Compose([
|
|
Resize(n_px, interpolation=InterpolationMode.NEAREST),
|
|
CenterCrop(n_px),
|
|
_convert_image_to_rgb,
|
|
ToTensor(),
|
|
Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
|
|
])
|
|
|
|
preprocess = _transform(224)
|
|
|
|
image_input = preprocess(Image.open("../../image.jpg")).unsqueeze(0).to("cpu")
|
|
|
|
outputs = ort_session.run(None, {input_name: image_input.numpy()})
|
|
|
|
print(outputs[0])
|
|
return outputs[0]
|
|
|
|
|
|
if __name__ == '__main__':
|
|
quant()
|
|
# res = test()
|
|
# python -m onnxruntime.tools.check_onnx_model_mobile_usability clip-text-encoder-quant-int8.onnx
|
|
# python -m onnxruntime.tools.convert_onnx_models_to_ort clip-image-encoder-quant-int8.onnx
|