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.')