mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-09 17:28:58 +00:00
Python module for dumping activation tensors when running an ONNX model (#12474)
Python module for dumping activation tensors when running an ONNX model This is the first step towards a quantization debugging tool. We dump the activation tensors. Next step would be to compare them: original model vs quantized model (running with same input) to see where the difference becomes significant.
This commit is contained in:
Chen Fu
GitHub
parent
2681648f5b
commit
47b787c28f
3 changed files with 329 additions and 7 deletions
|
|
@ -10,6 +10,7 @@ import itertools
|
|||
import uuid
|
||||
from enum import Enum
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional, Sequence
|
||||
|
||||
import numpy as np
|
||||
import onnx
|
||||
|
|
@ -36,12 +37,21 @@ class CalibrationDataReader(metaclass=abc.ABCMeta):
|
|||
"""generate the input data dict for ONNXinferenceSession run"""
|
||||
raise NotImplementedError
|
||||
|
||||
def __iter__(self):
|
||||
return self
|
||||
|
||||
def __next__(self):
|
||||
result = self.get_next()
|
||||
if result is None:
|
||||
raise StopIteration
|
||||
return result
|
||||
|
||||
|
||||
class CalibraterBase:
|
||||
def __init__(
|
||||
self,
|
||||
model,
|
||||
op_types_to_calibrate=[],
|
||||
op_types_to_calibrate: Optional[Sequence[str]] = None,
|
||||
augmented_model_path="augmented_model.onnx",
|
||||
symmetric=False,
|
||||
use_external_data_format=False,
|
||||
|
|
@ -106,7 +116,7 @@ class CalibraterBase:
|
|||
tensor_type_to_calibrate = set([TensorProto.FLOAT, TensorProto.FLOAT16])
|
||||
|
||||
for node in model.graph.node:
|
||||
if len(self.op_types_to_calibrate) == 0 or node.op_type in self.op_types_to_calibrate:
|
||||
if not self.op_types_to_calibrate or node.op_type in self.op_types_to_calibrate:
|
||||
for tensor_name in itertools.chain(node.input, node.output):
|
||||
if tensor_name in value_infos.keys():
|
||||
vi = value_infos[tensor_name]
|
||||
|
|
@ -150,7 +160,7 @@ class MinMaxCalibrater(CalibraterBase):
|
|||
def __init__(
|
||||
self,
|
||||
model,
|
||||
op_types_to_calibrate=[],
|
||||
op_types_to_calibrate: Optional[Sequence[str]] = None,
|
||||
augmented_model_path="augmented_model.onnx",
|
||||
symmetric=False,
|
||||
use_external_data_format=False,
|
||||
|
|
@ -320,7 +330,7 @@ class HistogramCalibrater(CalibraterBase):
|
|||
def __init__(
|
||||
self,
|
||||
model,
|
||||
op_types_to_calibrate=[],
|
||||
op_types_to_calibrate: Optional[Sequence[str]] = None,
|
||||
augmented_model_path="augmented_model.onnx",
|
||||
use_external_data_format=False,
|
||||
method="percentile",
|
||||
|
|
@ -429,7 +439,7 @@ class EntropyCalibrater(HistogramCalibrater):
|
|||
def __init__(
|
||||
self,
|
||||
model,
|
||||
op_types_to_calibrate=[],
|
||||
op_types_to_calibrate: Optional[Sequence[str]] = None,
|
||||
augmented_model_path="augmented_model.onnx",
|
||||
use_external_data_format=False,
|
||||
method="entropy",
|
||||
|
|
@ -463,7 +473,7 @@ class PercentileCalibrater(HistogramCalibrater):
|
|||
def __init__(
|
||||
self,
|
||||
model,
|
||||
op_types_to_calibrate=[],
|
||||
op_types_to_calibrate: Optional[Sequence[str]] = None,
|
||||
augmented_model_path="augmented_model.onnx",
|
||||
use_external_data_format=False,
|
||||
method="percentile",
|
||||
|
|
@ -810,7 +820,7 @@ class HistogramCollector(CalibrationDataCollector):
|
|||
|
||||
def create_calibrator(
|
||||
model,
|
||||
op_types_to_calibrate=[],
|
||||
op_types_to_calibrate: Optional[Sequence[str]] = None,
|
||||
augmented_model_path="augmented_model.onnx",
|
||||
calibrate_method=CalibrationMethod.MinMax,
|
||||
use_external_data_format=False,
|
||||
|
|
|
|||
147
onnxruntime/python/tools/quantization/save_activations.py
Normal file
147
onnxruntime/python/tools/quantization/save_activations.py
Normal file
|
|
@ -0,0 +1,147 @@
|
|||
# --------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft, Intel Corporation. All rights reserved.
|
||||
# Licensed under the MIT License. See License.txt in the project root for
|
||||
# license information.
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
"""Utilities to run a given ONNX model, while saving input/output tensors of
|
||||
eligible operator nodes.
|
||||
|
||||
A use case is to debug quantization induced accuracy drop. An AI engineer can
|
||||
run the original float32 model and the quantized model with the same inputs,
|
||||
then compare the corresponding activations between the two models to find
|
||||
where the divergence is.
|
||||
|
||||
Example Usage:
|
||||
|
||||
```python
|
||||
class ExampleDataReader(CalibrationDataReader):
|
||||
def __init__(self):
|
||||
...
|
||||
def get_next(self):
|
||||
...
|
||||
|
||||
input_data_reader = ExampleDataReader()
|
||||
|
||||
aug_model = modify_model_output_intermediate_tensors (path_to_onnx_model)
|
||||
augmented_model_path = str(Path(self._tmp_model_dir.name).joinpath("augmented_model.onnx"))
|
||||
onnx.save(
|
||||
aug_model,
|
||||
augmented_model_path,
|
||||
save_as_external_data=False,
|
||||
)
|
||||
|
||||
tensor_dict = collect_activations(augmented_model_path, data_reader)
|
||||
```
|
||||
|
||||
`tensor_dict` points to a dictionary where the keys are tensor names and each value
|
||||
is a list of tensors, one from each model run
|
||||
|
||||
"""
|
||||
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional, Sequence, Union
|
||||
|
||||
import numpy
|
||||
import onnx
|
||||
from onnx import ModelProto, TensorProto, helper, numpy_helper
|
||||
|
||||
import onnxruntime
|
||||
|
||||
from .calibrate import CalibraterBase, CalibrationDataReader
|
||||
from .quant_utils import clone_model_with_shape_infer
|
||||
|
||||
_TENSOR_SAVE_POSTFIX = "_ReshapedSavedOutput"
|
||||
_TENSOR_SAVE_POSTFIX_LEN = len(_TENSOR_SAVE_POSTFIX)
|
||||
|
||||
|
||||
def modify_model_output_intermediate_tensors(
|
||||
onnx_model: Union[str, Path, ModelProto], op_types_for_saving: Optional[Sequence[str]] = None
|
||||
) -> ModelProto:
|
||||
"""Augment a given ONNX model to save node input/output tensors.
|
||||
|
||||
Add all input/output tensors of operator nodes to model outputs
|
||||
so that their values can be retrieved for debugging purposes.
|
||||
|
||||
Args:
|
||||
model: An ONNX model or the path to load the model.
|
||||
op_types_for_saving: Operator types for which the
|
||||
input/output should be saved. By default, saving all the
|
||||
float32/float16 tensors.
|
||||
|
||||
Returns:
|
||||
The augmented ONNX model
|
||||
"""
|
||||
|
||||
if op_types_for_saving is None:
|
||||
op_types_for_saving = []
|
||||
saver = CalibraterBase(onnx_model, op_types_to_calibrate=op_types_for_saving)
|
||||
model: ModelProto = clone_model_with_shape_infer(saver.model)
|
||||
tensors, _ = saver.select_tensors_to_calibrate(model)
|
||||
reshape_shape_name = "LinearReshape_" + str(time.time())
|
||||
reshape_shape = numpy_helper.from_array(numpy.array([-1], dtype=numpy.int64), reshape_shape_name)
|
||||
model.graph.initializer.append(reshape_shape)
|
||||
|
||||
for tensor_name in tensors:
|
||||
reshape_output = tensor_name + _TENSOR_SAVE_POSTFIX
|
||||
reshape_node = onnx.helper.make_node(
|
||||
"Reshape",
|
||||
inputs=[tensor_name, reshape_shape_name],
|
||||
outputs=[reshape_output],
|
||||
name=reshape_output,
|
||||
)
|
||||
model.graph.node.append(reshape_node)
|
||||
reshape_output_value_info = helper.make_tensor_value_info(reshape_output, TensorProto.FLOAT, [1])
|
||||
model.graph.output.append(reshape_output_value_info)
|
||||
return model
|
||||
|
||||
|
||||
def collect_activations(
|
||||
augmented_model: str,
|
||||
input_reader: CalibrationDataReader,
|
||||
session_options=None,
|
||||
execution_providers: Optional[Sequence[str]] = None,
|
||||
) -> Dict[str, List[numpy.ndarray]]:
|
||||
"""Run augmented model and collect activations tensors.
|
||||
|
||||
Args:
|
||||
augmented_model: Path to augmented model created by modify_model_output_intermediate_tensors ()
|
||||
input_reader: Logic for reading input for the model, augmented model have the same
|
||||
input with the original model.
|
||||
session_options: Optional OnnxRuntime session options for controlling model run.
|
||||
By default graph optimization is turned off
|
||||
execution_providers: Collection of execution providers for running the model.
|
||||
Only CPU EP is used by default.
|
||||
|
||||
Returns:
|
||||
A dictionary where the key is tensor name and values are list of tensors from each batch
|
||||
"""
|
||||
|
||||
if session_options is None:
|
||||
session_options = onnxruntime.SessionOptions()
|
||||
session_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL
|
||||
if execution_providers is None:
|
||||
execution_providers = ["CPUExecutionProvider"]
|
||||
|
||||
inference_session = onnxruntime.InferenceSession(
|
||||
augmented_model,
|
||||
sess_options=session_options,
|
||||
providers=execution_providers,
|
||||
)
|
||||
|
||||
intermediate_outputs = []
|
||||
for input_d in input_reader:
|
||||
intermediate_outputs.append(inference_session.run(None, input_d))
|
||||
if not intermediate_outputs:
|
||||
raise RuntimeError("No data is collected while running augmented model!")
|
||||
|
||||
output_dict = {}
|
||||
output_info = inference_session.get_outputs()
|
||||
for batch in intermediate_outputs:
|
||||
for output, output_data in zip(output_info, batch):
|
||||
if output.name.endswith(_TENSOR_SAVE_POSTFIX):
|
||||
output_name = output.name[:-_TENSOR_SAVE_POSTFIX_LEN]
|
||||
output_dict.setdefault(output_name, []).append(output_data)
|
||||
|
||||
return output_dict
|
||||
165
onnxruntime/test/python/quantization/test_save_activations.py
Normal file
165
onnxruntime/test/python/quantization/test_save_activations.py
Normal file
|
|
@ -0,0 +1,165 @@
|
|||
# -------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
# Licensed under the MIT License. See License.txt in the project root for
|
||||
# license information.
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
"""Tests for the save_activations module."""
|
||||
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import onnx
|
||||
from onnx import TensorProto, helper, numpy_helper
|
||||
|
||||
import onnxruntime
|
||||
from onnxruntime.quantization.calibrate import CalibrationDataReader
|
||||
from onnxruntime.quantization.save_activations import collect_activations, modify_model_output_intermediate_tensors
|
||||
|
||||
|
||||
def generate_input_initializer(tensor_shape, tensor_dtype, input_name):
|
||||
"""
|
||||
Helper function to generate initializers for test inputs
|
||||
"""
|
||||
tensor = np.random.normal(0, 0.3, tensor_shape).astype(tensor_dtype)
|
||||
init = numpy_helper.from_array(tensor, input_name)
|
||||
return init
|
||||
|
||||
|
||||
def construct_test_model1(test_model_path):
|
||||
""" Create an ONNX model shaped as:
|
||||
```
|
||||
(input)
|
||||
|
|
||||
Relu
|
||||
/ \
|
||||
Conv \
|
||||
| \
|
||||
Relu Conv
|
||||
| |
|
||||
Conv |
|
||||
\ /
|
||||
Add
|
||||
|
|
||||
(X6)
|
||||
```
|
||||
We are keeping all intermediate tensors as output, just for test verification
|
||||
purposes
|
||||
"""
|
||||
|
||||
input_vi = helper.make_tensor_value_info("input", TensorProto.FLOAT, [1, 3, 1, 3])
|
||||
x1_output = helper.make_tensor_value_info("X1", TensorProto.FLOAT, [1, 3, 1, 3])
|
||||
x2_output = helper.make_tensor_value_info("X2", TensorProto.FLOAT, [1, 3, 1, 3])
|
||||
x3_output = helper.make_tensor_value_info("X3", TensorProto.FLOAT, [1, 3, 1, 3])
|
||||
x4_output = helper.make_tensor_value_info("X4", TensorProto.FLOAT, [1, 3, 1, 3])
|
||||
x5_output = helper.make_tensor_value_info("X5", TensorProto.FLOAT, [1, 3, 1, 3])
|
||||
x6_output = helper.make_tensor_value_info("X6", TensorProto.FLOAT, [1, 3, 1, 3])
|
||||
w1 = generate_input_initializer([3, 3, 1, 1], np.float32, "W1")
|
||||
b1 = generate_input_initializer([3], np.float32, "B1")
|
||||
w3 = generate_input_initializer([3, 3, 1, 1], np.float32, "W3")
|
||||
b3 = generate_input_initializer([3], np.float32, "B3")
|
||||
w5 = generate_input_initializer([3, 3, 1, 1], np.float32, "W5")
|
||||
b5 = generate_input_initializer([3], np.float32, "B5")
|
||||
relu_node_1 = helper.make_node("Relu", ["input"], ["X1"], name="Relu1")
|
||||
conv_node_1 = helper.make_node("Conv", ["X1", "W1", "B1"], ["X2"], name="Conv1")
|
||||
relu_node_2 = helper.make_node("Relu", ["X2"], ["X3"], name="Relu2")
|
||||
conv_node_2 = helper.make_node("Conv", ["X3", "W3", "B3"], ["X4"], name="Conv2")
|
||||
conv_node_3 = helper.make_node("Conv", ["X1", "W5", "B5"], ["X5"], name="Conv3")
|
||||
add_node = helper.make_node("Add", ["X4", "X5"], ["X6"], name="Add")
|
||||
|
||||
# we are keeping all tensors in the output anyway for verification purpose
|
||||
graph = helper.make_graph(
|
||||
[relu_node_1, conv_node_1, relu_node_2, conv_node_2, conv_node_3, add_node],
|
||||
"test_graph_4",
|
||||
[input_vi],
|
||||
[x1_output, x2_output, x3_output, x4_output, x5_output, x6_output],
|
||||
)
|
||||
graph.initializer.add().CopyFrom(w1)
|
||||
graph.initializer.add().CopyFrom(b1)
|
||||
graph.initializer.add().CopyFrom(w3)
|
||||
graph.initializer.add().CopyFrom(b3)
|
||||
graph.initializer.add().CopyFrom(w5)
|
||||
graph.initializer.add().CopyFrom(b5)
|
||||
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
|
||||
onnx.save(model, test_model_path)
|
||||
|
||||
|
||||
class TestDataReader(CalibrationDataReader):
|
||||
"""Random Data Input Generator"""
|
||||
|
||||
def __init__(self):
|
||||
self.preprocess_flag = True
|
||||
self.enum_data_dicts = []
|
||||
self.count = 2
|
||||
self.input_data_list = []
|
||||
for _ in range(self.count):
|
||||
self.input_data_list.append(np.random.normal(0, 0.33, [1, 3, 1, 3]).astype(np.float32))
|
||||
|
||||
def get_next(self):
|
||||
if self.preprocess_flag:
|
||||
self.preprocess_flag = False
|
||||
input_name = "input"
|
||||
self.enum_data_dicts = iter([{input_name: input_data} for input_data in self.input_data_list])
|
||||
return next(self.enum_data_dicts, None)
|
||||
|
||||
def rewind(self):
|
||||
self.preprocess_flag = True
|
||||
|
||||
|
||||
class TestSaveActivations(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls._tmp_model_dir = tempfile.TemporaryDirectory(prefix="test_save_activations.")
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
cls._tmp_model_dir.cleanup()
|
||||
|
||||
def test_saved_tensors_match_internal_tensors(self):
|
||||
test_model_path = str(Path(self._tmp_model_dir.name) / "augmented_model.onnx")
|
||||
construct_test_model1(test_model_path)
|
||||
data_reader = TestDataReader()
|
||||
|
||||
aug_model = modify_model_output_intermediate_tensors(test_model_path)
|
||||
augmented_model_path = str(Path(self._tmp_model_dir.name).joinpath("augmented_test_model_1.onnx"))
|
||||
|
||||
onnx.save(
|
||||
aug_model,
|
||||
augmented_model_path,
|
||||
save_as_external_data=False,
|
||||
)
|
||||
|
||||
tensor_dict = collect_activations(augmented_model_path, data_reader)
|
||||
|
||||
# run original model and compare the tensors
|
||||
sess_options = onnxruntime.SessionOptions()
|
||||
sess_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL
|
||||
infer_session = onnxruntime.InferenceSession(
|
||||
test_model_path,
|
||||
sess_options=sess_options,
|
||||
providers=["CPUExecutionProvider"],
|
||||
)
|
||||
|
||||
data_reader.rewind()
|
||||
oracle_outputs = []
|
||||
for input_d in data_reader:
|
||||
oracle_outputs.append(infer_session.run(None, input_d))
|
||||
|
||||
output_dict = {}
|
||||
output_info = infer_session.get_outputs()
|
||||
for batch in oracle_outputs:
|
||||
for output, output_data in zip(output_info, batch):
|
||||
output_dict.setdefault(output.name, []).append(output_data)
|
||||
|
||||
for output_name, model_outputs in output_dict.items():
|
||||
test_outputs = tensor_dict[output_name]
|
||||
for expected, actual in zip(model_outputs, test_outputs):
|
||||
exp = expected.reshape(-1)
|
||||
act = actual.reshape(-1)
|
||||
np.testing.assert_equal(exp, act)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Loading…
Reference in a new issue