Skip to content

Instantly share code, notes, and snippets.

@gaiar

gaiar/TF-TRT.py

Created January 29, 2020 11:08
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save gaiar/1fc48be388e6350890f9da600fa6f533 to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save gaiar/1fc48be388e6350890f9da600fa6f533 to your computer and use it in GitHub Desktop.
Download ZIP

Revisions

  1. gaiar created this gist Jan 29, 2020.
    16 changes: 16 additions & 0 deletions TF-TRT.py
    View file
    Original file line number Diff line number Diff line change
    @@ -0,0 +1,16 @@
    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)