diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py index 173d34896f69..96210d94d14a 100755 --- a/deepspeed/runtime/engine.py +++ b/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) diff --git a/requirements/requirements.txt b/requirements/requirements.txt index 2beb756048bb..80a4d727f7f5 100755 --- a/requirements/requirements.txt +++ b/requirements/requirements.txt @@ -1,6 +1,5 @@ torch tqdm -tensorboardX==1.8 ninja numpy psutil