forked from nudt_dsp/netrans
267 lines
9.0 KiB
Python
267 lines
9.0 KiB
Python
import os
|
|
import sys
|
|
import numpy as np
|
|
import tensorflow
|
|
import flatbuffers
|
|
main_version = tensorflow.__version__.split('.')[0]
|
|
if int(main_version) == 2:
|
|
import tensorflow.compat.v1 as tf
|
|
else:
|
|
import tensorflow as tf
|
|
from argparse import ArgumentParser
|
|
|
|
def arguments():
|
|
parser = ArgumentParser(description='TFLite Runner')
|
|
parser.add_argument(
|
|
'-m',
|
|
'--model',
|
|
required=True,
|
|
help='TFLite model to be executed',
|
|
)
|
|
parser.add_argument(
|
|
'-i',
|
|
'--input',
|
|
required=True,
|
|
help='TFLite model input file, each input path split by space',
|
|
)
|
|
parser.add_argument(
|
|
'--input-mean',
|
|
help='Input channel mean value list, every value split by space. Each mean value list split by #'
|
|
'Each entry in the list should match an entry in \'--input\'',
|
|
)
|
|
parser.add_argument(
|
|
'--input-std',
|
|
help='Input standard deviation. Each std split by #'
|
|
'Each entry in the list should match an entry in \'--input\'',
|
|
)
|
|
parser.add_argument(
|
|
'--dump-float',
|
|
action='store_true',
|
|
help='Force dump the output tensor as float32',
|
|
)
|
|
parser.add_argument(
|
|
'--detect-all',
|
|
action='store_true',
|
|
help="Dump all tensors in graph"
|
|
)
|
|
parser.add_argument(
|
|
'--output-dir',
|
|
default='./tflite_dump',
|
|
help="Specify the output dir of dump result, default by './tflite_dump'"
|
|
)
|
|
parser.add_argument(
|
|
'--detect-tensor',
|
|
default=None,
|
|
help="Dump specific tensor in graph"
|
|
)
|
|
args = parser.parse_args()
|
|
return args
|
|
|
|
def process_npy(input_detail, input_file):
|
|
_fdata = np.load(input_file)
|
|
if _fdata.dtype == np.float32 and input_detail['dtype'] != np.float32:
|
|
data = _quantize_data(_fdata, input_detail)
|
|
else:
|
|
data = _fdata
|
|
return data
|
|
|
|
def process_qtensor(input_detail, input_file):
|
|
data = np.loadtxt(input_file, dtype=input_detail['dtype'])
|
|
data = np.reshape(data, input_detail['shape'])
|
|
return data
|
|
|
|
def process_tensor(input_detail, input_file):
|
|
is_qtensor = input_file.endswith('.qnt.tensor')
|
|
if is_qtensor:
|
|
data = process_qtensor(input_detail, input_file)
|
|
else:
|
|
_fdata = np.loadtxt(input_file, dtype=np.float32)
|
|
_fdata = np.reshape(_fdata, input_detail['shape'])
|
|
if input_detail['dtype'] != np.float32:
|
|
data = _quantize_data(_fdata, input_detail)
|
|
else:
|
|
data = _fdata
|
|
return data
|
|
|
|
def process_image(input_detail, input_file, mean = None, std = None):
|
|
height = input_detail['shape'][1]
|
|
width = input_detail['shape'][2]
|
|
dtype = input_detail['dtype']
|
|
|
|
# decode the image file
|
|
img_tensor = tf.io.decode_image(tf.io.read_file(input_file))
|
|
img = img_tensor.numpy()
|
|
img_height = img.shape[0]
|
|
img_width = img.shape[1]
|
|
if img_height != height or img_width != width:
|
|
img_tensor = tf.image.resize(img_tensor, (height, width))
|
|
img = img_tensor.numpy()
|
|
_fdata = np.expand_dims(img, axis=0).astype(np.float32)
|
|
|
|
# process the mean and std
|
|
if mean is not None:
|
|
mean = np.array([float(m) for m in mean.split(' ')])
|
|
_fdata = _fdata - mean
|
|
if std is not None:
|
|
_std = float(std)
|
|
_fdata = _fdata / _std
|
|
|
|
# Convert float value to dtype value
|
|
if dtype != np.float32:
|
|
data = _quantize_data(_fdata, input_detail)
|
|
else:
|
|
data = _fdata.astype(dtype)
|
|
return data
|
|
|
|
def show_top5(data, n = 5):
|
|
res = data
|
|
res = np.reshape(res, [-1])
|
|
idx = np.argsort(res)[::-1]
|
|
# Top 5
|
|
print("Show Top 5")
|
|
if len(idx) < n:
|
|
n = len(idx)
|
|
for i in idx[:n]:
|
|
print('{}: {}'.format(i, res[i]))
|
|
|
|
def save_output(output_detail, data, output_dir=None):
|
|
shape = [str(s) for s in data.shape]
|
|
name = output_detail['name'].replace(':', '_').replace('@', '').replace('/', '_').replace(';', '_')
|
|
filename = name + '_' + '_'.join(shape) + '.tensor'
|
|
if output_dir is not None:
|
|
filename = os.path.join(output_dir, filename)
|
|
filename = filename[-255:] # handle file name too long
|
|
print("Dump result to file {}".format(filename))
|
|
data.tofile(filename, '\n')
|
|
|
|
def _quantize_data(data, tensor_detail):
|
|
quantize_parameter = tensor_detail['quantization_parameters']
|
|
scale = quantize_parameter['scales']
|
|
zp = quantize_parameter['zero_points']
|
|
dtype = tensor_detail['dtype']
|
|
# none quantized has empty scale and zp
|
|
if scale.size == 0 and zp.size == 0:
|
|
return data.astype(dtype)
|
|
else:
|
|
quantized_data = np.rint(data / scale).astype(dtype) + zp
|
|
return quantized_data.astype(dtype)
|
|
|
|
def _dequantize_data(data, tensor_detail):
|
|
quantize_parameter = tensor_detail['quantization_parameters']
|
|
scale = quantize_parameter['scales']
|
|
zp = quantize_parameter['zero_points']
|
|
# none quantized has empty scale and zp
|
|
if scale.size == 0 and zp.size == 0:
|
|
return data.astype(np.float32)
|
|
else:
|
|
dequantized_data = ((data - zp) * scale).astype(np.float32)
|
|
return dequantized_data
|
|
|
|
def _find_detect_tensor_details(interpreter, detect_tensor):
|
|
tensor_details = interpreter.get_tensor_details()
|
|
if detect_tensor == 'all':
|
|
return tensor_details
|
|
else:
|
|
details = list()
|
|
for name in detect_tensor:
|
|
for tensor in tensor_details:
|
|
if name == tensor['name']:
|
|
details.append(tensor)
|
|
return details
|
|
|
|
def dump_detect_tensors(interpreter, args, detect_tensor):
|
|
tensor_details = _find_detect_tensor_details(interpreter, detect_tensor)
|
|
for tensor in tensor_details:
|
|
data = interpreter.get_tensor(tensor['index'])
|
|
if args.dump_float and tensor['dtype'] != np.float32:
|
|
_data = data.astype(np.int32)
|
|
data = _dequantize_data(_data, tensor)
|
|
output_dir = args.output_dir
|
|
os.path.exists(output_dir) or os.mkdir(output_dir)
|
|
save_output(tensor, data, output_dir)
|
|
|
|
def post_process_output(interpreter, args):
|
|
output_details = interpreter.get_output_details()
|
|
for idx in range(len(output_details)):
|
|
data = interpreter.get_tensor(output_details[idx]['index'])
|
|
if args.dump_float and output_details[idx]['dtype'] != np.float32:
|
|
_data = data.astype(np.int32)
|
|
data = _dequantize_data(_data, output_details[idx])
|
|
save_output(output_details[idx], data, None)
|
|
show_top5(data)
|
|
|
|
def pre_process_input(interpreter, args):
|
|
input_files = args.input.split(' ')
|
|
input_details = interpreter.get_input_details()
|
|
if len(input_files) != len(input_details):
|
|
raise ValueError("This model need {} inputs, but model input file number is {}".
|
|
format(len(input_details), len(input_files)))
|
|
|
|
input_data = list()
|
|
if args.input_mean is not None:
|
|
mean_list = args.input_mean.split('#')
|
|
else:
|
|
mean_list = None
|
|
if args.input_std is not None:
|
|
std_list = args.input_std.split('#')
|
|
else:
|
|
std_list = None
|
|
for idx in range(len(input_files)):
|
|
shuffix = os.path.splitext(input_files[idx])[-1]
|
|
if shuffix == '.jpg' or shuffix == '.jpeg' or shuffix == '.bmp':
|
|
if mean_list is not None:
|
|
mean = mean_list[idx]
|
|
else:
|
|
mean = None
|
|
if std_list is not None:
|
|
std = std_list[idx]
|
|
else:
|
|
std = None
|
|
data = process_image(input_details[idx], input_files[idx], mean, std)
|
|
elif shuffix == '.tensor':
|
|
data = process_tensor(input_details[idx], input_files[idx])
|
|
elif shuffix == '.npy':
|
|
data = process_npy(input_details[idx], input_files[idx])
|
|
elif shuffix == '.qtensor':
|
|
data = process_qtensor(input_details[idx], input_files[idx])
|
|
else:
|
|
raise ValueError("TFLite Runner can't handle ({}) input file type".format(shuffix))
|
|
input_data.append(data)
|
|
|
|
return input_data
|
|
|
|
def inference(interpreter, inputs):
|
|
input_details = interpreter.get_input_details()
|
|
|
|
for idx in range(len(input_details)):
|
|
interpreter.set_tensor(input_details[idx]['index'], inputs[idx])
|
|
|
|
# Inference
|
|
interpreter.invoke()
|
|
|
|
def load_model(model, detect_tensor):
|
|
dump_tensor = False
|
|
if detect_tensor is not None:
|
|
dump_tensor = True
|
|
interpreter = tf.lite.Interpreter(model_path=model, experimental_preserve_all_tensors=dump_tensor)
|
|
interpreter.allocate_tensors()
|
|
return interpreter
|
|
|
|
if __name__ == '__main__':
|
|
args = arguments()
|
|
if args.detect_all:
|
|
detect_tensor = 'all'
|
|
elif args.detect_tensor is not None:
|
|
detect_tensor = list(args.detect_tensor.split(' '))
|
|
else:
|
|
detect_tensor = None
|
|
|
|
interpreter = load_model(args.model, detect_tensor)
|
|
inputs = pre_process_input(interpreter, args)
|
|
inference(interpreter, inputs)
|
|
|
|
if detect_tensor is not None:
|
|
dump_detect_tensors(interpreter, args, detect_tensor)
|
|
else:
|
|
post_process_output(interpreter, args)
|