diff options
| author |  Sami Kama <skama@nvidia.com> | 2018-06-11 15:10:41 -0700 |
|---|---|---|
| committer | Sami Kama <skama@nvidia.com> | 2018-06-11 15:10:41 -0700 |
| commit | 5fc6ebb1db3b7714f5449737438ac5007abbb410 (patch) | |
| tree | fa6fca888abe414639c02aedc571b7f52145688a /tensorflow/python/tools | |
| parent | a4b390bffbcb01d8f57f25c007277d457f752a69 (diff) | |
Try importing TRTOps to import_pb_to_tensorboard script to convert TF-TRT optimized graphs to tensorboard graphs
Diffstat (limited to 'tensorflow/python/tools')
| -rw-r--r-- | tensorflow/python/tools/import_pb_to_tensorboard.py | 7 |
1 files changed, 7 insertions, 0 deletions
diff --git a/tensorflow/python/tools/import_pb_to_tensorboard.py b/tensorflow/python/tools/import_pb_to_tensorboard.py index 00de044505..96f47c85da 100644 --- a/tensorflow/python/tools/import_pb_to_tensorboard.py +++ b/tensorflow/python/tools/import_pb_to_tensorboard.py @@ -29,6 +29,13 @@ from tensorflow.python.platform import app from tensorflow.python.platform import gfile from tensorflow.python.summary import summary +# Try importing TensorRT ops if available +# pylint: disable=unused-import,trailing-whitespace,g-import-not-at-top,wildcard-import +try: + from tensorflow.contrib.tensorrt.ops.gen_trt_engine_op import * +except ImportError: + pass +# pylint: enable=unused-import,trailing-whitespace,g-import-not-at-top,wildcard-import def import_to_tensorboard(model_dir, log_dir): """View an imported protobuf model (`.pb` file) as a graph in Tensorboard. |