feat(logging): add tensorboard logging (#61)

* add logger by tensorboard

* Fix small issue of Tensorboard Logger

* reformat the code with flake8 & black

* Add tensorboard to requirements.txt

* Reformat code by black

* Refactor the tensorboard logging module

* fix conflicts

* Add unit test script of tensorboard-logging

* Reformat the code

* Update CHANGELOG.rst

* update doc

* pin tensorboard version

* revert default value

Co-authored-by: zzzzwj <zwj@smail.nju.edu.cn>
Co-authored-by: Yi-Xuan Xu <xuyx@lamda.nju.edu.cn>

Author: 3 people

CHANGELOG.rst

@@ -4,6 +4,7 @@ Changelog
 Ver 0.1.*
 ---------
 
+* |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>`__
 * |Feature| |API| Improve the functionality of :meth:`evaluate` and :meth:`predict` | `@xuyxu <https://github.com/xuyxu>`__
 * |Feature| |API| Add :class:`FastGeometricClassifier` and :class:`FastGeometricRegressor` | `@xuyxu <https://github.com/xuyxu>`__

docs/quick_start.rst

@@ -57,7 +57,13 @@ Ensemble-PyTorch uses a global logger to track and print the intermediate inform
 
     logger = set_logger("classification_mnist_mlp")
 
-Using the logger, all intermediate information will be printed on the command line and saved to the specified text file: classification_mnist_mlp.
+With this logger, all intermediate information will be printed on the command line and saved to the specified text file: ``classification_mnist_mlp``.
+
+In addition, when passing ``use_tb_logger=True`` into the method :meth:`set_logger`, you can use tensorboard to have a better visualization result on training and evaluating the ensemble.
+
+.. code-block:: bash
+
+    tensorboard --logdir=logs/
 
 Choose the Ensemble
 -------------------

examples/classification_cifar10_cnn.py

@@ -96,7 +96,7 @@ def forward(self, x):
         shuffle=True,
     )
 
-    logger = set_logger("classification_cifar10_cnn")
+    logger = set_logger("classification_cifar10_cnn", use_tb_logger=True)
 
     # FusionClassifier
     model = FusionClassifier(

examples/fast_geometric_ensemble_cifar10_resnet18.py

@@ -136,7 +136,9 @@ def forward(self, x):
     )
 
     # Set the Logger
-    logger = set_logger("FastGeometricClassifier_cifar10_resnet")
+    logger = set_logger(
+        "FastGeometricClassifier_cifar10_resnet", use_tb_logger=True
+    )
 
     # Choose the Ensemble Method
     model = FastGeometricClassifier(
@@ -155,7 +157,11 @@ def forward(self, x):
     model.set_scheduler("CosineAnnealingLR", T_max=epochs)
 
     # Train
-    estimator = model.fit(train_loader, epochs=epochs, test_loader=test_loader)
+    estimator = model.fit(
+        train_loader,
+        epochs=epochs,
+        test_loader=test_loader,
+    )
 
     # Ensemble
     model.ensemble(

examples/regression_YearPredictionMSD_mlp.py

@@ -98,7 +98,7 @@ def forward(self, x):
     train_loader, test_loader = load_data(batch_size)
     print("Finish loading data...\n")
 
-    logger = set_logger("regression_YearPredictionMSD_mlp")
+    logger = set_logger("regression_YearPredictionMSD_mlp", use_tb_logger=True)
 
     # FusionRegressor
     model = FusionRegressor(

examples/snapshot_ensemble_cifar10_resnet18.py

@@ -136,7 +136,9 @@ def forward(self, x):
     )
 
     # Set the Logger
-    logger = set_logger("snapshot_ensemble_cifar10_resnet18")
+    logger = set_logger(
+        "snapshot_ensemble_cifar10_resnet18", use_tb_logger=True
+    )
 
     # Choose the Ensemble Method
     model = SnapshotEnsembleClassifier(
@@ -152,4 +154,8 @@ def forward(self, x):
     )
 
     # Train and Evaluate
-    model.fit(train_loader, epochs=epochs, test_loader=test_loader)
+    model.fit(
+        train_loader,
+        epochs=epochs,
+        test_loader=test_loader,
+    )

requirements.txt

@@ -1,3 +1,4 @@
 torch>=1.4.0
 torchvision>=0.2.2
-scikit-learn>=0.23.0
+scikit-learn>=0.23.0
+tensorboard==2.*

setup.py

@@ -8,14 +8,14 @@
 here = path.abspath(path.dirname(__file__))
 
 # Get the long description from README.rst
-with open(path.join(here, 'README.rst'), encoding='utf-8') as f:
+with open(path.join(here, "README.rst"), encoding="utf-8") as f:
     long_description = f.read()
 
 # get the dependencies and installs
-with open(path.join(here, 'requirements.txt'), encoding='utf-8') as f:
-    all_reqs = f.read().split('\n')
+with open(path.join(here, "requirements.txt"), encoding="utf-8") as f:
+    all_reqs = f.read().split("\n")
 
-install_requires = [x.strip() for x in all_reqs if 'git+' not in x]
+install_requires = [x.strip() for x in all_reqs if "git+" not in x]
 
 cmdclass = {}
 
@@ -28,12 +28,12 @@ def run(self):
         Clean.run(self)
         # Remove c files if we are not within a sdist package
         cwd = os.path.abspath(os.path.dirname(__file__))
-        remove_c_files = not os.path.exists(os.path.join(cwd, 'PKG-INFO'))
+        remove_c_files = not os.path.exists(os.path.join(cwd, "PKG-INFO"))
         if remove_c_files:
-            print('Will remove generated .c files')
-        if os.path.exists('build'):
-            shutil.rmtree('build')
-        for dirpath, dirnames, filenames in os.walk('sklearn'):
+            print("Will remove generated .c files")
+        if os.path.exists("build"):
+            shutil.rmtree("build")
+        for dirpath, dirnames, filenames in os.walk("sklearn"):
             for filename in filenames:
                 if any(
                     filename.endswith(suffix)
@@ -42,50 +42,50 @@ def run(self):
                     os.unlink(os.path.join(dirpath, filename))
                     continue
                 extension = os.path.splitext(filename)[1]
-                if remove_c_files and extension in ['.c', '.cpp']:
-                    pyx_file = str.replace(filename, extension, '.pyx')
+                if remove_c_files and extension in [".c", ".cpp"]:
+                    pyx_file = str.replace(filename, extension, ".pyx")
                     if os.path.exists(os.path.join(dirpath, pyx_file)):
                         os.unlink(os.path.join(dirpath, filename))
             for dirname in dirnames:
-                if dirname == '__pycache__':
+                if dirname == "__pycache__":
                     shutil.rmtree(os.path.join(dirpath, dirname))
 
 
-cmdclass.update({'clean': CleanCommand})
+cmdclass.update({"clean": CleanCommand})
 
 
 setup(
-    name='torchensemble',
-    maintainer='Yi-Xuan Xu',
-    maintainer_email='xuyx@lamda.nju.edu.cn',
+    name="torchensemble",
+    maintainer="Yi-Xuan Xu",
+    maintainer_email="xuyx@lamda.nju.edu.cn",
     description=(
-        'Implementations of scikit-learn like ensemble methods in Pytorch'
+        "Implementations of scikit-learn like ensemble methods in Pytorch"
     ),
-    license='BSD 3-Clause',
-    url='https://github.com/xuyxu/Ensemble-Pytorch',
+    license="BSD 3-Clause",
+    url="https://github.com/xuyxu/Ensemble-Pytorch",
     project_urls={
-        'Bug Tracker': 'https://github.com/xuyxu/Ensemble-Pytorch/issues',
-        'Documentation': 'https://ensemble-pytorch.readthedocs.io',
-        'Source Code': 'https://github.com/xuyxu/Ensemble-Pytorch',
+        "Bug Tracker": "https://github.com/xuyxu/Ensemble-Pytorch/issues",
+        "Documentation": "https://ensemble-pytorch.readthedocs.io",
+        "Source Code": "https://github.com/xuyxu/Ensemble-Pytorch",
     },
-    version='0.1.2',
+    version="0.1.2",
     long_description=long_description,
     classifiers=[
-        'Intended Audience :: Science/Research',
-        'Intended Audience :: Developers',
-        'Topic :: Software Development',
-        'Topic :: Scientific/Engineering',
-        'Development Status :: 4 - Beta',
-        'Operating System :: Microsoft :: Windows',
-        'Operating System :: POSIX',
-        'Operating System :: Unix',
-        'Operating System :: MacOS',
-        'Programming Language :: Python :: 3',
-        'Programming Language :: Python :: 3.6',
-        'Programming Language :: Python :: 3.7',
-        'Programming Language :: Python :: 3.8',
+        "Intended Audience :: Science/Research",
+        "Intended Audience :: Developers",
+        "Topic :: Software Development",
+        "Topic :: Scientific/Engineering",
+        "Development Status :: 4 - Beta",
+        "Operating System :: Microsoft :: Windows",
+        "Operating System :: POSIX",
+        "Operating System :: Unix",
+        "Operating System :: MacOS",
+        "Programming Language :: Python :: 3",
+        "Programming Language :: Python :: 3.6",
+        "Programming Language :: Python :: 3.7",
+        "Programming Language :: Python :: 3.8",
     ],
-    keywords=['PyTorch', 'Ensemble Learning'],
+    keywords=["PyTorch", "Ensemble Learning"],
     packages=find_packages(),
     cmdclass=cmdclass,
     python_requires=">=3.6",

torchensemble/_base.py

@@ -7,6 +7,7 @@
 import torch.nn as nn
 
 from . import _constants as const
+from .utils.logging import get_tb_logger
 
 
 def torchensemble_model_doc(header="", item="model"):
@@ -77,6 +78,7 @@ def __init__(
         self.device = torch.device("cuda" if cuda else "cpu")
         self.n_jobs = n_jobs
         self.logger = logging.getLogger()
+        self.tb_logger = get_tb_logger()
 
         self.estimators_ = nn.ModuleList()
         self.use_scheduler_ = False

torchensemble/adversarial_training.py

@@ -368,6 +368,12 @@ def _forward(estimators, data):
                             " % | Historical Best: {:.3f} %"
                         )
                         self.logger.info(msg.format(epoch, acc, best_acc))
+                        if self.tb_logger:
+                            self.tb_logger.add_scalar(
+                                "adversarial_training/Validation_Acc",
+                                acc,
+                                epoch,
+                            )
 
                 # Update the scheduler
                 with warnings.catch_warnings():
@@ -541,6 +547,12 @@ def _forward(estimators, data):
                             " {:.5f} | Historical Best: {:.5f}"
                         )
                         self.logger.info(msg.format(epoch, mse, best_mse))
+                        if self.tb_logger:
+                            self.tb_logger.add_scalar(
+                                "adversirial_training/Validation_MSE",
+                                mse,
+                                epoch,
+                            )
 
                 # Update the scheduler
                 with warnings.catch_warnings():