diff options
| author |  TensorFlower Gardener <gardener@tensorflow.org> | 2018-08-06 13:54:06 -0700 |
|---|---|---|
| committer | TensorFlower Gardener <gardener@tensorflow.org> | 2018-08-06 13:54:06 -0700 |
| commit | 827b04649c5ec55997ff3917a4b377eb49c4cbbb (patch) | |
| tree | d25eee583cac73d2eb88e0032d1b2a0d52510fbb /tensorflow/python/tools | |
| parent | f886e267ff644896a74c986d2b3b9945040681b6 (diff) | |
| parent | 8636c6dbbe1446cc139e6c23a435ca7cd0844fca (diff) | |
Merge pull request #19920 from samikama:TRTOpsForTensorBoard
PiperOrigin-RevId: 207600851
Diffstat (limited to 'tensorflow/python/tools')
| -rw-r--r-- | tensorflow/python/tools/BUILD | 1 | ||||
| -rw-r--r-- | tensorflow/python/tools/import_pb_to_tensorboard.py | 10 |
2 files changed, 11 insertions, 0 deletions
diff --git a/tensorflow/python/tools/BUILD b/tensorflow/python/tools/BUILD index 6c34b6aaf3..222f856511 100644 --- a/tensorflow/python/tools/BUILD +++ b/tensorflow/python/tools/BUILD @@ -64,6 +64,7 @@ py_binary( srcs_version = "PY2AND3", deps = [ "//tensorflow/core:protos_all_py", + "//tensorflow/python", "//tensorflow/python:client", "//tensorflow/python:framework", "//tensorflow/python:framework_ops", diff --git a/tensorflow/python/tools/import_pb_to_tensorboard.py b/tensorflow/python/tools/import_pb_to_tensorboard.py index 00de044505..6d2fec3ad6 100644 --- a/tensorflow/python/tools/import_pb_to_tensorboard.py +++ b/tensorflow/python/tools/import_pb_to_tensorboard.py @@ -29,6 +29,16 @@ from tensorflow.python.platform import app from tensorflow.python.platform import gfile from tensorflow.python.summary import summary +# Try importing TensorRT ops if available +# TODO(aaroey): ideally we should import everything from contrib, but currently +# tensorrt module would cause build errors when being imported in +# tensorflow/contrib/__init__.py. Fix it. +# pylint: disable=unused-import,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,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. |