SGDR_WRNs.py

"""
Lasagne implementation of SGDR on WRNs from "SGDR: Stochastic Gradient Descent with Restarts" (http://arxiv.org/abs/XXXX)
This code is based on Lasagne Recipes available at
https://github.com/Lasagne/Recipes/blob/master/papers/deep_residual_learning/Deep_Residual_Learning_CIFAR-10.py
and on WRNs implementation by Florian Muellerklein available at
https://gist.github.com/FlorianMuellerklein/3d9ba175038a3f2e7de3794fa303f1ee

"""

from __future__ import print_function

import sys
import os
import time
import string
import random
import pickle

import numpy as np
import theano
import theano.tensor as T
import lasagne
import math

from lasagne.nonlinearities import rectify, softmax
from lasagne.layers import InputLayer, DenseLayer, DropoutLayer, batch_norm, BatchNormLayer
from lasagne.layers import ElemwiseSumLayer, NonlinearityLayer, GlobalPoolLayer
from lasagne.layers.dnn import Conv2DDNNLayer as ConvLayer
from lasagne.init import HeNormal
from lasagne.layers import Conv2DLayer as ConvLayer


# for the larger networks (n>=9), we need to adjust pythons recursion limit
sys.setrecursionlimit(10000)

# ##################### Load data from CIFAR datasets #######################
# this code assumes the CIFAR dataset files have been extracted in current working directory
# from 'https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz' for CIFAR-10
# from 'https://www.cs.toronto.edu/~kriz/cifar-100-python.tar.gz' for CIFAR-100

def unpickle(file):
    import cPickle
    fo = open(file, 'rb')
    dict = cPickle.load(fo)
    fo.close()
    return dict

def load_data(dataset):
    xs = []
    ys = []
    if dataset == 'CIFAR-10':
        for j in range(5):
          d = unpickle('cifar-10-batches-py/data_batch_'+`j+1`)
          x = d['data']
          y = d['labels']
          xs.append(x)
          ys.append(y)

        d = unpickle('cifar-10-batches-py/test_batch')
        xs.append(d['data'])
        ys.append(d['labels'])
    if dataset == 'CIFAR-100':
        d = unpickle('cifar-100-python/train')
        x = d['data']
        y = d['fine_labels']
        xs.append(x)
        ys.append(y)

        d = unpickle('cifar-100-python/test')
        xs.append(d['data'])
        ys.append(d['fine_labels'])

    x = np.concatenate(xs)/np.float32(255)
    y = np.concatenate(ys)
    x = np.dstack((x[:, :1024], x[:, 1024:2048], x[:, 2048:]))
    x = x.reshape((x.shape[0], 32, 32, 3)).transpose(0,3,1,2)

    # subtract per-pixel mean
    pixel_mean = np.mean(x[0:50000],axis=0)
    x -= pixel_mean

    # create mirrored images
    X_train = x[0:50000,:,:,:]
    Y_train = y[0:50000]
    X_train_flip = X_train[:,:,:,::-1]
    Y_train_flip = Y_train
    X_train = np.concatenate((X_train,X_train_flip),axis=0)
    Y_train = np.concatenate((Y_train,Y_train_flip),axis=0)

    X_test = x[50000:,:,:,:]
    Y_test = y[50000:]

    return dict(
        X_train=lasagne.utils.floatX(X_train),
        Y_train=Y_train.astype('int32'),
        X_test = lasagne.utils.floatX(X_test),
        Y_test = Y_test.astype('int32'),)

# ##################### Build the neural network model #######################



def ResNet_FullPre_Wide(input_var=None, nout=10,  n=3, k=2, dropoutrate = 0):
    '''
    Adapted from https://gist.github.com/FlorianMuellerklein/3d9ba175038a3f2e7de3794fa303f1ee
    which was tweaked to be consistent with 'Identity Mappings in Deep Residual Networks', Kaiming He et al. 2016 (https://arxiv.org/abs/1603.05027)
    And 'Wide Residual Networks', Sergey Zagoruyko, Nikos Komodakis 2016 (http://arxiv.org/pdf/1605.07146v1.pdf)
    '''
    n_filters = {0:16, 1:16*k, 2:32*k, 3:64*k}

    # create a residual learning building block with two stacked 3x3 convlayers and dropout
    def residual_block(l, increase_dim=False, first=False, filters=16):
        if increase_dim:
            first_stride = (2,2)
        else:
            first_stride = (1,1)

        if first:
            # hacky solution to keep layers correct
            bn_pre_relu = l
        else:
            # contains the BN -> ReLU portion, steps 1 to 2
            bn_pre_conv = BatchNormLayer(l)
            bn_pre_relu = NonlinearityLayer(bn_pre_conv, rectify)

        # contains the weight -> BN -> ReLU portion, steps 3 to 5
        conv_1 = batch_norm(ConvLayer(bn_pre_relu, num_filters=filters, filter_size=(3,3), stride=first_stride, nonlinearity=rectify, pad='same', W=HeNormal(gain='relu')))

        if dropoutrate > 0:   # with dropout
            dropout = DropoutLayer(conv_1, p=dropoutrate)

            # contains the last weight portion, step 6
            conv_2 = ConvLayer(dropout, num_filters=filters, filter_size=(3,3), stride=(1,1), nonlinearity=None, pad='same', W=HeNormal(gain='relu'))
        else:   # without dropout
            conv_2 = ConvLayer(conv_1, num_filters=filters, filter_size=(3,3), stride=(1,1), nonlinearity=None, pad='same', W=HeNormal(gain='relu'))

        # add shortcut connections
        if increase_dim:
            # projection shortcut, as option B in paper
            projection = ConvLayer(l, num_filters=filters, filter_size=(1,1), stride=(2,2), nonlinearity=None, pad='same', b=None)
            block = ElemwiseSumLayer([conv_2, projection])
        elif first:
            # projection shortcut, as option B in paper
            projection = ConvLayer(l, num_filters=filters, filter_size=(1,1), stride=(1,1), nonlinearity=None, pad='same', b=None)
            block = ElemwiseSumLayer([conv_2, projection])
        else:
            block = ElemwiseSumLayer([conv_2, l])

        return block

    # Building the network
    l_in = InputLayer(shape=(None, 3, 32, 32), input_var=input_var)

    # first layer=
    l = batch_norm(ConvLayer(l_in, num_filters=n_filters[0], filter_size=(3,3), stride=(1,1), nonlinearity=rectify, pad='same', W=HeNormal(gain='relu')))

    # first stack of residual blocks
    l = residual_block(l, first=True, filters=n_filters[1])
    for _ in range(1,n):
        l = residual_block(l, filters=n_filters[1])

    # second stack of residual blocks
    l = residual_block(l, increase_dim=True, filters=n_filters[2])
    for _ in range(1,n):
        l = residual_block(l, filters=n_filters[2])

    # third stack of residual blocks
    l = residual_block(l, increase_dim=True, filters=n_filters[3])
    for _ in range(1,n):
        l = residual_block(l, filters=n_filters[3])

    bn_post_conv = BatchNormLayer(l)
    bn_post_relu = NonlinearityLayer(bn_post_conv, rectify)

    # average pooling
    avg_pool = GlobalPoolLayer(bn_post_relu)

    # fully connected layer
    network = DenseLayer(avg_pool, num_units=nout, W=HeNormal(), nonlinearity=softmax)

    return network

# ############################# Batch iterator ###############################

def iterate_minibatches(inputs, targets, batchsize, shuffle=False, augment=False):
    assert len(inputs) == len(targets)
    if shuffle:
        indices = np.arange(len(inputs))
        np.random.shuffle(indices)
    for start_idx in range(0, len(inputs) - batchsize + 1, batchsize):
        if shuffle:
            excerpt = indices[start_idx:start_idx + batchsize]
        else:
            excerpt = slice(start_idx, start_idx + batchsize)
        if augment:
            # as in paper :
            # pad feature arrays with 4 pixels on each side
            # and do random cropping of 32x32
            padded = np.pad(inputs[excerpt],((0,0),(0,0),(4,4),(4,4)),mode='constant')
            random_cropped = np.zeros(inputs[excerpt].shape, dtype=np.float32)
            crops = np.random.random_integers(0,high=8,size=(batchsize,2))
            for r in range(batchsize):
                random_cropped[r,:,:,:] = padded[r,:,crops[r,0]:(crops[r,0]+32),crops[r,1]:(crops[r,1]+32)]
            inp_exc = random_cropped
        else:
            inp_exc = inputs[excerpt]

        yield inp_exc, targets[excerpt]

# ############################## Main program ################################

def main(dataset = 'CIFAR-10', iscenario = 0, n=5, k = 1, num_epochs=82, model = None, irun = 0, Te = 2.0, E1 = 41, E2 = 61, E3 = 81,
         lr = 0.1, lr_fac = 0.1, reg_fac = 0.0005, t0 = math.pi/2.0, Estart = 0, dropoutrate = 0, multFactor = 1):
    # Check if CIFAR data exists
    if dataset == 'CIFAR-10':
        if not os.path.exists("./cifar-10-batches-py"):
            print("CIFAR-10 dataset can not be found. Please download the dataset from 'https://www.cs.toronto.edu/~kriz/cifar.html'.")
            return
        nout = 10
    if dataset == 'CIFAR-100':
        if not os.path.exists("./cifar-100-python"):
            print("CIFAR-100 dataset can not be found. Please download the dataset from 'https://www.cs.toronto.edu/~kriz/cifar.html'.")
            return
        nout = 100
    # Load the dataset
    print("Loading data...")
    data = load_data(dataset)
    X_train = data['X_train']
    Y_train = data['Y_train']
    X_test = data['X_test']
    Y_test = data['Y_test']

    # Prepare Theano variables for inputs and targets
    input_var = T.tensor4('inputs')
    target_var = T.ivector('targets')

    # Create neural network model
    print("Building model and compiling functions...")
    network = ResNet_FullPre_Wide(input_var, nout,  n, k, dropoutrate)
    print("number of parameters in model: %d" % lasagne.layers.count_params(network, trainable=True))


    if model is None:
        # Create a loss expression for training, i.e., a scalar objective we want
        # to minimize (for our multi-class problem, it is the cross-entropy loss):
        prediction = lasagne.layers.get_output(network)
        loss = lasagne.objectives.categorical_crossentropy(prediction, target_var)
        loss = loss.mean()
        # add weight decay
        all_layers = lasagne.layers.get_all_layers(network)
        sh_reg_fac = theano.shared(lasagne.utils.floatX(reg_fac))
        l2_penalty = lasagne.regularization.regularize_layer_params(all_layers, lasagne.regularization.l2) * sh_reg_fac
        loss = loss + l2_penalty

        # Create update expressions for training
        # Stochastic Gradient Descent (SGD) with momentum
        params = lasagne.layers.get_all_params(network, trainable=True)
        sh_lr = theano.shared(lasagne.utils.floatX(lr))
        updates = lasagne.updates.momentum(loss, params, learning_rate=sh_lr, momentum=0.9)

        # Compile a function performing a training step on a mini-batch (by giving
        # the updates dictionary) and returning the corresponding training loss:
        train_fn = theano.function([input_var, target_var], loss, updates=updates)

    # Create a loss expression for validation/testing
    test_prediction = lasagne.layers.get_output(network, deterministic=True)
    test_loss = lasagne.objectives.categorical_crossentropy(test_prediction, target_var)

    test_loss = test_loss.mean()
    test_acc = T.mean(T.eq(T.argmax(test_prediction, axis=1), target_var), dtype=theano.config.floatX)

    # Compile a second function computing the validation loss and accuracy:
    val_fn = theano.function([input_var, target_var], [test_loss, test_acc])

    # statistics file
    filename = "stat_{}_{}.txt".format(iscenario, irun)
    myfile=open(filename, 'w+', 0)
    start_time0 = time.time()

    tt = 0
    TeNext = Te
    batchsize = 128

    if model is None:
        # launch the training loop
        print("Starting training...")
        # We iterate over epochs:
        for epoch in range(num_epochs):
            # shuffle training data
            train_indices = np.arange(100000)
            np.random.shuffle(train_indices)
            X_train = X_train[train_indices,:,:,:]
            Y_train = Y_train[train_indices]

            # In each epoch, we do a full pass over the training data:
            train_err = 0
            train_batches = 0
            start_time = time.time()

            for batch in iterate_minibatches(X_train, Y_train, batchsize, shuffle=True, augment=True):
                inputs, targets = batch
                train_err += train_fn(inputs, targets)
                train_batches += 1

                if (epoch+1 >= Estart): # time to start adjust learning tate
                    dt = 2.0*math.pi/float(2.0*Te)
                    tt = tt + float(dt)/(len(Y_train)/float(batchsize))
                    if tt >= math.pi:
                        tt = tt - math.pi
                    curT = t0 + tt
                    new_lr = lr * (1.0 + math.sin(curT))/2.0    # lr_min = 0, lr_max = lr
                    sh_lr.set_value(lasagne.utils.floatX(new_lr))

            if (epoch+1 == TeNext):     # time to restart
                tt = 0                  # by setting to 0 we set lr to lr_max, see above
                Te = Te * multFactor    # change the period of restarts
                TeNext = TeNext + Te    # note the next restart's epoch

            # And a full pass over the validation data:
            val_err = 0
            val_acc = 0
            val_batches = 0
            for batch in iterate_minibatches(X_test, Y_test, 500, shuffle=False):
                inputs, targets = batch
                err, acc = val_fn(inputs, targets)
                val_err += err
                val_acc += acc
                val_batches += 1

            # Then we print the results for this epoch:
            print("Epoch {} of {} took {:.3f}s".format(epoch + 1, num_epochs, time.time() - start_time))
            print("  training loss:\t\t{:.6f}".format(train_err / train_batches))
            print("  validation loss:\t\t{:.6f}".format(val_err / val_batches))
            print("  validation accuracy:\t\t{:.2f} %".format(val_acc / val_batches * 100))

            # print some statistics
            myfile.write("{}\t{:.15g}\t{:.15g}\t{:.15g}\t{:.15g}\t{:.15g}\t{:.15g}\n".format(epoch, float(sh_lr.get_value()),
                    time.time() - start_time0, Te, train_err / train_batches, val_err / val_batches, val_acc / val_batches * 100))

            # dump the network weights to a file :
            if epoch % 10 == 0:
                filesave = "network_{}_{}_{}.npz".format(iscenario,irun,epoch)
                np.savez(filesave, *lasagne.layers.get_all_param_values(network))

            # adjust learning rate as in the original approach
            if (epoch+1) == E1 or (epoch+1) == E2 or (epoch+1) == E3:
                new_lr = sh_lr.get_value() * lr_fac
                print("New LR:"+str(new_lr))
                sh_lr.set_value(lasagne.utils.floatX(new_lr))

    else:
        # load network weights from model file
        with np.load(model) as f:
             param_values = [f['arr_%d' % i] for i in range(len(f.files))]
        lasagne.layers.set_all_param_values(network, param_values)

    myfile.close()

    # Calculate validation error of model:
    test_err = 0
    test_acc = 0
    test_batches = 0
    for batch in iterate_minibatches(X_test, Y_test, 500, shuffle=False):
        inputs, targets = batch
        err, acc = val_fn(inputs, targets)
        test_err += err
        test_acc += acc
        test_batches += 1
    print("Final results:")
    print("  test loss:\t\t\t{:.6f}".format(test_err / test_batches))
    print("  test accuracy:\t\t{:.2f} %".format(
        test_acc / test_batches * 100))


if __name__ == '__main__':

    # the only input is 'iscenario' index used to reproduce the experiments given in the paper
    # scenario #1 and #2 correspond to the original multi-step learning rate decay on CIFAR-10
    # scenarios [3-6] are 4 options for our SGDR
    # scenarios [7-10] are the same options but for 2 times wider WRNs, i.e., WRN-28-20
    # scenarios [11-20] are the same as [1-10] but for CIFAR-100
    iscenario = int(sys.argv[1])
    model = None

    dataset = 'CIFAR-10'

    iruns = [1,2,3,4,5]
    lr = 0.05
    lr_fac = 0.2
    reg_fac = 0.0005
    t0 = math.pi/2.0
    Te = 50
    dropoutrate = 0
    multFactor = 1
    num_epochs = 200
    E1 = -1;    E2 = -1;     E3 = -1;   Estart = -1

    # CIFAR-10
    if (iscenario == 1):    dataset = 'CIFAR-10';     n = 4;  k = 10;  E1 = 60;    E2 = 120;    E3 = 160;  Estart = 10000;   lr = 0.1;
    if (iscenario == 2):    dataset = 'CIFAR-10';     n = 4;  k = 10;  E1 = 60;    E2 = 120;    E3 = 160;  Estart = 10000;   lr = 0.05;
    if (iscenario == 3):    dataset = 'CIFAR-10';     n = 4;  k = 10;  Te = 50;
    if (iscenario == 4):    dataset = 'CIFAR-10';     n = 4;  k = 10;  Te = 100;
    if (iscenario == 5):    dataset = 'CIFAR-10';     n = 4;  k = 10;  Te = 1;      multFactor = 2;
    if (iscenario == 6):    dataset = 'CIFAR-10';     n = 4;  k = 10;  Te = 10;     multFactor = 2;
    if (iscenario == 7):    dataset = 'CIFAR-10';     n = 4;  k = 20;  Te = 50;
    if (iscenario == 8):    dataset = 'CIFAR-10';     n = 4;  k = 20;  Te = 100;
    if (iscenario == 9):    dataset = 'CIFAR-10';     n = 4;  k = 20;  Te = 1;      multFactor = 2;
    if (iscenario == 10):   dataset = 'CIFAR-10';     n = 4;  k = 20;  Te = 10;     multFactor = 2;

    # the same for CIFAR-100
    if (iscenario == 11):   dataset = 'CIFAR-100';    n = 4;  k = 10;  E1 = 60;    E2 = 120;    E3 = 160;  Estart = 10000;   lr = 0.1;
    if (iscenario == 12):   dataset = 'CIFAR-100';    n = 4;  k = 10;  E1 = 60;    E2 = 120;    E3 = 160;  Estart = 10000;   lr = 0.05;
    if (iscenario == 13):   dataset = 'CIFAR-100';    n = 4;  k = 10;  Te = 50;
    if (iscenario == 14):   dataset = 'CIFAR-100';    n = 4;  k = 10;  Te = 100;
    if (iscenario == 15):   dataset = 'CIFAR-100';    n = 4;  k = 10;  Te = 1;      multFactor = 2;
    if (iscenario == 16):   dataset = 'CIFAR-100';    n = 4;  k = 10;  Te = 10;     multFactor = 2;
    if (iscenario == 17):   dataset = 'CIFAR-100';    n = 4;  k = 20;  Te = 50;
    if (iscenario == 18):   dataset = 'CIFAR-100';    n = 4;  k = 20;  Te = 100;
    if (iscenario == 19):   dataset = 'CIFAR-100';    n = 4;  k = 20;  Te = 1;      multFactor = 2;
    if (iscenario == 20):   dataset = 'CIFAR-100';    n = 4;  k = 20;  Te = 10;     multFactor = 2;

    # very wide nets on CIFAR-10 and CIFAR-100
    if (iscenario == 21):   dataset = 'CIFAR-10';     n = 4;  k = 20;  E1 = 60;    E2 = 120;    E3 = 160;  Estart = 10000;   lr = 0.1;
    if (iscenario == 22):   dataset = 'CIFAR-10';     n = 4;  k = 20;  E1 = 60;    E2 = 120;    E3 = 160;  Estart = 10000;   lr = 0.05;
    if (iscenario == 23):   dataset = 'CIFAR-100';    n = 4;  k = 20;  E1 = 60;    E2 = 120;    E3 = 160;  Estart = 10000;   lr = 0.1;
    if (iscenario == 24):   dataset = 'CIFAR-100';    n = 4;  k = 20;  E1 = 60;    E2 = 120;    E3 = 160;  Estart = 10000;   lr = 0.05;
    if (iscenario == 25):   dataset = 'CIFAR-10';     n = 4;  k = 20;  E1 = 60;    E2 = 120;    E3 = 160;  Estart = 10000;   lr = 0.1;
    if (iscenario == 26):   dataset = 'CIFAR-10';     n = 4;  k = 20;  E1 = 60;    E2 = 120;    E3 = 160;  Estart = 10000;   lr = 0.05;
    if (iscenario == 27):   dataset = 'CIFAR-100';    n = 4;  k = 20;  E1 = 60;    E2 = 120;    E3 = 160;  Estart = 10000;   lr = 0.1;
    if (iscenario == 28):   dataset = 'CIFAR-100';    n = 4;  k = 20;  E1 = 60;    E2 = 120;    E3 = 160;  Estart = 10000;   lr = 0.05;

    for irun in iruns:
        main(dataset, iscenario, n, k, num_epochs, model, irun, Te, E1, E2, E3, lr, lr_fac, reg_fac, t0, Estart, dropoutrate, multFactor)