From 0beeeb008d56844a434d23eb91ef7612317f0d76 Mon Sep 17 00:00:00 2001 From: Jeff Rasley Date: Wed, 17 Nov 2021 15:06:15 -0800 Subject: [PATCH 1/3] remove hard tbx req --- deepspeed/runtime/engine.py | 2 +- requirements/requirements.txt | 1 - 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py index 173d34896f69..a2bfa89eb758 100755 --- a/deepspeed/runtime/engine.py +++ b/deepspeed/runtime/engine.py @@ -18,7 +18,7 @@ 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 torch.utils.tensorboard import SummaryWriter from typing import Callable, Dict, Optional, Union, Iterable 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 From 1b975365bf83519813da01d903c6056690ec6535 Mon Sep 17 00:00:00 2001 From: Jeff Rasley Date: Wed, 17 Nov 2021 15:15:50 -0800 Subject: [PATCH 2/3] delay importing sw unless needed --- deepspeed/runtime/engine.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py index a2bfa89eb758..b65474e9c410 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 torch.utils.tensorboard import SummaryWriter from typing import Callable, Dict, Optional, Union, Iterable @@ -533,6 +532,14 @@ def get_summary_writer( log_dir = os.path.join(base, summary_writer_dir_name, name) os.makedirs(log_dir, exist_ok=True) + try: + import tensorboard + from torch.utils.tensorboard import SummaryWriter + except ImportError as err: + print( + 'If you want to use tensorboard logging please `pip install tensorboard`' + ) + raise err return SummaryWriter(log_dir=log_dir) From 642d722c0e472cefab227f9e00f2042a07462b00 Mon Sep 17 00:00:00 2001 From: Jeff Rasley Date: Wed, 17 Nov 2021 15:59:34 -0800 Subject: [PATCH 3/3] address offline comments --- deepspeed/runtime/engine.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py index b65474e9c410..96210d94d14a 100755 --- a/deepspeed/runtime/engine.py +++ b/deepspeed/runtime/engine.py @@ -533,13 +533,15 @@ def get_summary_writer( 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 - from torch.utils.tensorboard import SummaryWriter - except ImportError as err: + except ImportError: print( 'If you want to use tensorboard logging please `pip install tensorboard`' ) - raise err + raise + from torch.utils.tensorboard import SummaryWriter return SummaryWriter(log_dir=log_dir)