fix the missing return in _get_quantize_input_nodes and format code with yapf (#3199)

* fix the missing return for function _get_quantize_input_nodes

* format quantization code with yapf
This commit is contained in:
Yufeng Li 2020-03-13 09:28:41 -07:00 committed by GitHub
parent d99554bea1
commit c69194ec4c
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 612 additions and 292 deletions

View file

@ -21,6 +21,7 @@ import re
import subprocess
import json
def augment_graph(model):
'''
Adds ReduceMin and ReduceMax nodes to all Conv and MatMul nodes in
@ -38,22 +39,33 @@ def augment_graph(model):
input_name = node.output[0]
# Adding ReduceMin nodes
reduce_min_name = node.name + '_ReduceMin'
reduce_min_node = onnx.helper.make_node('ReduceMin', [input_name],
[input_name + '_ReduceMin'], reduce_min_name, keepdims=0)
reduce_min_node = onnx.helper.make_node(
'ReduceMin', [input_name], [input_name + '_ReduceMin'],
reduce_min_name,
keepdims=0)
added_nodes.append(reduce_min_node)
added_outputs.append(helper.make_tensor_value_info(reduce_min_node.output[0], TensorProto.FLOAT, ()))
added_outputs.append(
helper.make_tensor_value_info(reduce_min_node.output[0],
TensorProto.FLOAT, ()))
# Adding ReduceMax nodes
reduce_max_name = node.name + '_ReduceMax'
reduce_max_node = onnx.helper.make_node('ReduceMax', [input_name],
[input_name + '_ReduceMax'], reduce_max_name, keepdims=0)
reduce_max_node = onnx.helper.make_node(
'ReduceMax', [input_name], [input_name + '_ReduceMax'],
reduce_max_name,
keepdims=0)
added_nodes.append(reduce_max_node)
added_outputs.append(helper.make_tensor_value_info(reduce_max_node.output[0], TensorProto.FLOAT, ()))
added_outputs.append(
helper.make_tensor_value_info(reduce_max_node.output[0],
TensorProto.FLOAT, ()))
model.graph.node.extend(added_nodes)
model.graph.output.extend(added_outputs)
return model
# Using augmented outputs to generate inputs to quantize.py
def get_intermediate_outputs(model_path, session, inputs, calib_mode='naive'):
'''
Gather intermediate model outputs after running inference
@ -69,33 +81,52 @@ def get_intermediate_outputs(model_path, session, inputs, calib_mode='naive'):
return: dictionary mapping added node names to (ReduceMin, ReduceMax) pairs
'''
model = onnx.load(model_path)
num_model_outputs = len(model.graph.output) # number of outputs in original model
# number of outputs in original model
num_model_outputs = len(model.graph.output)
num_inputs = len(inputs)
input_name = session.get_inputs()[0].name
intermediate_outputs = [session.run([], {input_name: inputs[i]}) for i in range(num_inputs)]
intermediate_outputs = [
session.run([], {input_name: inputs[i]}) for i in range(num_inputs)
]
# Creating dictionary with output results from multiple test inputs
node_output_names = [session.get_outputs()[i].name for i in range(len(intermediate_outputs[0]))]
output_dicts = [dict(zip(node_output_names, intermediate_outputs[i])) for i in range(num_inputs)]
node_output_names = [
session.get_outputs()[i].name
for i in range(len(intermediate_outputs[0]))
]
output_dicts = [
dict(zip(node_output_names, intermediate_outputs[i]))
for i in range(num_inputs)
]
merged_dict = {}
for d in output_dicts:
for k, v in d.items():
merged_dict.setdefault(k, []).append(v)
added_node_output_names = node_output_names[num_model_outputs:]
node_names = [added_node_output_names[i].rpartition('_')[0] for i in range(0, len(added_node_output_names), 2)] # output names
node_names = [
added_node_output_names[i].rpartition('_')[0]
for i in range(0, len(added_node_output_names), 2)
] # output names
# Characterizing distribution of a node's values across test data sets
clean_merged_dict = dict((i, merged_dict[i]) for i in merged_dict if i != list(merged_dict.keys())[0])
clean_merged_dict = dict((i, merged_dict[i]) for i in merged_dict
if i != list(merged_dict.keys())[0])
if calib_mode == 'naive':
pairs = [tuple([float(min(clean_merged_dict[added_node_output_names[i]])),
float(max(clean_merged_dict[added_node_output_names[i+1]]))])
for i in range(0, len(added_node_output_names), 2)]
pairs = [
tuple([
float(min(clean_merged_dict[added_node_output_names[i]])),
float(max(clean_merged_dict[added_node_output_names[i + 1]]))
]) for i in range(0, len(added_node_output_names), 2)
]
else:
raise ValueError('Unknown value for calib_mode. Currently only naive mode is supported.')
raise ValueError(
'Unknown value for calib_mode. Currently only naive mode is supported.'
)
final_dict = dict(zip(node_names, pairs))
return final_dict
def calculate_scale_zeropoint(node, next_node, rmin, rmax):
zp_and_scale = []
# adjust rmin and rmax such that 0 is included in the range. This is required
@ -117,7 +148,7 @@ def calculate_scale_zeropoint(node, next_node, rmin, rmax):
if rmin < 0:
rmin = 0
scale = np.float32((rmax - rmin)/255 if rmin != rmax else 1)
scale = np.float32((rmax - rmin) / 255 if rmin != rmax else 1)
initial_zero_point = (0 - rmin) / scale
zero_point = np.uint8(round(max(0, min(255, initial_zero_point))))
@ -125,6 +156,7 @@ def calculate_scale_zeropoint(node, next_node, rmin, rmax):
zp_and_scale.append(scale)
return zp_and_scale
def calculate_quantization_params(model, quantization_thresholds):
'''
Given a model and quantization thresholds, calculates the quantization params.
@ -145,16 +177,20 @@ def calculate_quantization_params(model, quantization_thresholds):
{
"param_name": [zero_point, scale]
}
'''
'''
if quantization_thresholds is None:
raise ValueError('quantization thresholds is required to calculate quantization params (zero point and scale)')
raise ValueError(
'quantization thresholds is required to calculate quantization params (zero point and scale)'
)
quantization_params = {}
for index, node in enumerate(model.graph.node):
node_output_name = node.output[0]
if node_output_name in quantization_thresholds:
node_thresholds = quantization_thresholds[node_output_name]
node_params = calculate_scale_zeropoint(node, model.graph.node[index+1], node_thresholds[0], node_thresholds[1])
node_params = calculate_scale_zeropoint(
node, model.graph.node[index + 1], node_thresholds[0],
node_thresholds[1])
quantization_params[node_output_name] = node_params
return quantization_params
@ -185,21 +221,38 @@ def load_pb_file(data_file_name, size_limit, samples, channels, height, width):
inputs = inputs[:size_limit]
dataset_size = size_limit
inputs = inputs.reshape(dataset_size, samples, channels, height, width)
inputs = inputs.reshape(dataset_size, samples, channels, height,
width)
except:
sys.exit("Input .pb file contains incorrect input size. \nThe required size is: (%s). The real size is: (%s)"
%((dataset_size, samples, channels, height, width), shape))
sys.exit(
"Input .pb file contains incorrect input size. \nThe required size is: (%s). The real size is: (%s)"
% ((dataset_size, samples, channels, height, width), shape))
return inputs
def main():
# Parsing command-line arguments
parser = argparse.ArgumentParser(description='parsing model and test data set paths')
parser = argparse.ArgumentParser(
description='parsing model and test data set paths')
parser.add_argument('--model_path', required=True)
parser.add_argument('--dataset_path', required=True)
parser.add_argument('--output_model_path', type=str, default='calibrated_quantized_model.onnx')
parser.add_argument('--dataset_size', type=int, default=0, help="Number of images or tensors to load. Default is 0 which means all samples")
parser.add_argument('--data_preprocess', type=str, required=True, choices=['preprocess_method1', 'preprocess_method2', 'None'], help="Refer to Readme.md for guidance on choosing this option.")
parser.add_argument('--output_model_path',
type=str,
default='calibrated_quantized_model.onnx')
parser.add_argument(
'--dataset_size',
type=int,
default=0,
help=
"Number of images or tensors to load. Default is 0 which means all samples"
)
parser.add_argument(
'--data_preprocess',
type=str,
required=True,
choices=['preprocess_method1', 'preprocess_method2', 'None'],
help="Refer to Readme.md for guidance on choosing this option.")
args = parser.parse_args()
model_path = args.model_path
output_model_path = args.output_model_path
@ -219,16 +272,24 @@ def main():
# Generating inputs for quantization
if args.data_preprocess == "None":
inputs = load_pb_file(images_folder, args.dataset_size, samples, channels, height, width)
inputs = load_pb_file(images_folder, args.dataset_size, samples,
channels, height, width)
else:
inputs = load_batch(images_folder, height, width, size_limit, args.data_preprocess)
inputs = load_batch(images_folder, height, width, size_limit,
args.data_preprocess)
print(inputs.shape)
dict_for_quantization = get_intermediate_outputs(model_path, session, inputs, calib_mode)
quantization_params_dict = calculate_quantization_params(model, quantization_thresholds=dict_for_quantization)
calibrated_quantized_model = quantize(onnx.load(model_path), quantization_mode=QuantizationMode.QLinearOps, quantization_params=quantization_params_dict)
dict_for_quantization = get_intermediate_outputs(model_path, session,
inputs, calib_mode)
quantization_params_dict = calculate_quantization_params(
model, quantization_thresholds=dict_for_quantization)
calibrated_quantized_model = quantize(
onnx.load(model_path),
quantization_mode=QuantizationMode.QLinearOps,
quantization_params=quantization_params_dict)
onnx.save(calibrated_quantized_model, output_model_path)
print("Calibrated, quantized model saved.")
if __name__ == '__main__':
main()
main()

View file

@ -10,16 +10,20 @@ import os
import sys
import numpy as np
def set_preprocess(preprocess_func_name):
'''
Set up the data preprocess function name and function dict.
parameter preprocess_func_name: name of the preprocess function
return: function pointer
'''
funcdict = {'preprocess_method1': preprocess_method1,
'preprocess_method2': preprocess_method2}
funcdict = {
'preprocess_method1': preprocess_method1,
'preprocess_method2': preprocess_method2
}
return funcdict[preprocess_func_name]
def preprocess_method1(image_filepath, height, width):
'''
Resizes image to NCHW format. Image is scaled to range [-1, 1].
@ -30,12 +34,13 @@ def preprocess_method1(image_filepath, height, width):
return: matrix characterizing image
'''
pillow_img = Image.open(image_filepath).resize((width, height))
input_data = np.float32(pillow_img)/127.5 - 1.0 # normalization
input_data -= np.mean(input_data) # normalization
input_data = np.float32(pillow_img) / 127.5 - 1.0 # normalization
input_data -= np.mean(input_data) # normalization
nhwc_data = np.expand_dims(input_data, axis=0)
nchw_data = nhwc_data.transpose(0, 3, 1, 2) # ONNX Runtime standard
nchw_data = nhwc_data.transpose(0, 3, 1, 2) # ONNX Runtime standard
return nchw_data
def preprocess_method2(image_filepath, height, width):
'''
Resizes and normalizes image to NCHW format.
@ -46,13 +51,18 @@ def preprocess_method2(image_filepath, height, width):
return: matrix characterizing image
'''
pillow_img = Image.open(image_filepath).resize((width, height))
input_data = np.float32(pillow_img) - np.array([123.68, 116.78, 103.94], dtype=np.float32)
input_data = np.float32(pillow_img) - \
np.array([123.68, 116.78, 103.94], dtype=np.float32)
nhwc_data = np.expand_dims(input_data, axis=0)
nchw_data = nhwc_data.transpose(0, 3, 1, 2) # ONNX Runtime standard
nchw_data = nhwc_data.transpose(0, 3, 1, 2) # ONNX Runtime standard
return nchw_data
def load_batch(images_folder, height, width, preprocess_func_name, size_limit=0):
def load_batch(images_folder,
height,
width,
preprocess_func_name,
size_limit=0):
'''
Loads a batch of images
parameter images_folder: path to folder storing images
@ -68,11 +78,13 @@ def load_batch(images_folder, height, width, preprocess_func_name, size_limit=0)
else:
batch_filenames = image_names
unconcatenated_batch_data = []
preprocess_func = set_preprocess(preprocess_func_name)
for image_name in batch_filenames:
image_filepath = images_folder + '/' + image_name
nchw_data = preprocess_func(image_filepath, height, width)
unconcatenated_batch_data.append(nchw_data)
batch_data = np.concatenate(np.expand_dims(unconcatenated_batch_data, axis=0), axis=0)
return batch_data
batch_data = np.concatenate(np.expand_dims(unconcatenated_batch_data,
axis=0),
axis=0)
return batch_data

File diff suppressed because it is too large Load diff

View file

@ -16,8 +16,8 @@ import calibrate
import numpy as np
from onnx import numpy_helper
class TestCalibrate(unittest.TestCase):
class TestCalibrate(unittest.TestCase):
def test_augment_graph(self):
# Creating graph
A = helper.make_tensor_value_info('A', TensorProto.FLOAT, [1, 1, 5, 5])
@ -26,10 +26,15 @@ class TestCalibrate(unittest.TestCase):
D = helper.make_tensor_value_info('D', TensorProto.FLOAT, [1, 1, 5, 5])
E = helper.make_tensor_value_info('E', TensorProto.FLOAT, [1, 1, 5, 1])
F = helper.make_tensor_value_info('F', TensorProto.FLOAT, [1, 1, 5, 1])
conv_node = onnx.helper.make_node('Conv', ['A', 'B'], ['C'], name='Conv', kernel_shape=[3, 3], pads=[1, 1, 1, 1])
conv_node = onnx.helper.make_node('Conv', ['A', 'B'], ['C'],
name='Conv',
kernel_shape=[3, 3],
pads=[1, 1, 1, 1])
clip_node = onnx.helper.make_node('Clip', ['C'], ['D'], name='Clip')
matmul_node = onnx.helper.make_node('MatMul', ['D', 'E'], ['F'], name='MatMul')
graph = helper.make_graph([conv_node, clip_node, matmul_node], 'test_graph', [A, B, E], [F])
matmul_node = onnx.helper.make_node('MatMul', ['D', 'E'], ['F'],
name='MatMul')
graph = helper.make_graph([conv_node, clip_node, matmul_node],
'test_graph', [A, B, E], [F])
model = helper.make_model(graph)
onnx.save(model, 'test_model.onnx')
@ -38,12 +43,23 @@ class TestCalibrate(unittest.TestCase):
onnx.save(augmented_model, 'augmented_test_model.onnx')
# Checking if each added ReduceMin and ReduceMax node and its output exists
augmented_model_node_names = [node.name for node in augmented_model.graph.node]
augmented_model_outputs = [output.name for output in augmented_model.graph.output]
added_node_names = ['Conv_ReduceMin', 'Conv_ReduceMax', 'MatMul_ReduceMin', 'MatMul_ReduceMax']
added_outputs = ['C_ReduceMin', 'C_ReduceMax', 'F_ReduceMin', 'F_ReduceMax']
self.assertEqual(len(augmented_model_node_names), 7) # original 3 nodes with 4 added ones
self.assertEqual(len(augmented_model_outputs), 5) # original single graph output with 4 added ones
augmented_model_node_names = [
node.name for node in augmented_model.graph.node
]
augmented_model_outputs = [
output.name for output in augmented_model.graph.output
]
added_node_names = [
'Conv_ReduceMin', 'Conv_ReduceMax', 'MatMul_ReduceMin',
'MatMul_ReduceMax'
]
added_outputs = [
'C_ReduceMin', 'C_ReduceMax', 'F_ReduceMin', 'F_ReduceMax'
]
# original 3 nodes with 4 added ones
self.assertEqual(len(augmented_model_node_names), 7)
# original single graph output with 4 added ones
self.assertEqual(len(augmented_model_outputs), 5)
for name in added_node_names:
self.assertTrue(name in augmented_model_node_names)
for output in added_outputs:
@ -55,9 +71,12 @@ class TestCalibrate(unittest.TestCase):
I = helper.make_tensor_value_info('I', TensorProto.FLOAT, [1, 1, 5, 1])
J = helper.make_tensor_value_info('J', TensorProto.FLOAT, [1, 1, 5, 1])
K = helper.make_tensor_value_info('K', TensorProto.FLOAT, [1, 1, 5, 1])
matmul_node_1 = onnx.helper.make_node('MatMul', ['G', 'H'], ['I'], name='MatMul_1')
matmul_node_2 = onnx.helper.make_node('MatMul', ['I', 'J'], ['K'], name='MatMul_2')
graph = helper.make_graph([matmul_node_1, matmul_node_2], 'test_graph_order', [G, H, J], [K])
matmul_node_1 = onnx.helper.make_node('MatMul', ['G', 'H'], ['I'],
name='MatMul_1')
matmul_node_2 = onnx.helper.make_node('MatMul', ['I', 'J'], ['K'],
name='MatMul_2')
graph = helper.make_graph([matmul_node_1, matmul_node_2],
'test_graph_order', [G, H, J], [K])
model = helper.make_model(graph)
onnx.save(model, 'test_model_matmul.onnx')
@ -65,7 +84,9 @@ class TestCalibrate(unittest.TestCase):
augmented_model = calibrate.augment_graph(model)
onnx.save(augmented_model, 'augmented_test_model_matmul.onnx')
augmented_model_outputs = [output.name for output in augmented_model.graph.output]
augmented_model_outputs = [
output.name for output in augmented_model.graph.output
]
self.assertEqual(augmented_model_outputs[0], 'K')
self.assertEqual(augmented_model_outputs[1], 'I_ReduceMin')
self.assertEqual(augmented_model_outputs[2], 'I_ReduceMax')
@ -76,20 +97,28 @@ class TestCalibrate(unittest.TestCase):
images_folder = 'test_images'
session = onnxruntime.InferenceSession('augmented_test_model.onnx')
(samples, channels, height, width) = session.get_inputs()[0].shape
batch_data = calibrate.load_batch(images_folder, height, width, preprocess_func_name="preprocess_method1")
self.assertEqual(len(batch_data.shape), 5) # for 2D images like the ones in test_images
batch_data = calibrate.load_batch(
images_folder,
height,
width,
preprocess_func_name="preprocess_method1")
# for 2D images like the ones in test_images
self.assertEqual(len(batch_data.shape), 5)
self.assertEqual(batch_data.shape[0], len(os.listdir(images_folder)))
self.assertEqual(batch_data.shape[2], 3) # checking for 3 channels for colored image
self.assertEqual(batch_data.shape[3], height) # checking if resized height is correct
self.assertEqual(batch_data.shape[4], width) # checking if resized width is correct
# checking for 3 channels for colored image
self.assertEqual(batch_data.shape[2], 3)
# checking if resized height is correct
self.assertEqual(batch_data.shape[3], height)
# checking if resized width is correct
self.assertEqual(batch_data.shape[4], width)
def test_load_pb(self):
numpy_array = np.random.randn(3, 1, 3, 5, 5).astype(np.float32)
tensor = numpy_helper.from_array(numpy_array)
test_file_name = 'test_tensor.pb'
with open(test_file_name, 'wb') as f:
f.write(tensor.SerializeToString())
# test size_limit < than number of samples in data set
# expecting to load size_limit number of samples
batch_data = calibrate.load_pb_file('test_tensor.pb', 2, 1, 3, 5, 5)
@ -97,9 +126,9 @@ class TestCalibrate(unittest.TestCase):
self.assertEqual(batch_data.shape[0], 2)
self.assertEqual(batch_data.shape[2], 3)
self.assertEqual(batch_data.shape[3], 5)
self.assertEqual(batch_data.shape[4], 5)
self.assertEqual(batch_data.shape[4], 5)
# test size_limit == 0
# test size_limit == 0
# expecting to load all samples
batch_data = calibrate.load_pb_file('test_tensor.pb', 0, 1, 3, 5, 5)
self.assertEqual(len(batch_data.shape), 5)
@ -120,9 +149,9 @@ class TestCalibrate(unittest.TestCase):
try:
os.remove('test_tensor.pb')
except:
print("Warning: Trying to remove test file {} failed.".format(test_file_name))
print("Warning: Trying to remove test file {} failed.".format(
test_file_name))
if __name__ == '__main__':
unittest.main()
unittest.main()