File size: 3,472 Bytes
7548bcf f20b933 7548bcf f20b933 7548bcf f20b933 7548bcf | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 | """The exact ONNX Runtime static-quantization baseline invocation used for the
comparison row in this repo's model card.
Usage: python quantize_baseline.py <fp32_model.onnx> <calib_dir> <out.onnx>
Deps: pip install onnxruntime numpy pillow onnx
Configuration: QDQ format, per-channel weights, quant_pre_process, and
onnxruntime's default MinMax calibration. onnxruntime also offers Percentile
and Entropy calibration, reduce_range, and per-node exclusion lists; none of
those are applied here, so this is a standard static configuration rather than
an exhaustively tuned one.
Preprocessing mirrors the eval transform stated on the card: bicubic resize of
the shorter side to size/CROP_PCT, center crop to the model's input size,
/255, ImageNet mean/std. onnxruntime version used for the published row: 1.27.
"""
import glob
import os
import sys
import numpy as np
from PIL import Image
CROP_PCT = 0.95 # this model's eval crop fraction (see the model card)
MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)
def preprocess(path, size):
img = Image.open(path).convert("RGB")
scale_size = int(round(size / CROP_PCT))
w, h = img.size
if w < h:
nw, nh = scale_size, int(round(h * scale_size / w))
else:
nw, nh = int(round(w * scale_size / h)), scale_size
img = img.resize((nw, nh), Image.BICUBIC)
left, top = (nw - size) // 2, (nh - size) // 2
img = img.crop((left, top, left + size, top + size))
x = ((np.asarray(img, np.float32) / 255.0 - MEAN) / STD).transpose(2, 0, 1)[None]
return x.astype(np.float32)
def main():
fp32_path, calib_dir, out_path = sys.argv[1], sys.argv[2], sys.argv[3]
import onnxruntime as ort
from onnxruntime.quantization import (
CalibrationDataReader, QuantFormat, QuantType, quantize_static,
)
from onnxruntime.quantization.shape_inference import quant_pre_process
sess = ort.InferenceSession(fp32_path, providers=["CPUExecutionProvider"])
size = sess.get_inputs()[0].shape[-1]
del sess
class Reader(CalibrationDataReader):
def __init__(self):
self.files = sorted(glob.glob(os.path.join(calib_dir, "*.jpg")))
self.i = 0
def get_next(self):
if self.i >= len(self.files):
return None
x = preprocess(self.files[self.i], size)
self.i += 1
return {"input": x}
pre_path = out_path + ".preprocessed.tmp.onnx"
try:
quant_pre_process(fp32_path, pre_path)
src = pre_path
except Exception as e:
# Same fallback as the pipeline that produced the published row: on
# some exports (EdgeNeXt, XCiT) onnxruntime 1.27's symbolic shape
# inference fails here, and the baseline quantizes the raw export.
print(f"quant_pre_process failed ({type(e).__name__}) -> quantizing the raw export", flush=True)
src = fp32_path
try:
quantize_static(
src, out_path, Reader(),
quant_format=QuantFormat.QDQ, per_channel=True,
activation_type=QuantType.QUInt8, weight_type=QuantType.QInt8,
)
finally:
if os.path.exists(pre_path):
os.remove(pre_path)
print(f"wrote {out_path} ({os.path.getsize(out_path):,} bytes)")
if __name__ == "__main__":
main()
|