Added TensorFlow Profiler (isarandi#44)

Young-Yoon · Aug 28, 2023 · 78006a9 · 78006a9
1 parent 1a16a34
commit 78006a9
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 0 deletions.
diff --git a/src/init.py b/src/init.py
@@ -231,4 +231,8 @@ def get_parser():
     parser.add_argument('--model-class', type=str, default='Metrabs')
     parser.add_argument('--input-float32', action=options.BoolAction, default=False)
     parser.add_argument('--final-transposed-conv', type=int, default=1)
+    parser.add_argument('--tfprofiler', action=options.BoolAction, default=False,
+                       help='tensorflow profiler')
+    parser.add_argument('--tfprofiler-dir', type=str, default='tfprofiler',
+                        help='Directory path of tensorflow profiler.')
     return parser
diff --git a/src/main.py b/src/main.py
@@ -141,6 +141,12 @@ def train():
         callbacks.TensorBoardCallback(global_step)
     ]
 
+    # `pip install -U tensorboard_plugin_profiler` or ensure the docker version metrabs:0.1.21 or higher
+    if FLAGS.tfprofiler:
+        profiler_dir = util.ensure_absolute_path(os.path.join('experiments', FLAGS.tfprofiler_dir))
+        c = keras.callbacks.TensorBoard(log_dir=profiler_dir, profile_batch=(30, 35))
+        cbacks.append(c)
+
     if FLAGS.finetune_in_inference_mode:
         switch_step = FLAGS.training_steps - FLAGS.finetune_in_inference_mode
         c = callbacks.SwitchToInferenceModeCallback(global_step, switch_step)