improved the performance of main_bayesian.py

main_bayesian.py

@@ -19,71 +19,140 @@
 # CUDA settings
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
+
 def getModel(net_type, inputs, outputs, priors, layer_type, activation_type):
-    if (net_type == 'lenet'):
+    if net_type == 'lenet':
         return BBBLeNet(outputs, inputs, priors, layer_type, activation_type)
-    elif (net_type == 'alexnet'):
+    elif net_type == 'alexnet':
         return BBBAlexNet(outputs, inputs, priors, layer_type, activation_type)
-    elif (net_type == '3conv3fc'):
+    elif net_type == '3conv3fc':
         return BBB3Conv3FC(outputs, inputs, priors, layer_type, activation_type)
     else:
-        raise ValueError('Network should be either [LeNet / AlexNet / 3Conv3FC')
+        raise ValueError(
+            'Network should be either [LeNet / AlexNet / 3Conv3FC]'
+        )
 
 
-def train_model(net, optimizer, criterion, trainloader, num_ens=1, beta_type=0.1, epoch=None, num_epochs=None):
+def train_model(
+    net,
+    optimizer,
+    criterion,
+    trainloader,
+    num_ens=1,
+    beta_type=0.1,
+    epoch=None,
+    num_epochs=None
+):
     net.train()
+
     training_loss = 0.0
     accs = []
     kl_list = []
+
     for i, (inputs, labels) in enumerate(trainloader, 1):
 
         optimizer.zero_grad()
 
         inputs, labels = inputs.to(device), labels.to(device)
-        outputs = torch.zeros(inputs.shape[0], net.num_classes, num_ens).to(device)
+
+        outputs = torch.zeros(
+            inputs.shape[0],
+            net.num_classes,
+            num_ens
+        ).to(device)
 
         kl = 0.0
+
         for j in range(num_ens):
             net_out, _kl = net(inputs)
             kl += _kl
             outputs[:, :, j] = F.log_softmax(net_out, dim=1)
-        
+
         kl = kl / num_ens
         kl_list.append(kl.item())
+
         log_outputs = utils.logmeanexp(outputs, dim=2)
 
-        beta = metrics.get_beta(i-1, len(trainloader), beta_type, epoch, num_epochs)
+        beta = metrics.get_beta(
+            i - 1,
+            len(trainloader),
+            beta_type,
+            epoch,
+            num_epochs
+        )
+
         loss = criterion(log_outputs, labels, kl, beta)
+
         loss.backward()
         optimizer.step()
 
         accs.append(metrics.acc(log_outputs.data, labels))
-        training_loss += loss.cpu().data.numpy()
-    return training_loss/len(trainloader), np.mean(accs), np.mean(kl_list)
+        training_loss += loss.item()
 
+    return (
+        training_loss / len(trainloader),
+        np.mean(accs),
+        np.mean(kl_list)
+    )
 
-def validate_model(net, criterion, validloader, num_ens=1, beta_type=0.1, epoch=None, num_epochs=None):
+
+def validate_model(
+    net,
+    criterion,
+    validloader,
+    num_ens=1,
+    beta_type=0.1,
+    epoch=None,
+    num_epochs=None
+):
     """Calculate ensemble accuracy and NLL Loss"""
-    net.train()
+
+    net.eval()
+
     valid_loss = 0.0
     accs = []
 
-    for i, (inputs, labels) in enumerate(validloader):
-        inputs, labels = inputs.to(device), labels.to(device)
-        outputs = torch.zeros(inputs.shape[0], net.num_classes, num_ens).to(device)
-        kl = 0.0
-        for j in range(num_ens):
-            net_out, _kl = net(inputs)
-            kl += _kl
-            outputs[:, :, j] = F.log_softmax(net_out, dim=1).data
+    with torch.no_grad():
 
-        log_outputs = utils.logmeanexp(outputs, dim=2)
+        for i, (inputs, labels) in enumerate(validloader):
+
+            inputs, labels = inputs.to(device), labels.to(device)
+
+            outputs = torch.zeros(
+                inputs.shape[0],
+                net.num_classes,
+                num_ens
+            ).to(device)
+
+            kl = 0.0
 
-        beta = metrics.get_beta(i-1, len(validloader), beta_type, epoch, num_epochs)
-        valid_loss += criterion(log_outputs, labels, kl, beta).item()
-        accs.append(metrics.acc(log_outputs, labels))
+            for j in range(num_ens):
+                net_out, _kl = net(inputs)
+                kl += _kl
+                outputs[:, :, j] = F.log_softmax(net_out, dim=1)
 
-    return valid_loss/len(validloader), np.mean(accs)
+            kl = kl / num_ens
+
+            log_outputs = utils.logmeanexp(outputs, dim=2)
+
+            beta = metrics.get_beta(
+                i - 1,
+                len(validloader),
+                beta_type,
+                epoch,
+                num_epochs
+            )
+
+            valid_loss += criterion(
+                log_outputs,
+                labels,
+                kl,
+                beta
+            ).item()
+
+            accs.append(metrics.acc(log_outputs, labels))
+
+    return valid_loss / len(validloader), np.mean(accs)
 
 
 def run(dataset, net_type):
@@ -103,40 +172,126 @@ def run(dataset, net_type):
     beta_type = cfg.beta_type
 
     trainset, testset, inputs, outputs = data.getDataset(dataset)
+
     train_loader, valid_loader, test_loader = data.getDataloader(
-        trainset, testset, valid_size, batch_size, num_workers)
-    net = getModel(net_type, inputs, outputs, priors, layer_type, activation_type).to(device)
+        trainset,
+        testset,
+        valid_size,
+        batch_size,
+        num_workers
+    )
+
+    net = getModel(
+        net_type,
+        inputs,
+        outputs,
+        priors,
+        layer_type,
+        activation_type
+    ).to(device)
 
     ckpt_dir = f'checkpoints/{dataset}/bayesian'
-    ckpt_name = f'checkpoints/{dataset}/bayesian/model_{net_type}_{layer_type}_{activation_type}.pt'
+    ckpt_name = (
+        f'checkpoints/{dataset}/bayesian/'
+        f'model_{net_type}_{layer_type}_{activation_type}.pt'
+    )
 
-    if not os.path.exists(ckpt_dir):
-        os.makedirs(ckpt_dir, exist_ok=True)
+    os.makedirs(ckpt_dir, exist_ok=True)
 
     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
 
-        train_loss, train_acc, train_kl = train_model(net, optimizer, criterion, train_loader, num_ens=train_ens, beta_type=beta_type, epoch=epoch, num_epochs=n_epochs)
-        valid_loss, valid_acc = validate_model(net, criterion, valid_loader, num_ens=valid_ens, beta_type=beta_type, epoch=epoch, num_epochs=n_epochs)
+    # Fixed for new PyTorch versions
+    lr_sched = lr_scheduler.ReduceLROnPlateau(
+        optimizer,
+        patience=6
+    )
+
+    valid_loss_max = np.inf
+
+    for epoch in range(n_epochs):
+
+        train_loss, train_acc, train_kl = train_model(
+            net,
+            optimizer,
+            criterion,
+            train_loader,
+            num_ens=train_ens,
+            beta_type=beta_type,
+            epoch=epoch,
+            num_epochs=n_epochs
+        )
+
+        valid_loss, valid_acc = validate_model(
+            net,
+            criterion,
+            valid_loader,
+            num_ens=valid_ens,
+            beta_type=beta_type,
+            epoch=epoch,
+            num_epochs=n_epochs
+        )
+
+        old_lr = optimizer.param_groups[0]['lr']
+
         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))
+        new_lr = optimizer.param_groups[0]['lr']
+
+        if old_lr != new_lr:
+            print(
+                f"Learning rate reduced from "
+                f"{old_lr:.6e} to {new_lr:.6e}"
+            )
+
+        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
+            )
+        )
 
-        # save model if validation accuracy has increased
         if valid_loss <= valid_loss_max:
-            print('Validation loss decreased ({:.6f} --> {:.6f}).  Saving model ...'.format(
-                valid_loss_max, valid_loss))
+            print(
+                'Validation loss decreased '
+                '({:.6f} --> {:.6f}). Saving model ...'.format(
+                    valid_loss_max,
+                    valid_loss
+                )
+            )
+
             torch.save(net.state_dict(), ckpt_name)
             valid_loss_max = valid_loss
 
+
 if __name__ == '__main__':
-    parser = argparse.ArgumentParser(description = "PyTorch Bayesian Model Training")
-    parser.add_argument('--net_type', default='lenet', type=str, help='model')
-    parser.add_argument('--dataset', default='MNIST', type=str, help='dataset = [MNIST/CIFAR10/CIFAR100]')
+
+    parser = argparse.ArgumentParser(
+        description="PyTorch Bayesian Model Training"
+    )
+
+    parser.add_argument(
+        '--net_type',
+        default='lenet',
+        type=str,
+        help='model'
+    )
+
+    parser.add_argument(
+        '--dataset',
+        default='MNIST',
+        type=str,
+        help='dataset = [MNIST/CIFAR10/CIFAR100]'
+    )
+
     args = parser.parse_args()
 
-    run(args.dataset, args.net_type)
+    run(args.dataset, args.net_type)