add specific param logging to wandb; update python logging

Author: slegroux

nbs/image.datasets.ipynb

@@ -66,7 +66,10 @@
     "\n",
     "from typing import Any, Dict, Optional, Tuple, List\n",
     "from nimrod.data.core import DataModule\n",
-    "from nimrod.utils import logger, set_seed"
+    "from nimrod.utils import set_seed\n",
+    "\n",
+    "import logging\n",
+    "logger = logging.getLogger(__name__)"
    ]
   },
   {

nbs/models.conv.ipynb

@@ -56,8 +56,11 @@
     "from omegaconf import OmegaConf\n",
     "\n",
     "from nimrod.image.datasets import MNISTDataModule\n",
-    "from nimrod.utils import get_device, logger\n",
-    "from nimrod.models.core import Classifier"
+    "from nimrod.utils import get_device\n",
+    "from nimrod.models.core import Classifier\n",
+    "\n",
+    "import logging\n",
+    "logger = logging.getLogger(__name__)"
    ]
   },
   {

nbs/models.core.ipynb

@@ -42,9 +42,11 @@
     "import torch\n",
     "\n",
     "from abc import ABC, abstractmethod\n",
-    "from nimrod.utils import logger\n",
+    "# from nimrod.utils import logger\n",
     "\n",
-    "from torchmetrics import Accuracy\n"
+    "from torchmetrics import Accuracy\n",
+    "import logging\n",
+    "logger = logging.getLogger(__name__)\n"
    ]
   },
   {

nbs/models.lm.ipynb

@@ -52,10 +52,12 @@
     "\n",
     "from nimrod.utils import get_device\n",
     "from nimrod.image.datasets import MNISTDataModule\n",
-    "from nimrod.utils import logger\n",
     "from nimrod.models.core import Classifier\n",
     "# torch.set_num_interop_threads(1)\n",
-    "# from IPython.core.debugger import set_trace"
+    "# from IPython.core.debugger import set_trace\n",
+    "\n",
+    "import logging\n",
+    "logger = logging.getLogger(__name__)"
    ]
   },
   {

nbs/models.mlp.ipynb

@@ -76,8 +76,11 @@
     "\n",
     "# nimrod\n",
     "# from nimrod.models.lm import Vocab\n",
-    "from nimrod.utils import set_seed, logger\n",
-    "from nimrod.data.core import DataModule\n"
+    "from nimrod.utils import set_seed\n",
+    "from nimrod.data.core import DataModule\n",
+    "\n",
+    "import logging\n",
+    "logger = logging.getLogger(__name__)\n"
    ]
   },
   {

nbs/text.datasets.ipynb

@@ -51,6 +51,7 @@
     "import random\n",
     "import os\n",
     "import logging\n",
+    "from rich.logging import RichHandler\n",
     "import lightning as L"
    ]
   },
@@ -87,8 +88,8 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "cpu\n",
-      "Is MPS (Metal Performance Shader) built? False\n"
+      "mps\n",
+      "Is MPS (Metal Performance Shader) built? True\n"
      ]
     }
    ],
@@ -123,7 +124,7 @@
     "    # # Set a fixed value for the hash seed\n",
     "    # os.environ[\"PYTHONHASHSEED\"] = str(seed)\n",
     "    # print(f\"Random seed set as {seed}\")\n",
-    "    L.seed_everything(seed)"
+    "    L.seed_everything(seed, workers=True)"
    ]
   },
   {
@@ -166,10 +167,15 @@
     "#| export\n",
     "\n",
     "# Configure the logger\n",
-    "logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')\n",
+    "logging.basicConfig(\n",
+    "    level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s',\n",
+    "    handlers=[RichHandler(rich_tracebacks=True)]\n",
+    "    )\n",
     "\n",
     "# Create a logger\n",
-    "logger = logging.getLogger(__name__)"
+    "# logger = logging.getLogger(__name__)\n",
+    "def get_logger(name=__name__):\n",
+    "    return logging.getLogger(name)"
    ]
   },
   {

nbs/utils.ipynb

@@ -4,7 +4,10 @@
 from omegaconf import DictConfig, OmegaConf
 import hydra
 from hydra.utils import instantiate
-import wandb
+from nimrod.utils import set_seed
+import json
+import logging
+log = logging.getLogger(__name__)
 
 @hydra.main(version_base="1.3",config_path="conf", config_name="train_mlp.yaml")
 def main(cfg: DictConfig) -> dict:
@@ -15,12 +18,10 @@ def main(cfg: DictConfig) -> dict:
     hp = OmegaConf.to_container(cfg, resolve=True, throw_on_missing=True)
 
     # SEED
-
-    L.seed_everything(cfg.seed, workers=True)
+    set_seed(cfg.seed)
 
     # MODEL
     model = instantiate(cfg.model)
-    # from IPython import embed; embed()
 
     # DATA
     datamodule = instantiate(cfg.datamodule)
@@ -31,42 +32,53 @@ def main(cfg: DictConfig) -> dict:
         callbacks.append(instantiate(cb_conf))
 
     loggers = []
+    # logger.info("Instantiating logger <{}>".format(cfg.logger._target_))
     for log_conf in cfg.loggers:
         logger = instantiate(cfg[log_conf])
         # wandb logger special setup
         if isinstance(logger, L.pytorch.loggers.WandbLogger):
             # deal with hangs when hp optim multirun training 
             # wandb.init(settings=wandb.Settings(start_method="thread"))
             # wandb requires dict not DictConfig
-            logger.experiment.config.update(hp)
+            logger.experiment.config.update(hp["datamodule"], allow_val_change=True)
+            logger.experiment.config.update(hp["model"], allow_val_change=True)
         loggers.append(logger)
+    
+    # print(json.dumps(hp, indent=4))
 
     # trainer
     profiler = instantiate(cfg.profiler)
-    trainer = instantiate(cfg.trainer, callbacks=callbacks, profiler=profiler, logger=[logger])
-    trainer.logger.log_hyperparams(hp)
-
-    # lr finder
-    # tuner = Tuner(trainer)
+    trainer = instantiate(cfg.trainer, callbacks=callbacks, profiler=profiler, logger=loggers)
+    # trainer.logger.log_hyperparams(hp)
 
-    # tuner.scale_batch_size(model, datamodule=datamodule, mode="power")
-    # lr_finder = tuner.lr_find(model,datamodule=datamodule)
-    # print(lr_finder.results)
-    # # Plot with
-    # fig = lr_finder.plot(suggest=True)
-    # fig.show()
-    # new_lr = lr_finder.suggestion()
-    # model.hparams.lr = new_lr
+    # batch size & lr optimization
+    tuner = Tuner(trainer)
+    if cfg.get("bs_finder"):
+        tuner.scale_batch_size(model, datamodule=datamodule, mode="power", init_val= 65536)
+        if isinstance(logger, L.pytorch.loggers.WandbLogger):
+            # bs is automatically updated in L datamodule but we need to manuallyupdate it in wandb
+            logger.experiment.config.update({"batch_size": datamodule.hparams.batch_size}, allow_val_change=True)
+
+    if cfg.get("lr_finder"):
+        lr_finder = tuner.lr_find(model, datamodule=datamodule)
+        print(lr_finder.results)
+        # Plot with
+        fig = lr_finder.plot(suggest=True)
+        fig.show()
+        new_lr = lr_finder.suggestion()
+        model.hparams.lr = new_lr
+        if isinstance(logger, L.pytorch.loggers.WandbLogger):
+            logger.experiment.config.update({"lr": new_lr}, allow_val_change=True)
 
 
     if cfg.get("train"):
         # trainer.fit(model=autoencoder_pl, train_dataloaders=train_dl, val_dataloaders=dev_dl, ckpt_path=cfg.get("ckpt_path"))
         trainer.fit(model, datamodule=datamodule, ckpt_path=cfg.get("ckpt_path"))
 
     # # TEST
-    # if cfg.get("test"):
+    if cfg.get("test"):
     #     # trainer.test(autoencoder_pl, dataloaders=test_dl)
-    #     trainer.test(datamodule=datamodule, ckpt_path="best")
+        trainer.test(datamodule=datamodule, ckpt_path="best")
 
     # wandb.finish()