Skip to content

Instantly share code, notes, and snippets.

@gaiar

gaiar/TF-TRT.py

Created January 29, 2020 11:08
Show Gist options
  • Save gaiar/1fc48be388e6350890f9da600fa6f533 to your computer and use it in GitHub Desktop.
Save gaiar/1fc48be388e6350890f9da600fa6f533 to your computer and use it in GitHub Desktop.
Download ZIP
TF-TRT 1.x Workflow With A SavedModel
Raw
TF-TRT.py
import tensorflow as tf
from tensorflow.python.compiler.tensorrt import trt_convert as trt
BATCH_SIZE = 2
#GRAPH_PATH = '/home/gaiar/developer/megadetector/megadetector_v3_optimized.pb'
input_saved_model_dir = "/home/gaiar/developer/megadetector/megadetector/exported_model/saved_model"
output_saved_model_dir = "/home/gaiar/developer/megadetector/model_optimization/trt"
converter = trt.TrtGraphConverter(
input_saved_model_dir=input_saved_model_dir,
max_batch_size=BATCH_SIZE,
max_workspace_size_bytes=2048<<20,
precision_mode="FP32",
maximum_cached_engines=16)
converter.convert()
converter.save(output_saved_model_dir)
@sheneman
Copy link

Thanks for writing this.

Using TF 1.15.5, In the call to converter.convert() I get this:

RuntimeError: MetaGraphDef associated with tags 'serve' could not be found in SavedModel. To inspect available tag-sets in the SavedModel, please use the SavedModel CLI: saved_model_cli
available_tags: [set()]

Have you seen this?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment