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

Add the missing onnx_proto import (#4705)

Browse source 
* add missing onnx_proto import
* Fix TensorProto usage in calibrate.py
* remove unused imports
This commit is contained in:
Vagif 2020-08-11 04:46:21 +09:00 committed by GitHub
parent 2e3ccc7518
commit 6499a38b7d
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 4 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

13
onnxruntime/python/tools/quantization/calibrate.py
View file

@ -6,19 +6,14 @@
# license information.
# --------------------------------------------------------------------------
import os
import sys
import numpy as np
import onnx
import onnxruntime
from onnx import helper, TensorProto, numpy_helper
import re
import subprocess
import json
from onnx import helper, TensorProto
import abc
class CalibrationDataReader(metaclass=abc.ABCMeta):
@classmethod
def __subclasshook__(cls,subclass):
@ -77,7 +72,7 @@ class ONNXCalibrater:
if input_tensor_name in value_infos.keys():
vi = value_infos[input_tensor_name]
if vi.type.HasField(
'tensor_type') and vi.type.tensor_type.elem_type == onnx_proto.TensorProto.FLOAT and (
'tensor_type') and vi.type.tensor_type.elem_type == TensorProto.FLOAT and (
input_tensor_name not in model.graph.initializer):
tensors_to_calibrate.add(input_tensor_name)
@ -85,7 +80,7 @@ class ONNXCalibrater:
if output_tensor_name in value_infos.keys():
vi = value_infos[output_tensor_name]
if vi.type.HasField(
'tensor_type') and vi.type.tensor_type.elem_type == onnx_proto.TensorProto.FLOAT:
'tensor_type') and vi.type.tensor_type.elem_type == TensorProto.FLOAT:
tensors_to_calibrate.add(output_tensor_name)
for tensor in tensors_to_calibrate: