update regularizer

nathanhubens · Mar 20, 2024 · 70d9977 · 70d9977
1 parent e8fbc1b
commit 70d9977
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 42 deletions.
diff --git a/fasterai/_modidx.py b/fasterai/_modidx.py
@@ -163,14 +163,6 @@
                                              'fasterai.quantize.quantizer.Quantizer.quantize': ( 'quantize.quantizer.html#quantizer.quantize',
                                                                                                  'fasterai/quantize/quantizer.py')},
             'fasterai.regularize.all': {},
-            'fasterai.regularize.regularization_callback': { 'fasterai.regularize.regularization_callback.RegularizationCallback': ( 'regularize.regularizer.html#regularizationcallback',
-                                                                                                                                     'fasterai/regularize/regularization_callback.py'),
-                                                             'fasterai.regularize.regularization_callback.RegularizationCallback.__init__': ( 'regularize.regularizer.html#regularizationcallback.__init__',
-                                                                                                                                              'fasterai/regularize/regularization_callback.py'),
-                                                             'fasterai.regularize.regularization_callback.RegularizationCallback.after_loss': ( 'regularize.regularizer.html#regularizationcallback.after_loss',
-                                                                                                                                                'fasterai/regularize/regularization_callback.py'),
-                                                             'fasterai.regularize.regularization_callback.RegularizationCallback.get_norm': ( 'regularize.regularizer.html#regularizationcallback.get_norm',
-                                                                                                                                              'fasterai/regularize/regularization_callback.py')},
             'fasterai.regularize.regularize_callback': { 'fasterai.regularize.regularize_callback.RegularizeCallback': ( 'regularize.regularizer.html#regularizecallback',
                                                                                                                          'fasterai/regularize/regularize_callback.py'),
                                                          'fasterai.regularize.regularize_callback.RegularizeCallback.__init__': ( 'regularize.regularizer.html#regularizecallback.__init__',

diff --git a/fasterai/regularize/regularization_callback.py b/fasterai/regularize/regularization_callback.py
diff --git a/fasterai/regularize/regularize_callback.py b/fasterai/regularize/regularize_callback.py
@@ -16,7 +16,7 @@
 # %% ../../nbs/05_regularize.regularizer.ipynb 4
 class RegularizeCallback(Callback):
     "Callback to apply grouped weight decay"
-    def __init__(self, g, wd=0.01):
+    def __init__(self, g, wd=0.01, layer_type=nn.Conv2d):
         store_attr()
 
     def after_loss(self):
@@ -25,4 +25,4 @@ def after_loss(self):
         self.learn.loss = self.learn.loss_grad.clone()
 
     def get_norm(self):
-        return self.wd*torch.stack([large_final.f(m.weight)[None].sum(Granularities.get_dim(m, self.g)).sum() for m in self.learn.modules() if isinstance(m, nn.Conv2d)]).sum()
+        return self.wd*torch.stack([large_final.f(m.weight)[None].sum(Granularities.get_dim(m, self.g)).sum() for m in self.learn.modules() if isinstance(m, self.layer_type)]).sum()
diff --git a/nbs/05_regularize.regularizer.ipynb b/nbs/05_regularize.regularizer.ipynb
@@ -64,7 +64,7 @@
     "#| export\n",
     "class RegularizeCallback(Callback):\n",
     "    \"Callback to apply grouped weight decay\"\n",
-    "    def __init__(self, g, wd=0.01):\n",
+    "    def __init__(self, g, wd=0.01, layer_type=nn.Conv2d):\n",
     "        store_attr()\n",
     "\n",
     "    def after_loss(self):\n",
@@ -73,7 +73,7 @@
     "        self.learn.loss = self.learn.loss_grad.clone()\n",
     "        \n",
     "    def get_norm(self):\n",
-    "        return self.wd*torch.stack([large_final.f(m.weight)[None].sum(Granularities.get_dim(m, self.g)).sum() for m in self.learn.modules() if isinstance(m, nn.Conv2d)]).sum()"
+    "        return self.wd*torch.stack([large_final.f(m.weight)[None].sum(Granularities.get_dim(m, self.g)).sum() for m in self.learn.modules() if isinstance(m, self.layer_type)]).sum()"
    ]
   },
   {
@@ -89,7 +89,8 @@
        "\n",
        "### RegularizeCallback\n",
        "\n",
-       ">      RegularizeCallback (g, wd=0.01)\n",
+       ">      RegularizeCallback (g, wd=0.01, layer_type=<class\n",
+       ">                          'torch.nn.modules.conv.Conv2d'>)\n",
        "\n",
        "Callback to apply grouped weight decay"
       ],
@@ -98,7 +99,8 @@
        "\n",
        "### RegularizeCallback\n",
        "\n",
-       ">      RegularizeCallback (g, wd=0.01)\n",
+       ">      RegularizeCallback (g, wd=0.01, layer_type=<class\n",
+       ">                          'torch.nn.modules.conv.Conv2d'>)\n",
        "\n",
        "Callback to apply grouped weight decay"
       ]