utils: Implement get_scheduler()

Author: synml

config.yaml

@@ -16,6 +16,7 @@ UNet:
     weight_decay: 0.00001
   scheduler:
     name: ReduceLROnPlateau
+    patience: 5
     min_lr: 0.0001
   pretrained_weights: weights/UNet_best.pth
 
@@ -29,6 +30,7 @@ Backbone:
     weight_decay: 0.00001
   scheduler:
     name: ReduceLROnPlateau
+    patience: 5
     min_lr: 0.0001
   pretrained_weights: weights/Backbone_best.pth
 
@@ -42,5 +44,6 @@ Proposed:
     weight_decay: 0.00001
   scheduler:
     name: ReduceLROnPlateau
+    patience: 5
     min_lr: 0.0001
   pretrained_weights: weights/Proposed_best.pth

train.py

@@ -23,8 +23,8 @@
 
     # 3. Loss function, optimizer, lr scheduler
     criterion = nn.CrossEntropyLoss()
-    optimizer = torch.optim.Adam(model.parameters(), lr=config[config['model']]['lr'])
-    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, min_lr=0.0001)
+    optimizer = utils.get_optimizer(config, model)
+    scheduler = utils.get_scheduler(config, optimizer)
 
     # 4. Tensorboard
     writer = torch.utils.tensorboard.SummaryWriter(os.path.join('runs', config['model']))

utils.py

@@ -44,7 +44,7 @@ def get_model(config: dict, pretrained=False) -> torch.nn.Module:
     return model
 
 
-def get_optimizer(config: dict, model: torch.nn.Module):
+def get_optimizer(config: dict, model: torch.nn.Module) -> torch.optim.Optimizer:
     cfg_optim: dict = config[config['model']]['optimizer']
 
     if cfg_optim['name'] == 'SGD':
@@ -58,6 +58,18 @@ def get_optimizer(config: dict, model: torch.nn.Module):
     return optimizer
 
 
+def get_scheduler(config: dict, optimizer: torch.optim.Optimizer):
+    cfg_scheduler: dict = config[config['model']]['scheduler']
+
+    if cfg_scheduler['name'] == 'ReduceLROnPlateau':
+        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=cfg_scheduler['patience'],
+                                                               min_lr=cfg_scheduler['min_lr'])
+    else:
+        raise NameError('Wrong scheduler name.')
+
+    return scheduler
+
+
 class Cityscapes:
     def __init__(self, config: dict):
         self.config = config