refactor lr find & 1-cycle

Author: slegroux

nbs/image.datasets.ipynb

@@ -925,6 +925,9 @@
     "            return list(self.train_ds[0][0].shape[-2:])\n",
     "        raise RuntimeError(\"train_ds is not initialized. Call prepare_data() first.\")\n",
     "\n",
+    "    @property\n",
+    "    def name(self)->str:\n",
+    "        return self.hparams.name\n",
     "\n",
     "    @property\n",
     "    def num_classes(self) -> int: # num of classes in dataset\n",

nbs/models.conv.ipynb

@@ -60,7 +60,7 @@
     "\n",
     "from matplotlib import pyplot as plt\n",
     "import pandas as pd\n",
-    "from typing import List, Optional, Type, Callable\n",
+    "from typing import List, Optional, Type, Callable, Any\n",
     "\n",
     "from nimrod.utils import get_device, set_seed\n",
     "from nimrod.models.core import Classifier\n",
@@ -1659,7 +1659,7 @@
     "            nnet:ConvNet, # model\n",
     "            num_classes:int, # number of classes\n",
     "            optimizer:Callable[...,torch.optim.Optimizer], # optimizer\n",
-    "            scheduler:Callable[...,torch.optim.lr_scheduler], # scheduler\n",
+    "            scheduler: Optional[Callable[...,Any]]=None, # scheduler\n",
     "            ):\n",
     "\n",
     "        logger.info(\"ConvNetX: init\")\n",

nbs/models.diffusion.ipynb

@@ -745,9 +745,13 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "python3",
+   "display_name": "nimrod",
    "language": "python",
    "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.11.8"
   }
  },
  "nbformat": 4,

nimrod/_modidx.py

@@ -174,6 +174,8 @@
                                                                                       'nimrod/image/datasets.py'),
                                        'nimrod.image.datasets.ImageDataModule.label_names': ( 'image.datasets.html#imagedatamodule.label_names',
                                                                                               'nimrod/image/datasets.py'),
+                                       'nimrod.image.datasets.ImageDataModule.name': ( 'image.datasets.html#imagedatamodule.name',
+                                                                                       'nimrod/image/datasets.py'),
                                        'nimrod.image.datasets.ImageDataModule.num_classes': ( 'image.datasets.html#imagedatamodule.num_classes',
                                                                                               'nimrod/image/datasets.py'),
                                        'nimrod.image.datasets.ImageDataModule.prepare_data': ( 'image.datasets.html#imagedatamodule.prepare_data',

nimrod/image/datasets.py

@@ -323,6 +323,9 @@ def dim(self)->List[int]:
             return list(self.train_ds[0][0].shape[-2:])
         raise RuntimeError("train_ds is not initialized. Call prepare_data() first.")
 
+    @property
+    def name(self)->str:
+        return self.hparams.name
 
     @property
     def num_classes(self) -> int: # num of classes in dataset

nimrod/models/conv.py

@@ -20,7 +20,7 @@
 
 from matplotlib import pyplot as plt
 import pandas as pd
-from typing import List, Optional, Type, Callable
+from typing import List, Optional, Type, Callable, Any
 
 from ..utils import get_device, set_seed
 from .core import Classifier
@@ -260,7 +260,7 @@ def __init__(
             nnet:ConvNet, # model
             num_classes:int, # number of classes
             optimizer:Callable[...,torch.optim.Optimizer], # optimizer
-            scheduler:Callable[...,torch.optim.lr_scheduler], # scheduler
+            scheduler: Optional[Callable[...,Any]]=None, # scheduler
             ):
 
         logger.info("ConvNetX: init")

tutorials/convnet.ipynb

@@ -408,11 +408,14 @@
    "source": [
     "# LR_FINDER\n",
     "\n",
-    "cfg_model.nnet.n_features = [1, 8, 16, 32, 16]\n",
+    "\n",
     "N_EPOCHS = 5\n",
+    "# set sched total steps\n",
     "cfg_sched.total_steps = len(dm.train_dataloader()) * N_EPOCHS\n",
-    "\n",
     "scheduler = instantiate(cfg_sched)\n",
+    "\n",
+    "# instantiate model with scheduler\n",
+    "cfg_model.nnet.n_features = [1, 8, 16, 32, 16]\n",
     "model = instantiate(cfg_model)(optimizer=optimizer, scheduler=scheduler)\n",
     "\n",
     "trainer = Trainer(\n",
@@ -748,14 +751,16 @@
    "source": [
     "# ONE-CYCLE TRAIN\n",
     "\n",
-    "cfg_model.nnet.n_features = [1, 8, 16, 32, 16]\n",
     "N_EPOCHS = 5\n",
+    "\n",
     "cfg_sched.total_steps = len(dm.train_dataloader()) * N_EPOCHS\n",
     "cfg_sched.max_lr = lr_finder.suggestion()\n",
     "\n",
+    "cfg_model.nnet.n_features = [1, 8, 16, 32, 16]\n",
+    "\n",
     "wandb_logger = WandbLogger(\n",
     "    project=\"MNIST Classification\",\n",
-    "    name=f\"resnetx-bs:{dm.batch_size}-epochs:{N_EPOCHS}-features:{cfg_model.nnet.n_features}\",\n",
+    "    name=f\"ResnetX-bs:{dm.batch_size}-epochs:{N_EPOCHS}-features:{cfg_model.nnet.n_features}\",\n",
     "    save_dir='wandb',\n",
     "    entity='slegroux',\n",
     "    tags=['arch', 'dev'],\n",