Add ReduceLRonPlateau

Author: Piyush-555

main_bayesian.py

@@ -5,7 +5,7 @@
 
 import torch
 import numpy as np
-from torch.optim import Adam
+from torch.optim import Adam, lr_scheduler
 from torch.nn import functional as F
 
 import data
@@ -119,13 +119,14 @@ def run(dataset, net_type):
 
     criterion = metrics.ELBO(len(trainset)).to(device)
     optimizer = Adam(net.parameters(), lr=lr_start)
+    lr_sched = lr_scheduler.ReduceLROnPlateau(optimizer, patience=6, verbose=True)
     valid_loss_max = np.Inf
     for epoch in range(n_epochs):  # loop over the dataset multiple times
         cfg.curr_epoch_no = epoch
-        utils.adjust_learning_rate(optimizer, metrics.lr_linear(epoch, 0, n_epochs, lr_start))
 
         train_loss, train_acc, train_kl = train_model(net, optimizer, criterion, train_loader, num_ens=train_ens, beta_type=beta_type)
         valid_loss, valid_acc = validate_model(net, criterion, valid_loader, num_ens=valid_ens)
+        lr_sched.step(valid_loss)
 
         print('Epoch: {} \tTraining Loss: {:.4f} \tTraining Accuracy: {:.4f} \tValidation Loss: {:.4f} \tValidation Accuracy: {:.4f} \ttrain_kl_div: {:.4f}'.format(
             epoch, train_loss, train_acc, valid_loss, valid_acc, train_kl))

main_frequentist.py

@@ -6,7 +6,7 @@
 import torch
 import numpy as np
 import torch.nn as nn
-from torch.optim import Adam
+from torch.optim import Adam, lr_scheduler
 
 import data
 import utils
@@ -82,12 +82,13 @@ def run(dataset, net_type):
 
     criterion = nn.CrossEntropyLoss()
     optimizer = Adam(net.parameters(), lr=lr)
+    lr_sched = lr_scheduler.ReduceLROnPlateau(optimizer, patience=6, verbose=True)
     valid_loss_min = np.Inf
     for epoch in range(1, n_epochs+1):
-        utils.adjust_learning_rate(optimizer, metrics.lr_linear(epoch, 0, n_epochs, lr))
 
         train_loss, train_acc = train_model(net, optimizer, criterion, train_loader)
         valid_loss, valid_acc = validate_model(net, criterion, valid_loader)
+        lr_sched.step(valid_loss)
 
         train_loss = train_loss/len(train_loader.dataset)
         valid_loss = valid_loss/len(valid_loader.dataset)

metrics.py

@@ -14,10 +14,10 @@ def forward(self, input, target, kl, beta):
         return F.nll_loss(input, target, reduction='mean') * self.train_size + beta * kl
 
 
-def lr_linear(epoch_num, decay_start, total_epochs, start_value):
-    if epoch_num < decay_start:
-        return start_value
-    return start_value*float(total_epochs-epoch_num)/float(total_epochs-decay_start)
+# def lr_linear(epoch_num, decay_start, total_epochs, start_value):
+#     if epoch_num < decay_start:
+#         return start_value
+#     return start_value*float(total_epochs-epoch_num)/float(total_epochs-decay_start)
 
 
 def acc(outputs, targets):