feat: relax tensorboard as a soft dependency (TorchEnsemble-Community#65

) * add code * Update CHANGELOG.rst * Update requirements.txt * Update logging.py
RyotaroOKabe · Apr 15, 2021 · f071ff9 · f071ff9
1 parent 4070b7e
commit f071ff9
Show file tree

Hide file tree

Showing 4 changed files with 9 additions and 7 deletions.
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
@@ -4,6 +4,7 @@ Changelog
 Ver 0.1.*
 ---------
 
+* |Enhancement| Relax :mod:`tensorboard` as a soft dependency | `@xuyxu <https://github.com/xuyxu>`__
 * |Enhancement| |API| Simplify the training workflow of :class:`FastGeometricClassifier` and :class:`FastGeometricRegressor` | `@xuyxu <https://github.com/xuyxu>`__
 * |Feature| |API| Support TensorBoard logging in :meth:`set_logger` | `@zzzzwj <https://github.com/zzzzwj>`__
 * |Enhancement| |API| Add ``use_reduction_sum`` parameter for :meth:`fit` of Gradient Boosting | `@xuyxu <https://github.com/xuyxu>`__

diff --git a/build_tools/requirements.txt b/build_tools/requirements.txt
@@ -1,3 +1,4 @@
 flake8
 pytest-cov
-black==20.8b1
+black==20.8b1
+tensorboard==2.*
diff --git a/requirements.txt b/requirements.txt
@@ -1,4 +1,3 @@
 torch>=1.4.0
 torchvision>=0.2.2
-scikit-learn>=0.23.0
-tensorboard==2.*
+scikit-learn>=0.23.0
diff --git a/torchensemble/utils/logging.py b/torchensemble/utils/logging.py
@@ -83,15 +83,16 @@ def _get_level(level):
 
 def init_tb_logger(log_dir):
     try:
-        from torch.utils.tensorboard import SummaryWriter
+        import tensorboard  # noqa: F401
     except ModuleNotFoundError:
         msg = (
-            "Cannot load the module torch when building the "
-            "ImageScanner. Please make sure that tensorboard is"
-            " installed."
+            "Cannot load the module tensorboard. Please make sure that"
+            " tensorboard is installed."
         )
         raise ModuleNotFoundError(msg)
 
+    from torch.utils.tensorboard import SummaryWriter
+
     global _tb_logger
 
     if not _tb_logger: