Add a way to set optimizer hyperparameters quickly.

easytorch/core/optimizer_builder.py

@@ -37,6 +37,14 @@ def build_optim(optim_cfg: Dict, model: nn.Module) -> optim.Optimizer:
         optim_cfg (Dict): optimizer config
         model (nn.Module): model defined by user
 
+    Option:
+        Add parameters with special optimizer hyperparameters by set _optim attribute
+        example:
+        net = nn.Parameter(torch.zeros(10))
+        setattr(net, "_optim", {'lr': 0.01,"weight_decay": 0.0})
+
+        Information of optimizer will be printed. You can check it.
+
     Returns:
         optimizer (optim.Optimizer)
     """
@@ -48,8 +56,40 @@ def build_optim(optim_cfg: Dict, model: nn.Module) -> optim.Optimizer:
             optim_type = getattr(optim, optim_cfg['TYPE'])
         else:
             optim_type = getattr(easyoptim, optim_cfg['TYPE'])
+
+    # Obtain general parameters
     optim_param = optim_cfg['PARAM'].copy()
-    optimizer = optim_type(model.parameters(), **optim_param)
+
+    # All parameters in the model
+    all_parameters = list(model.parameters())
+
+    # General parameters don't contain the special _optim key
+    params = [p for p in all_parameters if not hasattr(p, "_optim")]
+
+    # Create an optimizer with the general parameters
+    optimizer = optim_type(params, **optim_param)
+
+    # Add parameters with special hyperparameters
+    hps = [getattr(p, "_optim") for p in all_parameters if hasattr(p, "_optim")]
+    hps = [
+        # Create unique special hyperparameters dicts
+        dict(s) for s in sorted(list(dict.fromkeys(frozenset(hp.items()) for hp in hps)))
+    ]
+    for hp in hps:
+        params = [p for p in all_parameters if getattr(p, "_optim", None) == hp]
+        optimizer.add_param_group(
+            {"params": params, **hp}
+        )
+
+    # Print optimizer info
+    keys = sorted(set([k for hp in hps for k in hp.keys()]))
+    for i, g in enumerate(optimizer.param_groups):
+        group_hps = {k: g.get(k, None) for k in keys}
+        print(' | '.join([
+                             f"Optimizer group {i}",
+                             f"{len(g['params'])} tensors",
+                         ] + [f"{k} {v}" for k, v in group_hps.items()]))
+
     return optimizer
 
 

easytorch/core/runner.py

@@ -25,6 +25,28 @@
 class Runner(metaclass=ABCMeta):
     """Base EasyTorch Runner
     """
+    """Base EasyTorch Runner
+    init_logger()
+    define_model() unrealized
+    build_train_dataset unrealized
+    build_val_dataset() unrealized
+    get_ckpt_path()
+    build_model()
+    get_ckpt_path()
+    save_model()
+    load_model_resume()
+    load_model()
+    train()
+    init_training()
+    on_epoch_start()
+    on_training_end()  only close tensorboard
+    train_iters() unrealized
+    backward()
+    validate()
+    init_validation()
+    on_validating_start() unrealized
+    on_validating_end() unrealized
+    """
 
     def __init__(self, cfg: Config):
         # default logger
@@ -362,6 +384,16 @@ def train(self, cfg: Config):
 
         self.on_training_end()
 
+    def init_optim(self, cfg: Config):
+        """Initialize optimizer
+
+        Args:
+            cfg (Dict): config
+        """
+        # create lr_scheduler
+        self.optim = build_optim(cfg['TRAIN.OPTIM'], self.model)
+        self.logger.info('Set optim: {}'.format(self.optim))
+
     def init_lr_scheduler(self, cfg: Config):
         """Initialize lr_scheduler
 
@@ -397,8 +429,7 @@ def init_training(self, cfg: Config):
         self.register_epoch_meter('train_time', 'train', '{:.2f} (s)', plt=False)
 
         # create optim
-        self.optim = build_optim(cfg['TRAIN.OPTIM'], self.model)
-        self.logger.info('Set optim: {}'.format(self.optim))
+        self.init_optim(cfg)
 
         # create lr_scheduler
         self.init_lr_scheduler(cfg)
@@ -433,7 +464,10 @@ def on_epoch_start(self, epoch: int):
         self.logger.info('Epoch {:d} / {:d}'.format(epoch, self.num_epochs))
         # update lr meter
         if self.scheduler is not None:
-            self.update_epoch_meter('lr', self.scheduler.get_last_lr()[0])
+            try:
+                self.update_epoch_meter('lr', self.scheduler.get_last_lr()[0])
+            except NotImplementedError:
+                self.update_epoch_meter('lr', self.scheduler.get_lr()[0])
 
         # set epoch for sampler in distributed mode
         # see https://pytorch.org/docs/stable/data.html
@@ -611,6 +645,8 @@ def save_best_model(self, epoch: int, metric_name: str, greater_best: bool = Tru
                 '{}_best_{}.pt'.format(self.model_name, metric_name.replace('/', '_'))
             )
             save_ckpt(ckpt_dict, ckpt_path, self.logger)
+            return True
+        return False
 
     @master_only
     def register_epoch_meter(self, name, meter_type, fmt='{:f}', plt=True):