BayesCNN_pytorch_BBalpha.py

import numpy as np
import os, pickle, sys, time
import timeit

## import pytorch modules
import torch
from torch.autograd import Variable
import torch.nn.functional as Func
import torch.nn as nn
from torchvision import datasets, transforms

def log_sum_exp(tensor, dim  = None):
    xmax, _ = torch.max(tensor, dim = dim, keepdim = True)
    xmax_, _ = torch.max(tensor, dim = dim)
    return xmax_ + torch.log(torch.sum(torch.exp(tensor - xmax), dim = dim))

class MC_Loss(nn.Module):
    '''
    bbalpha softmax cross entropy with mc_logits
    '''

    def __init__(self, alpha= 1.0, k_mc = 20):
        super(MC_Loss, self).__init__()
        self.alpha = alpha
        if alpha == 0:
            print "alpha = 0"
        self.k_mc = k_mc

    def forward(self, mc_logits, y_true):
        len_y_true = len(y_true)
        y_true = y_true.expand(self.k_mc, -1, -1).contiguous().permute(1, 0, 2)
        if self.alpha != 0:
            # log(p_ij), p_ij = softmax(logit_ij)
            #assert mc_logits.ndim == 3
            temp, _ = torch.max(mc_logits, dim=2, keepdim=True)
            mc_log_softmax = mc_logits - temp
            mc_log_softmax = mc_log_softmax - torch.log(torch.sum(torch.exp(mc_log_softmax), dim=2, keepdim=True))
            mc_ll = torch.sum(y_true * mc_log_softmax, dim = -1)  # N x K
            # print mc_ll.size()
            out = - 1. / self.alpha * (log_sum_exp(self.alpha * mc_ll, 1) + np.log(1.0 / self.k_mc))
            # print out.size()
            # sys.exit()
            return torch.sum(out)
        else:
            predictions = Func.log_softmax(mc_logits, dim=2)
            # print y_true, predictions
            out = - torch.sum(torch.mean(y_true * predictions, dim=1))
            # print out.size()
            # sys.exit()
            return out

class BayesCNN(nn.Module):
    def __init__(self, alpha=1.0, k_mc=20, wd = 10**-6):
        super(BayesCNN, self).__init__()
        # input is 28x28
        # padding=2 for same padding
        self.wd = wd
        self.conv1 = nn.Conv2d(1, 32, 3, padding=1)
        # feature map size is 14*14 by pooling
        # padding=2 for same padding
        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
        # feature map size is 7*7 by pooling
        self.fc1 = nn.Linear(64*14*14, 128)
        self.fc2 = nn.Linear(128, 10)
        self.k_mc = k_mc
        self.alpha = alpha
        self.bbalpha_loss = MC_Loss(self.alpha, self.k_mc)

    def lin(self, x):
        x = Func.relu(self.conv1.forward(x))
        # print x.size()
        x = Func.relu(self.conv2.forward(x))
        x = Func.max_pool2d(x, 2)
        x = Func.dropout(x, p = 0.5, training=self.training)
        x = x.view(-1, 64*14*14)   # reshape Variable
        x = Func.relu(self.fc1.forward(x))
        x = Func.dropout(x, p = 0.5, training=self.training)
        out = self.fc2.forward(x)
        return out

    def generate_MC_samples(self, x):
        # k_mc: number of samples
        if self.k_mc == 1:
            out = self.lin(x)
            mc_logits = out.view(len(out), 1, -1) # nb_batch x K_mc x nb_classes
            # print mc_logits.size()
            return mc_logits
        else:
            output_list = []
            for _ in xrange(self.k_mc):
                output_list += [self.lin(x)]  # THIS IS BAD!!! we create new dense layers at every call!!!!
            # print len(output_list)
            # sys.exit()
            output = torch.stack(output_list) # K_mc x nb_batch x nb_classes
            # print output.size()
            mc_logits = output.permute(1, 0, 2) # nb_batch x K_mc x nb_classes
            # print mc_logits.size()
            # sys.exit()
            return mc_logits

    def forward(self, x):
        lin_out = self.generate_MC_samples(x)
        out = Func.softmax(lin_out, dim=-1)
        # print out.size()
        out = torch.mean(out, dim=1).squeeze()
        return out

    def cal_bbalpha_loss(self, x, true_y): # negative log likelihood
        mc_logits = self.generate_MC_samples(x)
        loss = self.bbalpha_loss.forward(mc_logits, true_y)
        # negative log-likelihood
        return loss

    def cal_priors(self): # dropout may cause two forwards different
        prior = 0
        for param in self.parameters():
            prior = prior + self.wd * torch.sum(param**2)
        return prior

    def cal_npos(self, x, true_y): # negative log posterior
        return self.cal_bbalpha_loss(x, true_y) + self.cal_priors()



def main():

    """ Step 0: Compiling Model """
    # torch.manual_seed(123)
    # print bayescnn

    batch_size = 150
    num_epochs = 5

    bayescnn = BayesCNN(alpha = 0.5, k_mc = 3)
    optimizer = torch.optim.SGD(bayescnn.parameters(), lr=0.001)

    if torch.cuda.is_available():
        print "Run on GPU"
        bayescnn.cuda()
    else:
        print "Run on CPU"


    """ Step 1: Preprocessing """
    train_loader = torch.utils.data.DataLoader(
        datasets.MNIST('data', train=True, download=True, transform=transforms.ToTensor()),
        batch_size=batch_size, shuffle=True)

    test_loader = torch.utils.data.DataLoader(
        datasets.MNIST('data', train=False, transform=transforms.ToTensor()),
        batch_size=100)

    # '''
    """ step 2: Training """
    bayescnn.train()
    start = timeit.default_timer()
    for epoch in xrange(num_epochs):
        for i, (images, labels) in enumerate(train_loader):
            # print images.size()
            # print labels.size()
            images = Variable(images)
            d_labels = Variable(torch.zeros(len(images), 10).scatter_(1, labels.view(-1, 1), 1))
            # print labels, torch.max(d_labels, 1)[1]
            if torch.cuda.is_available():
                images = images.cuda()
                d_labels = d_labels.cuda()

            # Forward + Backward + Optimize
            bayescnn.zero_grad()
            #loss = bayescnn.cal_nllloss(images, labels)
            loss = bayescnn.cal_npos(images, d_labels)
            loss.backward()
            optimizer.step()

            if (i+1) % 100 == 0:
                print 'Epoch [{}/{}], Iter [{}/{}]'.format(epoch+1, num_epochs, i+1, len(train_loader))

    stop = timeit.default_timer()
    print("Time: ", stop - start )
    # '''

    """ Step 2: Model test """
    bayescnn.eval()  # Change model to 'eval' mode (BN uses moving mean/var).
    correct = 0
    total = 0
    for images, labels in test_loader:
        images = Variable(images)
        # labels = Variable(labels)
        if torch.cuda.is_available():
            images = images.cuda()
            labels = labels.cuda()
        outputs = bayescnn.forward(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum()

    print 'Test Accuracy of the model on the 10000 test images: {:0.2f}%'.format(100.0 * correct / total)

if __name__ == "__main__":
    main()