def dump_tensor(self, activation=True, weight=False):
if "QuantizeLinear" in [
node.op_type for node in self.model.graph.node
]:
self.augment_nodes = ["DequantizeLinear"]
self.already_quantized = True
activation_only = not weight
self.augment_graph(activation_only=activation_only, output_only=True)
_, output_dicts_list = self.get_intermediate_outputs()
output_dicts = {}
for output_dicts_iter in output_dicts_list:
for output_name in output_dicts_iter:
if output_name not in output_dicts:
output_dicts[output_name] = []
output_dicts[output_name].append(
output_dicts_iter[output_name])
iters = len(output_dicts_list)
map_node_activation = [{} for _ in range(iters)]
map_node_weight = {}
self.white_nodes = [
node.replace('_quant', '') for node in self.white_nodes
]
augmengted_wrapper = ONNXModel(self.augmented_model)
map_output = augmengted_wrapper.output_name_to_node
map_input = augmengted_wrapper.input_name_to_nodes
model_output_names = [t.name for t in self.model.graph.output]
model_initializer_names = [
t.name for t in self.model.graph.initializer
]
for tensor_name, tensors in output_dicts.items():
if tensor_name.endswith('_scale') or tensor_name.endswith(
'_zero_point'):
continue # don't dump scale and zero_point
if tensor_name in model_initializer_names:
nodes = [node for node in map_input[tensor_name] \
if node.name.replace('_quant', '') in self.white_nodes]
else:
nodes = [map_output[tensor_name]]
for node in nodes:
node_name = node.name.replace('_quant', '')
if tensor_name in model_output_names and node_name not in self.white_nodes:
continue
while node_name not in self.white_nodes:
node = augmengted_wrapper.get_parents(
node, output_name_to_node=map_output)[0]
node_name = node.name.replace('_quant', '')
if node_name not in map_node_weight:
map_node_weight[node_name] = {}
if tensor_name not in model_initializer_names:
for i in range(iters):
map_node_activation[i][node_name] = \
{tensor_name.replace('_quantized', ''): tensors[i]}
else:
map_node_weight[node_name].update({tensor_name.replace('_quantized', ''): \
tensors[0]})
dumped_tensors_map = {}
if weight:
dumped_tensors_map.update({"weight": map_node_weight})
if activation:
dumped_tensors_map.update({"activation": map_node_activation})
return dumped_tensors_map
class TestOnnxModel(unittest.TestCase):
def setUp(self):
# Relu
# | \
# Conv \
# | \
# Relu |
# | Conv
# Conv /
# \ /
# |
# Add
input0 = helper.make_tensor_value_info('input0', TensorProto.FLOAT,
[1, 3, 1, 3])
output = helper.make_tensor_value_info('output', TensorProto.FLOAT,
[1, 3, 1, 3])
X1_weight = generate_input_initializer([3, 3, 1, 1], np.float32,
'X1_weight')
X1_bias = generate_input_initializer([3], np.float32, 'X1_bias')
X3_weight = generate_input_initializer([3, 3, 1, 1], np.float32,
'X3_weight')
X3_bias = generate_input_initializer([3], np.float32, 'X3_bias')
X5_weight = generate_input_initializer([3, 3, 1, 1], np.float32,
'X5_weight')
X5_bias = generate_input_initializer([3], np.float32, 'X5_bias')
relu_node_1 = onnx.helper.make_node('Relu', ['input0'], ['X1'],
name='Relu1')
conv_node_1 = onnx.helper.make_node('Conv',
['X1', 'X1_weight', 'X1_bias'],
['X2'],
name='Conv1')
relu_node_2 = onnx.helper.make_node('Relu', ['X2'], ['X3'],
name='Relu2')
conv_node_2 = onnx.helper.make_node('Conv',
['X3', 'X3_weight', 'X3_bias'],
['X4'],
name='Conv2')
conv_node_3 = onnx.helper.make_node('Conv',
['X1', 'X5_weight', 'X5_bias'],
['X5'],
name='Conv3')
add_node = onnx.helper.make_node('Add', ['X4', 'X5'], ['output'],
name='Add')
graph = helper.make_graph([
relu_node_1, conv_node_1, relu_node_2, conv_node_2, conv_node_3,
add_node
], 'test_graph_6', [input0], [output])
graph.initializer.add().CopyFrom(X1_weight)
graph.initializer.add().CopyFrom(X1_bias)
graph.initializer.add().CopyFrom(X3_weight)
graph.initializer.add().CopyFrom(X3_bias)
graph.initializer.add().CopyFrom(X5_weight)
graph.initializer.add().CopyFrom(X5_bias)
model = helper.make_model(graph)
test_model_path = './test_model_6.onnx'
onnx.save(model, test_model_path)
model = onnx.load(test_model_path)
self.model = ONNXModel(model)
# QuantizeLinear
# |
# QLinearConv
# |
# DequantizeLinear
A = helper.make_tensor_value_info('A', TensorProto.FLOAT, [1, 1, 5, 5])
A_scale = helper.make_tensor_value_info('A_scale', TensorProto.FLOAT,
[1])
a_scale = generate_input_initializer([1], np.float32, 'A_scale')
A_zo = helper.make_tensor_value_info('A_zero_point', TensorProto.INT8,
[1])
a_zero_point = generate_input_initializer([1], np.int8, 'A_zero_point')
B_scale = helper.make_tensor_value_info('B_scale', TensorProto.FLOAT,
[1])
b_scale = generate_input_initializer([1], np.float32, 'B_scale')
B_zo = helper.make_tensor_value_info('B_zero_point', TensorProto.INT8,
[1])
b_zero_point = generate_input_initializer([1], np.int8, 'B_zero_point')
C = helper.make_tensor_value_info('C', TensorProto.INT8, [1, 1, 5, 5])
c = generate_input_initializer([1, 1, 5, 5], np.int8, 'C')
C_scale = helper.make_tensor_value_info('C_scale', TensorProto.FLOAT,
[1])
c_scale = generate_input_initializer([1], np.float32, 'C_scale')
C_zo = helper.make_tensor_value_info('C_zero_point', TensorProto.INT8,
[1])
c_zero_point = generate_input_initializer([1], np.int8, 'C_zero_point')
E = helper.make_tensor_value_info('E', TensorProto.INT32, [1])
e = generate_input_initializer([1], np.int32, 'E')
D_scale = helper.make_tensor_value_info('D_scale', TensorProto.FLOAT,
[1])
d_scale = generate_input_initializer([1], np.float32, 'D_scale')
D_zo = helper.make_tensor_value_info('D_zero_point', TensorProto.INT8,
[1])
d_zero_point = generate_input_initializer([1], np.int8, 'D_zero_point')
D = helper.make_tensor_value_info('D', TensorProto.FLOAT, [1, 1, 5, 5])
quantize_node = onnx.helper.make_node('QuantizeLinear',
['A', 'A_scale', 'A_zero_point'],
['B_quantized'],
name='A_QuantizeLinear')
conv_node = onnx.helper.make_node('QLinearConv', [
'B_quantized', 'B_scale', 'B_zero_point', 'C_quantized', 'C_scale',
'C_zero_point', 'D_scale', 'D_zero_point', 'E'
], ['D_quantized'],
name='conv_quant',
kernel_shape=[3, 3],
pads=[1, 1, 1, 1])
dequantize_node = onnx.helper.make_node(
'DequantizeLinear', ['D_quantized', 'D_scale', 'D_zero_point'],
['D'],
name='D_DequantizeLinear')
graph = helper.make_graph(
[quantize_node, conv_node, dequantize_node], 'test_graph_7',
[A, A_scale, A_zo, C, C_scale, C_zo, E, D_scale, D_zo], [D])
graph.initializer.add().CopyFrom(a_scale)
graph.initializer.add().CopyFrom(a_zero_point)
graph.initializer.add().CopyFrom(b_scale)
graph.initializer.add().CopyFrom(b_zero_point)
graph.initializer.add().CopyFrom(c)
graph.initializer.add().CopyFrom(c_scale)
graph.initializer.add().CopyFrom(c_zero_point)
graph.initializer.add().CopyFrom(e)
graph.initializer.add().CopyFrom(d_scale)
graph.initializer.add().CopyFrom(d_zero_point)
model = helper.make_model(graph)
self.q_model = ONNXModel(model)
def test_nodes(self):
self.assertEqual(len(self.model.nodes()), 6)
nodes_name = [node.name for node in self.model.nodes()]
nodes = ["Relu1", "Conv1", "Relu2", "Conv2", "Conv3", "Add"]
for node in nodes:
self.assertTrue(node in nodes_name)
def test_initializer(self):
self.assertEqual(len(self.model.initializer()), 6)
inits_name = [init.name for init in self.model.initializer()]
inits = [
'X1_weight', 'X1_bias', 'X3_weight', 'X3_bias', 'X5_weight',
'X5_bias'
]
for init in inits:
self.assertTrue(init in inits_name)
def test_remove_node(self):
for node in self.model.nodes():
if node.op_type == "Add":
self.model.remove_node(node)
self.assertEqual(len(self.model.nodes()), 5)
nodes_name = [node.name for node in self.model.nodes()]
nodes = ["Relu1", "Conv1", "Relu2", "Conv2", "Conv3"]
for node in nodes:
self.assertTrue(node in nodes_name)
def test_remove_nodes(self):
nodes_to_remove = []
for node in self.model.nodes():
if node.name == "Conv3" or node.name == "Add":
nodes_to_remove.append(node)
self.model.remove_nodes(nodes_to_remove)
self.assertEqual(len(self.model.nodes()), 4)
nodes_name = [node.name for node in self.model.nodes()]
nodes = ["Relu1", "Conv1", "Relu2", "Conv2"]
for node in nodes:
self.assertTrue(node in nodes_name)
def test_add_node(self):
node_to_add = onnx.helper.make_node('Relu', ['output'], ['output1'],
keepdims=0)
self.model.add_node(node_to_add)
last_node = self.model.nodes()[-1]
self.assertEqual(last_node.op_type, 'Relu')
def test_add_nodes(self):
nodes_to_add = []
for i in range(2):
node_to_add = onnx.helper.make_node(
'Relu', ["add_node{}_input".format(str(i))],
["add_node{}_output".format(str(i))],
keepdims=0)
nodes_to_add.append(node_to_add)
self.model.add_nodes(nodes_to_add)
self.assertEqual(self.model.nodes()[-1].input, ['add_node1_input'])
self.assertEqual(self.model.nodes()[-2].input, ['add_node0_input'])
self.assertEqual(self.model.nodes()[-1].output, ['add_node1_output'])
self.assertEqual(self.model.nodes()[-2].output, ['add_node0_output'])
def test_get_initializer(self):
inits = [
'X1_weight', 'X1_bias', 'X3_weight', 'X3_bias', 'X5_weight',
'X5_bias'
]
for init in inits:
self.assertIsNotNone(self.model.get_initializer(init))
def test_remove_initializer(self):
for init in self.model.initializer():
if init.name == "X1_weight":
self.model.remove_initializer(init)
self.assertEqual(len(self.model.initializer()), 5)
inits_name = [init.name for init in self.model.initializer()]
inits = ['X1_bias', 'X3_weight', 'X3_bias', 'X5_weight', 'X5_bias']
for init in inits:
self.assertTrue(init in inits_name)
def test_remove_initializers(self):
init_to_remove = []
for init in self.model.initializer():
if "bias" in init.name:
init_to_remove.append(init)
self.model.remove_initializers(init_to_remove)
self.assertEqual(len(self.model.initializer()), 3)
inits_name = [init.name for init in self.model.initializer()]
inits = ['X1_weight', 'X3_weight', 'X5_weight']
for init in inits:
self.assertTrue(init in inits_name)
def test_input_name_to_nodes(self):
self.assertEqual(len(self.model.input_name_to_nodes), 12)
ipts_name = [name for name in self.model.input_name_to_nodes]
ipts = [
'input0', 'X1', 'X2', 'X3', 'X3_weight', 'X3_bias', 'X5_weight',
'X5_bias', 'X4', 'X5'
]
for ipt in ipts:
self.assertTrue(ipt in ipts_name)
def test_output_name_to_node(self):
self.assertEqual(len(self.model.output_name_to_node), 6)
opts_name = [name for name in self.model.output_name_to_node]
opts = ['X1', 'X2', 'X3', 'X4', 'X5', 'output']
for opt in opts:
self.assertTrue(opt in opts_name)
def test_get_children(self):
for node in self.model.nodes():
if node.name == "Relu1":
children = self.model.get_children(node)
self.assertEqual(len(children), 2)
children_name = [child.name for child in children]
names = ["Conv1", "Conv3"]
for name in names:
self.assertTrue(name in children_name)
def test_get_parents(self):
for node in self.model.nodes():
if node.op_type == "Add":
parents = self.model.get_parents(node)
self.assertEqual(len(parents), 2)
parents_name = [parent.name for parent in parents]
names = ["Conv2", "Conv3"]
for name in names:
self.assertTrue(name in parents_name)
def test_get_parent(self):
for node in self.model.nodes():
if node.op_type == "Add":
node_to_get_parent = node
parent = self.model.get_parent(node, 0)
self.assertEqual(parent.name, "Conv2")
parent = self.model.get_parent(node, 1)
self.assertEqual(parent.name, "Conv3")
parent = self.model.get_parent(node, 2)
self.assertIsNone(parent)
def test_find_nodes_by_initializer(self):
for init in self.model.initializer():
if init.name == "X1_weight":
initializer = init
nodes = self.model.find_nodes_by_initializer(self.model.graph(),
initializer)
self.assertEqual(len(nodes), 1)
self.assertEqual(nodes[0].name, "Conv1")
def test_get_scale_zo(self):
input_scale, input_zo = self.q_model.get_scale_zo('B_quantized')
weight_scale, weight_zo = self.q_model.get_scale_zo('C_quantized')
bias_scale, bias_zo = self.q_model.get_scale_zo('E')
def test_save(self):
self.model.save_model_to_file('./test_model_6.onnx',
use_external_data_format=True)