saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-28 20:11:22 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

clean up calibration model (#12255)

Browse source
This commit is contained in:
Yufeng Li 2022-07-21 14:50:28 -07:00 committed by GitHub
parent 45c0be8a25
commit a18b080513
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 13 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

22
onnxruntime/python/tools/quantization/quantize.py
View file

@ -4,6 +4,7 @@
# license information.
# --------------------------------------------------------------------------
import logging
import tempfile
from pathlib import Path
from onnx import onnx_pb as onnx_proto
@ -142,15 +143,18 @@ def quantize_static(
calib_extra_options = {
key: extra_options.get(name) for (name, key) in calib_extra_options_keys if name in extra_options
}
calibrator = create_calibrator(
model,
op_types_to_quantize,
calibrate_method=calibrate_method,
use_external_data_format=use_external_data_format,
extra_options=calib_extra_options,
)
calibrator.collect_data(calibration_data_reader)
tensors_range = calibrator.compute_range()
with tempfile.TemporaryDirectory(prefix="ort.quant.") as quant_tmp_dir:
calibrator = create_calibrator(
model,
op_types_to_quantize,
augmented_model_path=Path(quant_tmp_dir).joinpath("augmented_model.onnx").as_posix(),
calibrate_method=calibrate_method,
use_external_data_format=use_external_data_format,
extra_options=calib_extra_options,
)
calibrator.collect_data(calibration_data_reader)
tensors_range = calibrator.compute_range()
check_static_quant_arguments(quant_format, activation_type, weight_type)