diff --git a/checkpoints/MNIST/bayesian/model_lenet_lrt_softplus.pt b/checkpoints/MNIST/bayesian/model_lenet_lrt_softplus.pt new file mode 100644 index 00000000..1e6122d7 Binary files /dev/null and b/checkpoints/MNIST/bayesian/model_lenet_lrt_softplus.pt differ diff --git a/main_bayesian.py b/main_bayesian.py index b14c6433..60961ab0 100755 --- a/main_bayesian.py +++ b/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) \ No newline at end of file