add profiler wrapper

utils/benchmark.py

@@ -3,6 +3,7 @@
 import utils.misc as utils
 import utils.dataset as utils_ds
 from tqdm.auto import tqdm
+from utils.profiler_wrapper import TBTracer
 
 intra_op_parallelism_threads = None
 
@@ -111,7 +112,7 @@ def run_model(single_pass_func, runner, dataset, batch_size, num_of_runs, timeou
         if dataset.available_instances < num_of_runs * batch_size:
             utils.print_goodbye_message_and_die(
                 f"Number of runs requested exceeds number of instances available in dataset!")
-
+    tracer = TBTracer()
     try:
         if num_of_runs is None:
             single_pass_func(runner, dataset)
@@ -123,7 +124,7 @@ def run_model(single_pass_func, runner, dataset, batch_size, num_of_runs, timeou
                 single_pass_func(runner, dataset)
     except utils_ds.OutOfInstances:
         pass
-
+    tracer.write()
     return dataset.summarize_accuracy(), runner.print_performance_metrics(batch_size)
 
 

utils/profiler_wrapper.py

@@ -0,0 +1,25 @@
+import tensorflow as tf
+import os
+from datetime import datetime
+
+def print_prof():
+  if os.getenv("DLS_PROFILER", 0):
+    try:
+      tf.DLS.print_profile_data()  
+    except AttributeError:
+      print("Non dls tf")
+
+class TBTracer:
+  def __init__(self):
+    self.should_trace = os.getenv("TRACE", 0)
+    if self.should_trace: 
+      # Set up logging.
+      options = tf.profiler.experimental.ProfilerOptions()
+      stamp = datetime.now().strftime("%Y%m%d-%H%M%S")
+      self.logdir = 'logs/func/%s' % stamp
+      tf.profiler.experimental.start(self.logdir, options = options)
+
+  def write(self):
+    if self.should_trace:
+      tf.profiler.experimental.stop()
+    print_prof()