Remove hard tensorboardX requirement

deepspeed/runtime/engine.py

@@ -18,7 +18,6 @@
 from torch.optim import Optimizer
 from torch.optim.lr_scheduler import _LRScheduler
 from torch.distributed.distributed_c10d import _get_global_rank
-from tensorboardX import SummaryWriter
 
 from typing import Callable, Dict, Optional, Union, Iterable
 
@@ -533,6 +532,16 @@ def get_summary_writer(
             log_dir = os.path.join(base, summary_writer_dir_name, name)
 
         os.makedirs(log_dir, exist_ok=True)
+        try:
+            # torch.utils.tensorboard will fail if `tensorboard` is not available,
+            # see their docs for more details: https://pytorch.org/docs/1.8.0/tensorboard.html
+            import tensorboard
+        except ImportError:
+            print(
+                'If you want to use tensorboard logging please `pip install tensorboard`'
+            )
+            raise
+        from torch.utils.tensorboard import SummaryWriter
 
         return SummaryWriter(log_dir=log_dir)
 

requirements/requirements.txt

@@ -1,6 +1,5 @@
 torch
 tqdm
-tensorboardX==1.8
 ninja
 numpy
 psutil