Migrate off of tensorboardX

classy_train.py

@@ -128,12 +128,12 @@ def configure_hooks(args, config):
 
     if not args.skip_tensorboard:
         try:
-            from tensorboardX import SummaryWriter
+            from torch.utils.tensorboard import SummaryWriter
 
             tb_writer = SummaryWriter(log_dir=Path(base_folder) / "tensorboard")
             hooks.append(TensorboardPlotHook(tb_writer))
         except ImportError:
-            logging.warning("tensorboardX not installed, skipping tensorboard hooks")
+            logging.warning("tensorboard not installed, skipping tensorboard hooks")
 
     args_dict = vars(args)
     args_dict["config"] = config

classy_vision/generic/visualize.py

@@ -15,7 +15,7 @@
 
 try:
     import visdom
-    from tensorboardX import SummaryWriter
+    from torch.utils.tensorboard import SummaryWriter
 except ImportError:
     pass
 

classy_vision/hooks/model_tensorboard_hook.py

@@ -14,7 +14,7 @@
 
 
 try:
-    from tensorboardX import SummaryWriter  # noqa F401
+    from torch.utils.tensorboard import SummaryWriter  # noqa F401
 
     tbx_available = True
 except ImportError:
@@ -45,7 +45,7 @@ def __init__(self, tb_writer) -> None:
         super().__init__()
         if not tbx_available:
             raise RuntimeError(
-                "tensorboardX not installed, cannot use ModelTensorboardHook"
+                "tensorboard not installed, cannot use ModelTensorboardHook"
             )
 
         self.tb_writer = tb_writer

classy_vision/hooks/tensorboard_plot_hook.py

@@ -14,11 +14,11 @@
 
 
 try:
-    from tensorboardX import SummaryWriter  # noqa F401
+    from torch.utils.tensorboard import SummaryWriter  # noqa F401
 
-    tbx_available = True
+    tb_available = True
 except ImportError:
-    tbx_available = False
+    tb_available = False
 
 
 log = logging.getLogger()
@@ -45,9 +45,9 @@ def __init__(self, tb_writer, log_period: int = 10) -> None:
             SummaryWriter>`_ instance
         """
         super().__init__()
-        if not tbx_available:
+        if not tb_available:
             raise RuntimeError(
-                "tensorboardX not installed, cannot use TensorboardPlotHook"
+                "tensorboard not installed, cannot use TensorboardPlotHook"
             )
         assert isinstance(log_period, int), "log_period must be an int"
 
@@ -121,7 +121,7 @@ def on_phase_end(self, task) -> None:
         logging.info(f"Plotting to Tensorboard for {phase_type} phase {phase_type_idx}")
 
         phase_type = task.phase_type
-        learning_rate_key = f"Learning Rate/{phase_type}"
+        learning_rate_key = f"Learning_Rate/{phase_type}"
 
         if task.train:
             for learning_rate, global_step, wall_time in zip(

scripts/formatter.sh

@@ -43,7 +43,7 @@ then
     fi
 
     # run isort
-    cmd="isort -o classy_vision -o torchelastic -o visdom -o tensorboardX \
+    cmd="isort -o classy_vision -o torchelastic -o visdom \
     -o progressbar $CHANGED_FILES"
     echo "Running command \"$cmd\""
     ($cmd)

test/manual/hooks_model_tensorboard_hook_test.py

@@ -11,7 +11,7 @@
 
 from classy_vision.hooks import ModelTensorboardHook
 from classy_vision.models import build_model
-from tensorboardX import SummaryWriter
+from torch.utils.tensorboard import SummaryWriter
 
 
 class TestModelTensorboardHook(HookTestBase):

test/manual/hooks_tensorboard_plot_hook_test.py

@@ -18,7 +18,7 @@
 from classy_vision.optim.param_scheduler import UpdateInterval
 from classy_vision.tasks import build_task
 from classy_vision.trainer import LocalTrainer
-from tensorboardX import SummaryWriter
+from torch.utils.tensorboard import SummaryWriter
 
 
 class TestTensorboardPlotHook(HookTestBase):

tutorials/getting_started.ipynb

@@ -72,7 +72,7 @@
    },
    "outputs": [],
    "source": [
-    "! pip install tensorboard tensorboardX"
+    "! pip install tensorboard"
    ]
   },
   {