def main(argv): if len(argv) > 1: raise tf.app.UsageError('Too many command-line arguments.') export_saved_model_tpu_lib.export(FLAGS.pipeline_config_file, FLAGS.ckpt_path, FLAGS.export_dir, FLAGS.input_placeholder_name, FLAGS.input_type, FLAGS.use_bfloat16)
def testExportAndLoad(self, pipeline_config_file, input_type='image_tensor', use_bfloat16=False, repeat=1): input_placeholder_name = 'placeholder_tensor' export_dir = os.path.join(FLAGS.test_tmpdir, 'tpu_saved_model') if tf.io.gfile.exists(export_dir): tf.io.gfile.rmtree(export_dir) ckpt_path = None export_saved_model_tpu_lib.export(pipeline_config_file, ckpt_path, export_dir, input_placeholder_name, input_type, use_bfloat16) inputs = np.random.rand(256, 256, 3) tensor_dict_out = export_saved_model_tpu_lib.run_inference_from_saved_model( inputs, export_dir, input_placeholder_name, repeat) for k, v in tensor_dict_out.items(): tf.compat.v1.logging.info('{}: {}'.format(k, v))