fix(optim): fix optim params adj util to mod any

Author: okunator

cellseg_models_pytorch/optimizers/tests/test_optim_setup.py

@@ -7,15 +7,11 @@
 @pytest.mark.parametrize("optim", ["adam"])
 def test_optim_setup(optim):
     model = cellpose_base(type_classes=3)
-    params = adjust_optim_params(
-        model,
-        encoder_lr=0.5,
-        encoder_wd=0.4,
-        decoder_lr=0.3,
-        decoder_wd=0.2,
-        remove_bias_wd=True,
-    )
-
+    optim_params = {
+        "encoder": {"lr": 0.004, "weight_decay": 0.3},
+        "decoder": {"lr": 0.003, "weight_decay": 0.2},
+    }
+    params = adjust_optim_params(model, optim_params)
     optimizer = OPTIM_LOOKUP[optim](params)
 
     assert all([p_g["lr"] == p["lr"] for p_g, p in zip(optimizer.param_groups, params)])

cellseg_models_pytorch/optimizers/utils.py

@@ -1,67 +1,38 @@
-from typing import Dict, List
+from typing import Any, Dict, List
 
 import torch.nn as nn
 
 
 def adjust_optim_params(
-    model: nn.Module,
-    encoder_lr: float,
-    encoder_wd: float,
-    decoder_lr: float,
-    decoder_wd: float,
-    remove_bias_wd: bool = True,
+    model: nn.Module, optim_params: Dict[str, Dict[str, Any]]
 ) -> List[Dict[str, Dict]]:
     """Adjust the learning parameters for optimizer.
 
-    1. Adjust learning rate and weight decay in the pre-trained
-        encoder and decoders.
-    2. Remove weight decay from bias terms to reduce overfitting.
-
-    "Bag of Tricks for Image Classification with Convolutional Neural Networks"
-    - https://arxiv.org/pdf/1812.01187
-
     Parameters
     ----------
         model : nn.Module
             The encoder-decoder segmentation model.
-        encoder_lr : float
-            Learning rate of the model encoder.
-        encoder_wd : float
-            Weight decay for the model encoder.
-        decoder_lr : float
-            Learning rate of the model decoder.
-        decoder_wd : float
-            Weight decay for the model decoder.
-        remove_bias_wd : bool, default=True
-            If True, the weight decay from the bias terms is removed from the model
-            params. Ignored if `remove_wd`=True.
+        optim_params : Dict[str, Dict[str, Any]]
+            optim paramas like learning rates, weight decays etc for diff parts of
+            the network. E.g.
+            {"encoder": {"weight_decay: 0.1, "lr":0.1}, "sem": {"lr": 0.1}}
 
     Returns
     -------
         List[Dict[str, Dict]]:
             a list of kwargs (str, Dict pairs) containing the model params.
     """
     params = list(model.named_parameters())
-    encoder_params = {"encoder": {"lr": encoder_lr, "weight_decay": encoder_wd}}
-    decoder_params = {"decoder": {"lr": decoder_lr, "weight_decay": decoder_wd}}
 
     adjust_params = []
     for name, parameters in params:
         opts = {}
-        for enc, enc_opts in encoder_params.items():
-            if enc in name:
-                for key, item in enc_opts.items():
-                    opts[key] = item
 
-        for dec, dec_opts in decoder_params.items():
-            if dec in name:
-                for key, item in dec_opts.items():
+        for block, block_params in optim_params.items():
+            if block in name:
+                for key, item in block_params.items():
                     opts[key] = item
 
-        if remove_bias_wd:
-            if name.endswith("bias"):
-                opts["weight_decay"] = 0.0
-
         adjust_params.append({"params": parameters, **opts})
 
     return adjust_params

changelog.d/20221006_115604_oskari.lehtonen.md

@@ -0,0 +1,3 @@
+## Fixes
+
+- Modify the optimizer adjustment utility function to adjust any optim/weight params.