tensorboard add_graph

Summary: Add the API to log graph in tensorboard

Reviewed By: wat3rBro

Differential Revision: D27855774

fbshipit-source-id: 415c469c5de0c56fc828d1b95f4be697e0acac84

d2go/runner/default_runner.py

@@ -276,6 +276,13 @@ class Detectron2GoRunner(BaseRunner):
             if cfg.MODEL_EMA.ENABLED and cfg.MODEL_EMA.USE_EMA_WEIGHTS_FOR_EVAL_ONLY:
                 model_ema.apply_model_ema(model)
 
+        # Note: the _visualize_model API is experimental
+        if comm.is_main_process():
+            if hasattr(model, "_visualize_model"):
+                logger.info("Adding model visualization ...")
+                tbx_writer = _get_tbx_writer(get_tensorboard_log_dir(cfg.OUTPUT_DIR))
+                model._visualize_model(tbx_writer)
+
         return model
 
     def build_checkpointer(self, cfg, model, save_dir, **kwargs):