Import tensorrt if available to import_pb_to_tensorboard.py for displaying TensorRT ops
authorSami Kama <skama@nvidia.com>
Wed, 30 May 2018 04:00:22 +0000 (21:00 -0700)
committerSami Kama <skama@nvidia.com>
Wed, 30 May 2018 04:00:22 +0000 (21:00 -0700)
tensorflow/python/tools/import_pb_to_tensorboard.py

index 00de044..d1f9cd8 100755 (executable)
@@ -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
+try:
+  import tensorflow.contrib.tensorrt as trt 
+except ImportError:
+  pass
+# pylint: enable=unused-import,trailing-whitespace
 
 def import_to_tensorboard(model_dir, log_dir):
   """View an imported protobuf model (`.pb` file) as a graph in Tensorboard.