def export(self,
output_dir: Text,
tflite_path: Text = None,
tensorrt: Text = None):
"""Export a saved model, frozen graph, and potential tflite/tensorrt model.
Args:
output_dir: the output folder for saved model.
tflite_path: the path for saved tflite file.
tensorrt: If not None, must be {'FP32', 'FP16', 'INT8'}.
"""
signitures = self.signitures
signature_def_map = {
'serving_default':
tf.saved_model.predict_signature_def(
{signitures['image_arrays'].name: signitures['image_arrays']},
{signitures['prediction'].name: signitures['prediction']}),
}
b = tf.saved_model.Builder(output_dir)
b.add_meta_graph_and_variables(
self.sess,
tags=['serve'],
signature_def_map=signature_def_map,
assets_collection=tf.get_collection(tf.GraphKeys.ASSET_FILEPATHS),
clear_devices=True)
b.save()
logging.info('Model saved at %s', output_dir)
# also save freeze pb file.
graphdef = self.freeze()
pb_path = os.path.join(output_dir, self.model_name + '_frozen.pb')
tf.io.gfile.GFile(pb_path, 'wb').write(graphdef.SerializeToString())
logging.info('Frozen graph saved at %s', pb_path)
if tflite_path:
height, width = utils.parse_image_size(self.params['image_size'])
input_name = signitures['image_arrays'].op.name
input_shapes = {input_name: [None, height, width, 3]}
converter = tf.lite.TFLiteConverter.from_saved_model(
output_dir,
input_arrays=[input_name],
input_shapes=input_shapes,
output_arrays=[signitures['prediction'].op.name])
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS]
tflite_model = converter.convert()
tf.io.gfile.GFile(tflite_path, 'wb').write(tflite_model)
logging.info('TFLite is saved at %s', tflite_path)
if tensorrt:
from tensorflow.python.compiler.tensorrt import trt # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
sess_config = tf.ConfigProto(gpu_options=tf.GPUOptions(allow_growth=True))
trt_path = os.path.join(output_dir, 'tensorrt_' + tensorrt.lower())
trt.create_inference_graph(
None,
None,
precision_mode=tensorrt,
input_saved_model_dir=output_dir,
output_saved_model_dir=trt_path,
session_config=sess_config)
logging.info('TensorRT model is saved at %s', trt_path)
def __init__(self,
graph_path,
target_size=(320, 240),
tf_config=None,
trt_bool=False,
verbose=False):
self.target_size = target_size
# load graph
logger.info('loading graph from %s(default size=%dx%d)' %
(graph_path, target_size[0], target_size[1]))
with tf.io.gfile.GFile(graph_path, 'rb') as f:
graph_def = tf.compat.v1.GraphDef()
graph_def.ParseFromString(f.read())
if trt_bool is True:
output_nodes = ["Openpose/concat_stage7"]
graph_def = trt.create_inference_graph(
graph_def,
output_nodes,
max_batch_size=1,
max_workspace_size_bytes=1 << 20,
precision_mode="FP16",
# precision_mode="INT8",
minimum_segment_size=3,
is_dynamic_op=True,
maximum_cached_engines=int(1e3),
use_calibration=True,
)
self.graph = tf.compat.v1.get_default_graph()
tf.import_graph_def(graph_def, name='TfPoseEstimator')
self.persistent_sess = tf.compat.v1.Session(graph=self.graph,
config=tf_config)
if verbose:
for ts in [
n.name for n in
tf.compat.v1.get_default_graph().as_graph_def().node
]:
print(ts)
self.tensor_image = self.graph.get_tensor_by_name(
'TfPoseEstimator/image:0')
self.tensor_output = self.graph.get_tensor_by_name(
'TfPoseEstimator/Openpose/concat_stage7:0')
self.tensor_heatMat = self.tensor_output[:, :, :, :19]
self.tensor_pafMat = self.tensor_output[:, :, :, 19:]
self.upsample_size = tf.placeholder(dtype=tf.int32,
shape=(2, ),
name='upsample_size')
self.tensor_heatMat_up = tf.image.resize_area(
self.tensor_output[:, :, :, :19],
self.upsample_size,
align_corners=False,
name='upsample_heatmat')
self.tensor_pafMat_up = tf.image.resize_area(
self.tensor_output[:, :, :, 19:],
self.upsample_size,
align_corners=False,
name='upsample_pafmat')
if trt_bool is True:
smoother = Smoother({'data': self.tensor_heatMat_up}, 25, 3.0, 19)
else:
smoother = Smoother({'data': self.tensor_heatMat_up}, 25, 3.0)
gaussian_heatMat = smoother.get_output()
max_pooled_in_tensor = tf.nn.pool(gaussian_heatMat,
window_shape=(3, 3),
pooling_type='MAX',
padding='SAME')
self.tensor_peaks = tf.where(
tf.equal(gaussian_heatMat, max_pooled_in_tensor), gaussian_heatMat,
tf.zeros_like(gaussian_heatMat))
self.heatMat = self.pafMat = None
# warm-up
self.persistent_sess.run(
tf.compat.v1.variables_initializer([
v for v in tf.global_variables() if v.name.split(':')[0] in [
x.decode('utf-8') for x in self.persistent_sess.run(
tf.report_uninitialized_variables())
]
]))
self.persistent_sess.run(
[self.tensor_peaks, self.tensor_heatMat_up, self.tensor_pafMat_up],
feed_dict={
self.tensor_image: [
np.ndarray(shape=(target_size[1], target_size[0], 3),
dtype=np.float32)
],
self.upsample_size: [target_size[1], target_size[0]]
})
self.persistent_sess.run(
[self.tensor_peaks, self.tensor_heatMat_up, self.tensor_pafMat_up],
feed_dict={
self.tensor_image: [
np.ndarray(shape=(target_size[1], target_size[0], 3),
dtype=np.float32)
],
self.upsample_size: [target_size[1] // 2, target_size[0] // 2]
})
self.persistent_sess.run(
[self.tensor_peaks, self.tensor_heatMat_up, self.tensor_pafMat_up],
feed_dict={
self.tensor_image: [
np.ndarray(shape=(target_size[1], target_size[0], 3),
dtype=np.float32)
],
self.upsample_size: [target_size[1] // 4, target_size[0] // 4]
})
# logs
if self.tensor_image.dtype == tf.quint8:
logger.info('quantization mode enabled.')
