add lr_step

Author: ankitdhall

network/ethec_experiments.py

@@ -212,11 +212,12 @@ def __init__(self, data_loaders, labelmap, criterion, lr,
                  load_wt=False,
                  model_name=None,
                  optimizer_method='adam',
-                 use_grayscale=False):
+                 use_grayscale=False,
+                 lr_step=[]):
 
         CIFAR10.__init__(self, data_loaders, labelmap, criterion, lr, batch_size, evaluator, experiment_name,
                          experiment_dir, n_epochs, eval_interval, feature_extracting, use_pretrained,
-                         load_wt, model_name, optimizer_method)
+                         load_wt, model_name, optimizer_method, lr_step=lr_step)
 
         if use_grayscale:
             if model_name in ['alexnet', 'vgg']:
@@ -421,7 +422,8 @@ def ETHEC_train_model(arguments):
                                         load_wt=arguments.resume,
                                         model_name=arguments.model,
                                         optimizer_method=arguments.optimizer_method,
-                                        use_grayscale=arguments.use_grayscale)
+                                        use_grayscale=arguments.use_grayscale,
+                                        lr_step=arguments.lr_step)
     ETHEC_trainer.prepare_model()
     #if arguments.use_2d and arguments.resume:
     #    ETHEC_trainer.plot_label_representations()
@@ -458,6 +460,7 @@ def ETHEC_train_model(arguments):
                         required=True)
     parser.add_argument("--level_weights", help='List of weights for each level', nargs=4, default=None, type=float)
     parser.add_argument("--use_2d", help='Use model with 2d features', action='store_true')
+    parser.add_argument("--lr_step", help='List of epochs to make multiple lr by 0.1', nargs='*', default=[], type=int)
     args = parser.parse_args()
 
     ETHEC_train_model(args)

network/experiment.py

@@ -23,7 +23,7 @@
 class Experiment:
 
     def __init__(self, model, dataloaders, criterion, classes, experiment_name, n_epochs, eval_interval, batch_size,
-                 exp_dir, load_wt, evaluator):
+                 exp_dir, load_wt, evaluator, lr_step=[]):
         self.epoch = 0
         self.exp_dir = exp_dir
         self.load_wt = load_wt
@@ -52,6 +52,8 @@ def __init__(self, model, dataloaders, criterion, classes, experiment_name, n_ep
 
         self.writer = SummaryWriter(log_dir=os.path.join(self.log_dir, 'tensorboard'))
 
+        self.lr_step = lr_step
+
     @staticmethod
     def make_dir_if_non_existent(dir):
         if not os.path.exists(dir):
@@ -152,6 +154,7 @@ def pass_samples(self, phase, save_to_tensorboard=True):
 
     def run_model(self, optimizer):
         self.optimizer = optimizer
+        scheduler = torch.optim.lr_scheduler.MultiStepLR(self.optimizer, milestones=self.lr_step, gamma=0.1)
 
         if self.load_wt:
             self.find_existing_weights()
@@ -174,6 +177,8 @@ def run_model(self, optimizer):
                 self.pass_samples(phase='test')
                 self.eval.disable_plotting()
 
+            scheduler.step()
+
         time_elapsed = time.time() - since
         print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))
         print('Best val score: {:4f}'.format(self.best_score))

network/finetuner.py

@@ -98,7 +98,8 @@ def __init__(self, data_loaders, labelmap, criterion, lr,
                  use_pretrained=True,
                  load_wt=False,
                  model_name=None,
-                 optimizer_method='adam'):
+                 optimizer_method='adam',
+                 lr_step=[]):
 
         self.classes = labelmap.classes
         self.n_classes = labelmap.n_classes
@@ -111,6 +112,7 @@ def __init__(self, data_loaders, labelmap, criterion, lr,
         self.optimal_thresholds = np.zeros(self.n_classes)
         self.optimizer_method = optimizer_method
         self.labelmap = labelmap
+        self.lr_step = lr_step
 
         if model_name == 'alexnet':
             model = models.alexnet(pretrained=use_pretrained)
@@ -127,7 +129,7 @@ def __init__(self, data_loaders, labelmap, criterion, lr,
 
         Experiment.__init__(self, model, data_loaders, criterion, self.classes, experiment_name, n_epochs,
                             eval_interval,
-                            batch_size, experiment_dir, load_wt, evaluator)
+                            batch_size, experiment_dir, load_wt, evaluator, lr_step=self.lr_step)
         self.model_name = model_name
 
     def prepare_model(self, loading=False):