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

fix issues caused by quantize/calibrate changes (#6802)

Browse source
This commit is contained in:
stevenlix 2021-02-25 05:41:21 -08:00 committed by GitHub
parent 7465673e33
commit d5f292ab73
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 10 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

17
onnxruntime/python/tools/quantization/E2E_example_model/image_classification/trt/resnet50/e2e_tensorrt_resnet_example.py
View file

@ -7,7 +7,7 @@ import logging
from PIL import Image
import onnx
import onnxruntime
from onnxruntime.quantization import CalibrationDataReader, create_calibrater, write_calibration_table
from onnxruntime.quantization import CalibrationDataReader, create_calibrator, write_calibration_table
class ImageNetDataReader(CalibrationDataReader):
@ -258,9 +258,13 @@ class ImageClassificationEvaluator:
def convert_model_batch_to_dynamic(model_path):
model = onnx.load(model_path)
input = model.graph.input
input_name = input[0].name
shape = input[0].type.tensor_type.shape
initializers = [node.name for node in model.graph.initializer]
inputs = []
for node in model.graph.input:
if node.name not in initializers:
inputs.append(node)
input_name = inputs[0].name
shape = inputs[0].type.tensor_type.shape
dim = shape.dim
if not dim[0].dim_param:
dim[0].dim_param = 'N'
@ -329,6 +333,8 @@ if __name__ == '__main__':
# Generate INT8 calibration table
if calibration_table_generation_enable:
calibrator = create_calibrator(new_model_path, [], augmented_model_path=augmented_model_path)
calibrator.set_execution_providers(["CUDAExecutionProvider"])
data_reader = ImageNetDataReader(ilsvrc2012_dataset_path,
start_index=0,
end_index=calibration_dataset_size,
@ -336,9 +342,6 @@ if __name__ == '__main__':
batch_size=batch_size,
model_path=augmented_model_path,
input_name=input_name)
# For TensorRT calibration, augment all FP32 tensors (empty op_types), disable ORT graph optimization and skip quantization parameter calculation
calibrator = create_calibrater(new_model_path)
calibrator.set_execution_providers(["CUDAExecutionProvider"])
calibrator.collect_data(data_reader)
write_calibration_table(calibrator.compute_range())