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