How to use nvtx-plugins with tf.session calls

[mankeyboy]

The code I'm working on is this.
The code uses tf.session call to take in a graph for object detection tasks. Link
My aim here is to profile this code for Nvidia GPUs using the nvtx-plugins-tf to analyze the time taken for different ops. Link to docs

The plugin library provides a function hook for a tf.train.MonitoredSession as given in their example code here.
The code linked above uses tf.session along with a tf.config and when I try to modify the tf.session call to a tf.train.MonitoredSession call, I can't get my code to work and it fails with an error that graph can't be modified. I went through the tensorflow APIs and it turns out that tf.session doesn't support hook callbacks and tf.train.MonitoredSession doesn't support tf_config as a function argument.

Traceback (most recent call last):
  File "/home/mayroy13/anaconda3/envs/trt-py36/lib/python3.6/runpy.py", line 193, in _run_module_as_main
    "__main__", mod_spec)
  File "/home/mayroy13/anaconda3/envs/trt-py36/lib/python3.6/runpy.py", line 85, in _run_code
    exec(code, run_globals)
  File "/home/mayroy13/Mayank/Mayank/test/tensorrt/tftrt/examples/object_detection/test.py", line 105, in <module>
    test(args.test_config_path)
  File "/home/mayroy13/Mayank/Mayank/test/tensorrt/tftrt/examples/object_detection/test.py", line 81, in test
    **test_config['benchmark_config'])
  File "/home/mayroy13/Mayank/Mayank/test/tensorrt/tftrt/examples/object_detection/object_detection.py", line 608, in benchmark_model
    tf.import_graph_def(frozen_graph, name='')
  File "/home/mayroy13/anaconda3/envs/trt-py36/lib/python3.6/site-packages/tensorflow/python/util/deprecation.py", line 507, in new_func
    return func(*args, **kwargs)
  File "/home/mayroy13/anaconda3/envs/trt-py36/lib/python3.6/site-packages/tensorflow/python/framework/importer.py", line 443, in import_graph_def
    _ProcessNewOps(graph)
  File "/home/mayroy13/anaconda3/envs/trt-py36/lib/python3.6/site-packages/tensorflow/python/framework/importer.py", line 236, in _ProcessNewOps
    for new_op in graph._add_new_tf_operations(compute_devices=False):  # pylint: disable=protected-access
  File "/home/mayroy13/anaconda3/envs/trt-py36/lib/python3.6/site-packages/tensorflow/python/framework/ops.py", line 3751, in _add_new_tf_operations
    for c_op in c_api_util.new_tf_operations(self)
  File "/home/mayroy13/anaconda3/envs/trt-py36/lib/python3.6/site-packages/tensorflow/python/framework/ops.py", line 3751, in <listcomp>
    for c_op in c_api_util.new_tf_operations(self)
  File "/home/mayroy13/anaconda3/envs/trt-py36/lib/python3.6/site-packages/tensorflow/python/framework/ops.py", line 3640, in _create_op_from_tf_operation
    self._check_not_finalized()
  File "/home/mayroy13/anaconda3/envs/trt-py36/lib/python3.6/site-packages/tensorflow/python/framework/ops.py", line 3225, in _check_not_finalized
    raise RuntimeError("Graph is finalized and cannot be modified.")
RuntimeError: Graph is finalized and cannot be modified.

Any directions to go in would be appreciated. If there are ways in tensorflow to use hooks in conjunction with tf.session, that will also work for me.

Stack Overflow question related to this: https://stackoverflow.com/questions/58162110/how-to-port-a-tf-session-to-a-tf-train-monitoredsession-call-while-allowing-grap

[ahmadki]

I think what you are looking for in your case are nvtx_tf.ops.start() and nvtx_tf.ops.end() which wrap individual or multiple layers and operations. nvtx.plugins.tf.estimator.NVTXHook() is intended to annotate an entire monitored_session.run() call.

tf.session is a high level API and doesn't support hooks so naturally nvtx.plugins.tf.estimator.NVTXHook() will not work. However, all that NVTXHook does is open/close nvtx ranges before/after each monitored_session.run() call. This is done using the nvtx C API.
To replicate this behaviour look at https://github.com/NVIDIA/nvtx-plugins/blob/master/nvtx_plugins/python/nvtx/plugins/tf/base_callbacks.py and wrap your session.run() with open_marker() and close_marker()

[mankeyboy]

Can you elaborate on how nvtx_tf.ops.start() and nvtx_tf.ops.end() would be used in this situation? I've tried putting it using tf_input but it throws an error. I also tried using the base open_marker and close_marker around my call but it didn't work.