mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-30 20:18:08 +00:00
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:
parent
d99554bea1
commit
c69194ec4c
4 changed files with 612 additions and 292 deletions
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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
|
|
@ -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()
|
||||
|
|
|
|||
Loading…
Reference in a new issue