Flag to disable Lightning Logs (#367)

* added flag to disable lightning logs

* changed default value to suppress logs as false

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: manujosephv

src/pytorch_tabular/tabular_model.py

@@ -51,11 +51,12 @@
 from pytorch_tabular.utils import (
     OOMException,
     OutOfMemoryHandler,
+    count_parameters,
     get_logger,
     getattr_nested,
     pl_load,
+    suppress_lightning_logs,
 )
-from pytorch_tabular.utils.nn_utils import count_parameters
 
 try:
     import captum.attr
@@ -79,6 +80,7 @@ def __init__(
         model_callable: Optional[Callable] = None,
         model_state_dict_path: Optional[Union[str, Path]] = None,
         verbose: bool = True,
+        suppress_lightning_logger: bool = False,
     ) -> None:
         """The core model which orchestrates everything from initializing the datamodule, the model, trainer, etc.
 
@@ -111,8 +113,13 @@ def __init__(
                 If provided, will load the state dict after initializing the model from config.
 
             verbose (bool): turns off and on the logging. Defaults to True.
+
+            suppress_lightning_logger (bool): If True, will suppress the default logging from PyTorch Lightning.
+                Defaults to False.
         """
         super().__init__()
+        if suppress_lightning_logger:
+            suppress_lightning_logs()
         self.verbose = verbose
         self.exp_manager = ExperimentRunManager()
         if config is None:

src/pytorch_tabular/tabular_model_tuner.py

@@ -22,7 +22,7 @@
     TrainerConfig,
 )
 from pytorch_tabular.tabular_model import TabularModel
-from pytorch_tabular.utils import OOMException, OutOfMemoryHandler, get_logger
+from pytorch_tabular.utils import OOMException, OutOfMemoryHandler, get_logger, suppress_lightning_logs
 
 logger = get_logger(__name__)
 
@@ -45,6 +45,7 @@ def __init__(
         trainer_config: Optional[Union[TrainerConfig, str]] = None,
         model_callable: Optional[Callable] = None,
         model_state_dict_path: Optional[Union[str, Path]] = None,
+        suppress_lightning_logger: bool = True,
         **kwargs,
     ):
         """Tabular Model Tuner helps you tune the hyperparameters of a TabularModel.
@@ -53,21 +54,30 @@ def __init__(
             data_config (Optional[Union[DataConfig, str]], optional): The DataConfig for the TabularModel.
                 If str is passed, will initialize the DataConfig using the yaml file in that path.
                 Defaults to None.
+
             model_config (Optional[Union[ModelConfig, str]], optional): The ModelConfig for the TabularModel.
                 If str is passed, will initialize the ModelConfig using the yaml file in that path.
                 Defaults to None.
+
             optimizer_config (Optional[Union[OptimizerConfig, str]], optional): The OptimizerConfig for the
                 TabularModel. If str is passed, will initialize the OptimizerConfig using the yaml file in
                 that path. Defaults to None.
+
             trainer_config (Optional[Union[TrainerConfig, str]], optional): The TrainerConfig for the TabularModel.
                 If str is passed, will initialize the TrainerConfig using the yaml file in that path.
                 Defaults to None.
+
             model_callable (Optional[Callable], optional): A callable that returns a PyTorch Tabular Model.
                 If provided, will ignore the model_config and use this callable to initialize the model.
                 Defaults to None.
+
             model_state_dict_path (Optional[Union[str, Path]], optional): Path to the state dict of the model.
+
                 If provided, will ignore the model_config and use this state dict to initialize the model.
                 Defaults to None.
+
+            suppress_lightning_logger (bool, optional): Whether to suppress the lightning logger. Defaults to True.
+
             **kwargs: Additional keyword arguments to be passed to the TabularModel init.
         """
         if trainer_config.profiler is not None:
@@ -84,6 +94,7 @@ def __init__(
         self.model_config = model_config
         self.optimizer_config = optimizer_config
         self.trainer_config = trainer_config
+        self.suppress_lightning_logger = suppress_lightning_logger
         self.tabular_model_init_kwargs = {
             "model_callable": model_callable,
             "model_state_dict_path": model_state_dict_path,
@@ -208,6 +219,8 @@ def tune(
         assert mode in ["max", "min"], "mode must be one of ['max', 'min']"
         assert metric is not None, "metric must be specified"
         assert isinstance(search_space, dict) and len(search_space) > 0, "search_space must be a non-empty dict"
+        if self.suppress_lightning_logger:
+            suppress_lightning_logs()
         if cv is not None and validation is not None:
             warnings.warn(
                 "Both validation and cv are provided. Ignoring validation and using cv. Use "