mnist_2d_routenet_1to1_output_banks.py

# mnist_2d_routenet_1to1_output_banks.py

from __future__ import print_function
import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from torch.autograd import Variable
import numpy as np
import time
import sys
import ConfigParser
import matplotlib.pyplot as plt
from sklearn.metrics import confusion_matrix
import scipy.io as sio

import routenet as rn
import pdb
import pickle

plt.ion()

#############################################################################
# COMMENTS:
#
# 5/21/18
#
# Gates are applied to connections between banks of source neurons and banks
# of target neurons. The notion is that a bank of source neurons is capable of
# sending its outputs to multiple destinations, but the outputs themselves
# dictate where the information needs to go. That is, neural data are
# dynamically routed in a local manner. The gates are applied to the outputs
# of the  source neurons using hard sigmoid functions. A target bank may be
# receive all or none of the outputs from its source banks.
#
# Network connectivity architecture is relevant. If banks have full
# connectivity there is no reason for the network to learn meaningful routing
# paths. In contrast, if the network has partial connectivity (e.g., layered
# banks, but with banks having limited fan-in and fan-out) the network may
# learn distint paths for distinct types of data that share features, allowing
# for network re-use while only activiting parts of the network useful for
# processing a certain type of data (e.g., faces versus places). This may only
# be relevant in a network in which network outputs are spatially distinct as
# well, as in brain. E.g., neural outputs relevant for reaching ("where"
# information) need to ultimately route to motor circuits while outputs
# relevant to classification and memory ("what" information) need to route to
# hippocampal circuits.
#
# The cost function includes a component equal to the sum/mean of the
# activations (hard sigmoid outputs) of all the gates. This can promote
# sparsity but more importantly promotes self-routing. Additionally, if all
# inputs to a target bank are gated off, then that target bank effectively
# sends no outputs.
#
# Hypothetically this can lead to efficient processing of data because
# activations of the banks only have to be computed for a subset of the total
# number of banks. One can maintain a list/table of which banks are active,
# and all sources of a target bank are inactive (gated off), the activations
# of the target bank do not need to be computed/updated. In a recurrent
# network, asynchronous and parallel computation might be effective--that is,
# the activation/gate table could be shared amongst different processers,
# which iteratively update activations of a subset of the banks. And/or, one
# could use a random updating schedule in  which a random target bank is
# chosen and updated. Finally, this may also be a method of adaptively
# selecting the size of the network. During training, banks that are never or
# rarely activated could be removed from the network altogether.
#
# In comparison to Hinton's capsules ("routing by agreement") I believe a key
# difference is that in capsule networks, some information is always routed
# forward (need to confirm), whereas that is not the case here. This might
# be compared to missed perceptions. E.g., some phonetic sounds that are not
# behaviorally relevant during childhood development are simply "not heard"
# (perceived) in adulthood.
#
# 
# Aside from efficiency, what else might these networks be good at?
#
# - They might be resistant to Adversarial Examples, in that all outputs
#   may be gated off. This could be especially true if networks are
#   trained with probabilistic gating.
# - In a recurrent network, top-down information may gate bottom-up information.
#
#
# How might these networks be built upon?
#
# - Gating may be dictated not solely by data from a source bank to a target
#   bank, but by multiple source banks to a target bank (e.g., winner-take-
#   all). "Collective gating" rather than "neighbor gating."
# - In a network with iterative processing, feed-forward data could be
#   gated by both feedfoward and feedback data.
# - Both of the above, together.
# - Recurrent networks in which gating is independently applied to feedforward
#   data and feedback data.
# - Not pertinent only to routing networks: Asynchronous updates with gating,
#   in which auto-feedback from a bank (a memory, e.g., activity decay of 0.9)
#   is used to allow for only a subset of inputs being used to update the
#   output in each iteration. The auto-feedback "holds" that activity that
#   was generated by the previously processed inputs.
# - Banks could be arranged with patterned connectivity (e.g., hexagons) and
#   furthermore, connected to a central structure (e.g., thalamus) that
#   accomodates gating between more distant areas. Probably need a different
#   learning rule for the two types of gating.
#
#
# What about regularization?
#
# - Gating itself may provide a form of regularization.
# - This regularization might be stronger if gates are implemented
#   probabilistically.
# - What about dropout? It may be hard to train models in which inputs to 
#   gates
# 
#############################################################################

#############################################################################
# FOR ARXIV PAPER, POSSIBLE THINGS TO DEMONSTRATE:
# 
# 1. Effiency: Number of operations needed on gated versus ungated model
#    applied to generic MNIST. First build minimal model without gates (set
#    gates to gain of 1.0 and bias of 0 to simulate no gates). Strive for
#    95% test accuracy. Then train with gates using same architecture.
#    What is efficiency gain?
#
# 2. Train model with more banks in the final bank layer than the number
#    of output nodes. Show that gates that feed the unconnected final banks
#    are almost always closed even for the test set.
#    Should show "functional connectivity map" for both training and test sets.
#
# 3. Incremental learning without forgetting?
#    Train on one digit at a time, using samples of all digit types, but 
#    a cost function that is binary on the digit under training. Is
#    forgetting reduced compared to an ungated model under the same
#    training regimen?
#
# 4. Spatial routing
#    Use MNIST digits in larger field. Parcellate inputs into input banks
#    in 2D grid-wise manner. Maintain a 2D spatial arrangement of banks in layers
#    but take outputs only from a single row or column. After training,
#    show "flow" of digit info from its input location to the proper
#    output bank/node. If possible, also have model predict digit location,
#    using output neurons that are spatially distant from the classification
#    neurons. This is analogous to what/where pathways in the brain, with
#    those different types of information being used for different purposes.
#    In a hypothetical model with feedback gating, that task at had could
#    gate off pathways the provide unnecessary information.
#
# 5. Spatial routing and counting
#    Use inputs and architecture similar to in (4) but use multiple instance
#    of two digit classes in the input images, randomly positioned. Outputs
#    are two nodes that give the count of the number of instances of each
#    class. Does the model route in a manner that can accomodate this 
#    parallelized counting?
#
# 6. Different routing paths for different information?
#    Use standard MNIST but with colored digits (on background of colored
#    noise?). Use 2D spatial arrangement. Output neurons for digit classification
#    should be spatially isolated from those That output color information.
#    Show that the two paths are both activated. But what's the point?
#    It's only half an experiment in which feedback gating might control
#    what information is needed/wanted, and thus gate-off routes that are
#    unnecessary.
#############################################################################


# TODO: Use FF net with limited connectivity. Train/test with MNIST
#       expanded to larger area, with digit in random location. Network
#       has three outputs: Class, x-location, y-location. Do we see
#       divergence of information as neurons get closer to the output?
#       What if a new class is used in testing?  Does location information
#       get routed, but ID information gets gated off before hitting the
#       classification layer/output?
# Related neuroscience experiment:
#       Put two objects in image. Do two classifications and two locations
#       emerge at network outputs? Is there neural activity that can
#       "connects" each object to it's location? Is there ever a confusion
#       (e.g. swap) of locations?  If noise is added to the neurons?

# TODO: Train/test on CIFAR, and mixed CIFAR-MNIST
# TODO: On mixed CIFAR-MNIST, do we see divergence of routing paths?
# TODO: What fraction of banks are never gated and can thus be removed? Method for adapting network size to fit the data?

# TODO: Does trained hardgate network ever give None as output?
# TODO: Not in near term, but test random ordering of gate and bank updates. Converges to solution?

# TODO: Probabilistic gating.
# TODO: Allow for lambda scheduling: Start with low weight on gating
#       activation and increase with each epoch.
# TODO: Add activation loss. All activations or just gates?

# IDEA: Add loss based on distance between neurons (wiring cost). Helpful in neuromimetic processor (keep processing local)?
# IDEA: Hierarchical routing?  Fractal/hierarchical connectivity patterns/modularity?
# IDEA: Accompanying mechanisms to modulate learning?

data_set = 'random_location_mnist'  # 'mnist', or 'cifar10', or 'random_location_mnist'

# Read in path where raw and processed data are stored
configParser = ConfigParser.RawConfigParser()
configParser.readfp(open(r'config.txt'))

data_section = 'Data Directories'
dirMnistData = configParser.get(data_section, 'dirMnistData')
dirCifar10Data = configParser.get(data_section, 'dirCifar10Data')
fullRootFilenameSoftModel = configParser.get(data_section, 'fullRootFilenameSoftModel')
fullRootFilenameHardModel = configParser.get(data_section, 'fullRootFilenameHardModel')
fullRootFilenameSoftModelLastEpoch = configParser.get(data_section, 'fullRootFilenameSoftModelLastEpoch')

## Get training settings
parser = argparse.ArgumentParser(description='PyTorch MNIST Example')
parser.add_argument('--batch-size', type=int, default=100, metavar='N',
                    help='input batch size for training (default: 100)')
parser.add_argument('--test-batch-size', type=int, default=100, metavar='N',
                    help='input batch size for testing (default: 1000)')
parser.add_argument('--epochs', type=int, default=10, metavar='N',
                    help='number of epochs to train (default: 10)')
parser.add_argument('--lr', type=float, default=0.01, metavar='LR',
                    help='learning rate (default: 0.01)')
parser.add_argument('--momentum', type=float, default=0.5, metavar='M',
                    help='SGD momentum (default: 0.5)')
parser.add_argument('--no-cuda', action='store_true', default=False,
                    help='disables CUDA training')
parser.add_argument('--seed', type=int, default=0, metavar='S',
                    help='random seed (default: 1)')
parser.add_argument('--log-interval', type=int, default=100, metavar='N',
                    help='how many batches to wait before logging training status')
parser.add_argument('--lambda-nll', type=float, default=1.0, metavar='N',
                    help='weighting on nll loss. weight on gate activation loss is 1-lambda_nll.')
parser.add_argument('--load', action='store_true', default=False,
                    help='load stored model')
parser.add_argument('--no-save', action='store_true', default=False,
                    help='disables saving of model during training')
parser.add_argument('--always-save', action='store_true', default=False,
                    help='saves the model after each epoch (overwrites) during training')
parser.add_argument('--no-gates', action='store_true', default=False,
                    help='trains model with all gates fully open')
parser.add_argument('--neg-gate-loss', action='store_true', default=False,
                    help='trains model with negative gate loss, promoting open gates')
args = parser.parse_args()

args.cuda = not args.no_cuda and torch.cuda.is_available()
if args.cuda:
    print('\nUsing CUDA.\n')
else:
    print('\nNot using CUDA.\n')

lambda_nll = args.lambda_nll
lambda_gate = 1 - args.lambda_nll

torch.manual_seed(args.seed)
if args.cuda:
    torch.cuda.manual_seed(args.seed)


## Define training and testing functions
def train_hardgate(epoch):
    model.train()
    loss_sum = 0.0
    prob_open_gate_sum = 0.0
    cnt = 0
    loss_nll_train_hist = np.asarray([])
    loss_gate_train_hist = np.asarray([])
    loss_total_train_hist = np.asarray([])
    t_start = time.time()
    for batch_idx, (data, target) in enumerate(train_loader):
        if args.cuda:
            data, target = data.cuda(), target.cuda()
        data, target = Variable(data), Variable(target)
        optimizer.zero_grad()
        output, total_gate_act, prob_open_gate, _ = model.forward_hardgate(data)

        if output is not None:
            # loss_nll = F.nll_loss(output, target) # Use if log_softmax *is* applied in model's forward method
            loss_nll = F.cross_entropy(output, target) # Use if log_softmax *is not* applied in model's forward method
            loss_gate = torch.mean(total_gate_act)
            loss = lambda_nll*loss_nll + lambda_gate*loss_gate

            loss.backward()
            optimizer.step()
            # TODO?: Currently have to use batch size of one. Could we accumulate
            # gradients over a number of samples and then update weights, without
            # using the optimizer? Can't use usual pytorch approach because
            # the constructed graph can be difference for each sample, during
            # training.

            loss_sum = loss_sum + loss.data.cpu().numpy()[0]
            prob_open_gate_sum += prob_open_gate
            cnt += 1

            loss_nll_train_hist = np.append(loss_nll_train_hist, loss_nll.data.cpu().numpy()[0])
            loss_gate_train_hist = np.append(loss_gate_train_hist, loss_gate.data.cpu().numpy()[0])
            loss_total_train_hist = np.append(loss_total_train_hist, loss.data.cpu().numpy()[0])

        if batch_idx % args.log_interval == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%), {:d} valid]\tLoss: {:.6f}\tProb open gate: {:.6f}\t{:.2f} seconds'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                # 100. * batch_idx / len(train_loader), loss.data[0]))
                100. * batch_idx / len(train_loader), cnt, loss_sum/cnt, prob_open_gate_sum/cnt, time.time()-t_start))
            loss_sum = 0.0
            prob_open_gate_sum = 0.0
            cnt = 0
    # return loss_sum/cnt
    return loss_total_train_hist, loss_nll_train_hist, loss_gate_train_hist

def train_softgate(epoch, weight=None):
    return_gate_status = False
    b_batch_norm = True

    model.train()

    cnt = 0
    loss_sum = 0.0
    loss_gate_sum = 0.0
    loss_nll_sum = 0.0
    prob_open_gate_sum = 0.0
    correct = 0
    pred_all = np.asarray([])
    targets_all = np.asarray([])

    if weight is None:
        weight = torch.ones(10)

    t_start = time.time()
    for batch_idx, (data, target) in enumerate(train_loader):
        target = target[:,0].contiguous()    # not using position informtion, just class label
        
        if args.cuda:
            weight = weight.cuda()
            target = target.cuda()
            for i_input_group in range(len(data)):
                data[i_input_group] = data[i_input_group].cuda()

        target = Variable(target)
        for i_input_group in range(len(data)):
                data[i_input_group] = Variable(data[i_input_group])

        optimizer.zero_grad()

        if return_gate_status:
            output, total_gate_act, prob_open_gate, gate_status_ = model.forward_softgate(data,
                                                                                          return_gate_status = return_gate_status,
                                                                                          b_batch_norm = b_batch_norm,
                                                                                          b_use_cuda = args.cuda,
                                                                                          b_no_gates = args.no_gates,
                                                                                          b_neg_gate_loss = args.neg_gate_loss
                                                                                          )
        else:
            output, total_gate_act, prob_open_gate = model.forward_softgate(data,
                                                            return_gate_status = return_gate_status,
                                                            b_batch_norm = b_batch_norm,
                                                            b_use_cuda = args.cuda,
                                                            b_no_gates = args.no_gates,
                                                            b_neg_gate_loss = args.neg_gate_loss)

        loss_nll = F.cross_entropy(output, target, weight=weight)
        loss_gate = torch.mean(total_gate_act)
        loss = lambda_nll*loss_nll + lambda_gate*loss_gate
        # loss = torch.mean(total_gate_act)

        loss.backward()
        optimizer.step()
        # model.bias_limit(None, 0)   # Don't allow positive biases on hidden or output banks/nodes

        loss_sum += rn.item(loss)
        loss_gate_sum += rn.item(loss_gate)
        loss_nll_sum += rn.item(loss_nll)
        # loss_gate_sum += rn.item(loss)
        prob_open_gate_sum += prob_open_gate

        # Compute accuracy and accumulate
        pred = output.data.max(1, keepdim=True)[1] # get the index of the max log-probability
        correct += pred.eq(target.data.view_as(pred)).long().cpu().sum()
        pred_all = np.append(pred_all, pred[:,0], axis=0)
        targets_all = np.append(targets_all, target.data.cpu().numpy(), axis=0)

        cnt += 1

        if (batch_idx+1)*args.batch_size % 10000 == 0:
            cm = confusion_matrix(targets_all, pred_all)
            weight = torch.FloatTensor( np.sum(np.diag(cm)) - np.diag(cm) )**2
            weight = weight/sum(weight)
            pred_all = np.asarray([])
            targets_all = np.asarray([])

        if batch_idx % args.log_interval == 0:
            acc = (100. * float(correct)) / (cnt*args.batch_size)
            print('Train Epoch: {} [{:05d}/{} ({:.0f}%), Loss: {:.6f}\tGate loss: {:.4f}\tProb open gate: {:.4f}\tAcc: {:.2f}\t{:.2f} seconds'.format(
                epoch, (batch_idx+1)*len(data[0]), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), loss_sum/cnt, loss_gate_sum/cnt, prob_open_gate_sum/cnt, acc, time.time()-t_start))
            cnt = 0
            loss_sum = 0.0
            loss_gate_sum = 0.0
            prob_open_gate_sum = 0.0
            correct = 0.0


    acc = (100. * float(correct)) / (cnt*args.batch_size)
    print('Train Epoch: {} [{}/{} ({:.0f}%), Loss: {:.6f}\tGate loss: {:.4f}\tProb open gate: {:.4f}\tAcc: {:.2f}\t{:.2f} seconds'.format(
        epoch, (batch_idx+1)*len(data[0]), len(train_loader.dataset),
        100. * batch_idx / len(train_loader), loss_sum/cnt, loss_gate_sum/cnt, prob_open_gate_sum/cnt, acc, time.time()-t_start))
    return loss_sum/cnt, loss_nll_sum/cnt, loss_gate_sum/cnt, prob_open_gate/cnt, acc

def test_hardgate_speed():
    # Just get NN output. No other processing. To assess speed.
    model.eval()
    t_start = time.time()
    for data, target in test_loader:
        if args.cuda:
            data, target = data.cuda(), target.cuda()
        data, target = Variable(data, volatile=True), Variable(target)
        output, _, = model.forward_hardgate(data, return_gate_status=False)
    return time.time() - t_start

def test_hardgate():
    model.eval()
    test_loss_nll = 0
    test_loss_gate = 0
    test_prob_open_gate = 0
    correct = 0
    cnt = 0
    t_start = time.time()
    for data, target in test_loader:
        if args.cuda:
            data, target = data.cuda(), target.cuda()
        data, target = Variable(data, volatile=True), Variable(target)
        output, total_gate_act, prob_open_gate, gate_status = model.forward_hardgate(data)

        # Store target labels and gate status for all samples
        if cnt==0:
            gates_all = gate_status
            targets_all = target.data.cpu().numpy()
        else:
            gates_all = np.append(gates_all, gate_status, axis=0)
            targets_all = np.append(targets_all, target.data.cpu().numpy(), axis=0)

        if output is not None:
            # Accumulate losses, etc.
            # test_loss_nll += F.nll_loss(output, target).data[0] # sum up batch loss
            test_loss_nll += F.cross_entropy(output, target).data[0] # sum up batch loss
            test_loss_gate += torch.mean(total_gate_act).data[0]
            test_prob_open_gate += prob_open_gate

            # Compute accuracy and accumulate
            pred = output.data.max(1, keepdim=True)[1] # get the index of the max output
            correct += pred.eq(target.data.view_as(pred)).long().cpu().sum()

            if cnt%1000 == 0:
                print(cnt)
            cnt += 1

    test_loss_nll /= len(test_loader.dataset)
    test_loss_gate /= len(test_loader.dataset)
    test_prob_open_gate /= cnt
    test_loss = lambda_nll*test_loss_nll + lambda_gate*test_loss_gate
    acc = 100. * float(correct) / len(test_loader.dataset)
    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%), Duration: {:0.4f} seconds\n'.format(
        test_loss, correct, len(test_loader.dataset),
        100. * correct / len(test_loader.dataset), time.time()-t_start))
    # return test_loss_nll, test_loss_act
    return test_loss, test_loss_nll, test_loss_gate, test_prob_open_gate, acc, gates_all, targets_all

def test_softgate_speed():
    # Just get NN output. No other processing. To assess speed.
    model.eval()
    t_start = time.time()
    for data, target in test_loader:
        if args.cuda:
            data, target = data.cuda(), target.cuda()
        data, target = Variable(data, volatile=True), Variable(target)
        output, _, = model.forward_softgate(data, return_gate_status=False)
    return time.time() - t_start

def test_softgate():
    # TODO: Match to train_softgate. Don't get/report gate_status/prob in train, but do it in test.
    model.eval()

    return_gate_status = True
    b_batch_norm = True

    test_loss_nll = 0
    test_loss_gate = 0
    test_prob_open_gate = 0
    correct = 0
    cnt_batches = 0
    cnt_samples = 0
    pred_all = np.asarray([])
    t_start = time.time()
    for data, target in test_loader:
        target = target[:,0].contiguous()    # not using position informtion, just class label

        if args.cuda:
            target = target.cuda()
            for i_input_group in range(len(data)):
                data[i_input_group] = data[i_input_group].cuda()

        target = Variable(target)
        for i_input_group in range(len(data)):
            # data[i_input_group] = Variable(data[i_input_group], volatile=True)
            data[i_input_group] = Variable(data[i_input_group])

        # DEBUGGING: Convert everything to number and save
        '''
        targ = target.data.cpu().numpy()
        im = []
        for d in data:
            im.append(d.data.cpu().numpy())
        '''
        
        output, total_gate_act, prob_open_gate, gate_status = model.forward_softgate(data,
        #output, total_gate_act, prob_open_gate, gate_status, gate_status_value = model.forward_softgate(data,
                                                                                     return_gate_status = return_gate_status,
                                                                                     b_batch_norm = b_batch_norm,
                                                                                     b_use_cuda = args.cuda,
                                                                                     b_no_gates = args.no_gates,
                                                                                     b_neg_gate_loss = args.neg_gate_loss)
        # DEBUGGING: Convert everything to number and save
        '''
        out = output.data.cpu().numpy()
        #gate_data_dict = {'im':im, 'targ':targ, 'gate_status':gate_status, 'output':out}
        gate_data_dict = {'im':im, 'targ':targ, 'gate_status':gate_status, 'gate_status_value':gate_status_value, 'output':out}
        sio.savemat('gate_data.mat',gate_data_dict)
        pdb.set_trace()
        '''
        
        # Store target labels and gate status for all samples
        if cnt_batches==0:
            gates_all = gate_status
            targets_all = target.data.cpu().numpy()
        else:
            gates_all = np.append(gates_all, gate_status, axis=0)
            targets_all = np.append(targets_all, target.data.cpu().numpy(), axis=0)

        # Accumulate losses, etc.
        # test_loss_nll += F.nll_loss(output, target).data[0] # sum up batch loss
        test_loss_nll += rn.item(F.cross_entropy(output, target))
        test_loss_gate += rn.item(torch.mean(total_gate_act))
        # test_loss_nll += F.cross_entropy(output, target).data[0]  # sum up batch loss
        # test_loss_gate += torch.mean(total_gate_act).data[0]
        test_prob_open_gate += prob_open_gate

        # Compute accuracy and accumulate
        pred = output.data.max(1, keepdim=True)[1] # get the index of the max log-probability
        correct += pred.eq(target.data.view_as(pred)).long().cpu().sum()
        pred_all = np.append(pred_all, pred[:,0], axis=0)

        cnt_batches += 1
        cnt_samples += len(target)
        
    # test_loss_nll /= len(test_loader.dataset)
    # test_loss_gate /= len(test_loader.dataset)
    test_loss_nll /= cnt_batches
    test_loss_gate /= cnt_batches
    test_prob_open_gate /= cnt_batches
    test_loss = lambda_nll*test_loss_nll + lambda_gate*test_loss_gate
    acc = 100. * float(correct) / cnt_samples
    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.2f}%), Duration: {:0.4f} seconds\n'.format(
        test_loss, correct, cnt_samples,
        100. * float(correct) / cnt_samples, time.time()-t_start))
    cm = confusion_matrix(targets_all, pred_all)
    print('Confusion Matrix:')
    print(cm)
    print('\n')
    return test_loss, test_loss_nll, test_loss_gate, test_prob_open_gate, acc, gates_all, targets_all, pred_all, cm


def test_compare():
    model.eval()
    cnt = 0
    for data, target in test_loader:
        if args.cuda:
            data, target = data.cuda(), target.cuda()
        data, target = Variable(data, volatile=True), Variable(target)
        output1, total_gate_act1, prob_open_gate1, gate_status1 = model.forward_softgate(data,
                                                                                     return_gate_status=return_gate_status,
                                                                                     b_batch_norm=b_batch_norm)
        output2, total_gate_act2, prob_open_gate2, gate_status2 = model.forward_hardgate(data,
                                                                                     return_gate_status=return_gate_status,
                                                                                     b_batch_norm=b_batch_norm)
        if output2 is None:
            if not np.all(output1==0):
                print('Mismatch.')
        elif not np.array_equal(output1.data.cpu().numpy(), output2.data.cpu().numpy()):
            print('Mismatch.')

        if cnt % 1000 == 0:
            print(cnt)
        cnt += 1


## Set up DataLoaders
kwargs = {'num_workers': 1, 'pin_memory': True} if args.cuda else {}

if data_set == 'cifar10':
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
        ])
    f_datasets = datasets.CIFAR10
    dir_dataset = dirCifar10Data
    n_input_neurons = 3 * 32 * 32
elif data_set == 'mnist':
    transform=transforms.Compose([
        transforms.ToTensor(),
        # transforms.Normalize((0.1307,), (0.3081,)),
        transforms.Normalize((0.5,), (0.5,)),
        ])
    f_datasets = datasets.MNIST
    dir_dataset = dirMnistData
    n_input_neurons = 28 * 28
    # n_input_neurons = 32 * 32
elif data_set == 'random_location_mnist':
    transform=transforms.Compose([
        transforms.ToTensor(),
        # transforms.Normalize((0.1307,), (0.3081,)),
        #transforms.Normalize((0.5,), (0.5,)),
        ])
    f_datasets = rn.RandomLocationMNIST
    dir_dataset = dirMnistData
    # n_input_neurons = 28 * 28
    expanded_size = 112
    group_size_per_dim = 16
    #expanded_size = 28
    #group_size_per_dim = 28/7
    n_neurons_per_input_group = group_size_per_dim**2
    kwargs_datasets = {'expanded_size':expanded_size, 'group_size_1D':group_size_per_dim}
else:
    print('Unknown dataset.')
    sys.exit()

kwargs = {'num_workers': 1, 'pin_memory': True} if args.cuda else {}

train_loader = torch.utils.data.DataLoader(
                    f_datasets(dir_dataset,
                                train = True,
                                download = False,
                                transform = transform,
                                **kwargs_datasets
                                ),
                    batch_size = args.batch_size,
                    shuffle = False,
                    **kwargs)
test_loader = torch.utils.data.DataLoader(
                    f_datasets(dir_dataset,
                                train = False,
                                download = False,
                                transform = transform,
                                **kwargs_datasets
                                ),
                    batch_size = args.test_batch_size,
                    shuffle = False,
                    **kwargs)

sys.stdout.flush()
## Instantiate network model
# n_layers = 3
# n_banks_per_layer = 20
# n_fan_out = 5
# banks_per_layer = [n_banks_per_layer] * n_layers
# # banks_per_layer = np.asarray(banks_per_layer)
# # bank_conn = rn.make_conn_matrix_ff_full(banks_per_layer)
# # bank_conn = rn.make_conn_matrix_ff_part(n_layers, n_banks_per_layer, n_fan_out)
n_layers = 7
n_banks_per_layer_per_dim = expanded_size/group_size_per_dim
n_fan_out_per_dim = 3
bank_conn = rn.make_conn_matrix_ff_part_2d(n_layers, n_banks_per_layer_per_dim, n_fan_out_per_dim)
banks_per_layer = n_banks_per_layer_per_dim**2
param_dict = {'n_neurons_per_input_group':n_neurons_per_input_group,
             'idx_input_banks':np.arange(banks_per_layer),
             'bank_conn':bank_conn,
             'idx_output_banks':np.arange(n_layers*banks_per_layer-10, n_layers*banks_per_layer),
             'n_neurons_per_hidd_bank':128,
            }
if args.load:
    print('Loading existing model and weights from files. This can take several minutes...')
    t_start = time.time()
    model = rn.RouteNetOneToOneOutputGroupedInputs.init_from_files(fullRootFilenameSoftModel)
    print('\tDone. %d seconds.' % (time.time() - t_start))
else:
    # model = rn.RouteNet(**param_dict)
    # model = rn.RouteNetRecurrentGate(**param_dict)
    model = rn.RouteNetOneToOneOutputGroupedInputs(**param_dict)
    model.init_gate_bias(0.5)
if args.cuda:
    model.cuda()
# model.bias_limit(None, 0)   # Don't allow positive biases on hidden or output banks/nodes

#weights = []
#for param in model.parameters():
#    if param.data.cpu().numpy().size==1:
#        weights.append(param.data)
#pdb.set_trace()

optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.97)

## Train it, get results on test set, and save the model
loss_total_train = np.zeros(args.epochs)
loss_nll_train = np.zeros(args.epochs)
loss_gate_train = np.zeros(args.epochs)
prob_open_gate_train = np.zeros(args.epochs)
acc_train = np.zeros(args.epochs)

loss_total_test = np.zeros(args.epochs)
loss_nll_test = np.zeros(args.epochs)
loss_gate_test = np.zeros(args.epochs)
prob_open_gate_test = np.zeros(args.epochs)
acc_test = np.zeros(args.epochs)

t_start = time.time()
acc_best = 0
acc_best_epoch = 0


#test_softgate()
##########################################################
## Run the main training and testing loop
# weight = torch.ones(10)/10.
# model.freeze_data_params()
# model.freeze_gate_params()
for ep in range(0, args.epochs):
    scheduler.step()
    for param_group in optimizer.param_groups:
        print('LR = %f' % param_group['lr'])
    loss_total_train[ep], loss_nll_train[ep], loss_gate_train[ep], prob_open_gate_train[ep], acc_train[ep] = train_softgate(ep+1)
    t_test = time.time()
    with torch.no_grad():
        loss_total_test[ep], loss_nll_test[ep], loss_gate_test[ep], prob_open_gate_test[ep], acc_test[ep], gate_status, target, predicted, cm = test_softgate()
    print('Testing duration = %f' % (time.time() - t_test))
    # weight = torch.FloatTensor( np.sum(np.diag(cm)) - np.diag(cm))**2
    # weight = weight/sum(weight)
    # print(weight.view(1,-1))

    # Save model architecture and params, if it's the best so far on the test set
    if (acc_test[ep] > acc_best) and not args.no_save:
        acc_best_epoch = ep
        acc_best = acc_test[ep]
        print('Highest test set accuracy so far. Saving model.\n')
        model.save_model(fullRootFilenameSoftModel)

    if args.always_save:
        print('args.always_save: Saving model.\n')
        model.save_model(fullRootFilenameSoftModel)

    model.save_model(fullRootFilenameSoftModelLastEpoch)
        
dur = time.time()-t_start
print('Time = %f, %f sec/epoch' % (dur, dur/args.epochs))
sys.stdout.flush()
##########################################################
sys.exit()

## Compare soft and hard gating on test set
#test_compare()


## Start plotting results...
fn = 1  # Figure number


## Plot losses for test set
plt.figure(fn)
fn = fn + 1
plt.clf()

f_plot = plt.plot
# f_plot = plt.semilogy
h_subplots = 3
v_subplots = 2
i_subplot = 1

plt.subplot(h_subplots, v_subplots, i_subplot)
i_subplot += 1
f_plot(loss_nll_train, 'o-')
f_plot(loss_nll_test,'o-')
plt.legend(('Train','Test'))
plt.title('NLL loss')
plt.xlabel('Epoch')
plt.grid()

plt.subplot(h_subplots, v_subplots, i_subplot)
i_subplot += 1
f_plot(loss_gate_train, 'o-')
f_plot(loss_gate_test, 'o-')
plt.legend(('Train','Test'))
plt.title('Activation loss')
plt.xlabel('Epoch')
plt.grid()

plt.subplot(h_subplots, v_subplots, i_subplot)
i_subplot += 1
f_plot(loss_total_train, 'o-')
f_plot(loss_total_test, 'o-')
plt.legend(('Train','Test'))
plt.title('Total loss')
plt.xlabel('Epoch')
plt.grid()

plt.subplot(h_subplots, v_subplots, i_subplot)
i_subplot += 1
plt.plot(acc_train, 'o-')
plt.plot(acc_test, 'o-')
plt.legend(('Train','Test'))
plt.title('Classification accuracy')
plt.xlabel('Epoch')
plt.grid()

plt.subplot(h_subplots, v_subplots, i_subplot)
i_subplot += 1
plt.plot(100*prob_open_gate_train, 'o-')
plt.plot(100*prob_open_gate_test, 'o-')
plt.legend(('Train','Test'))
plt.title('Percentage of gates open')
plt.xlabel('Epoch')
plt.grid()


## Compute and print the test set confusion matrix
cm = confusion_matrix(target, predicted)
print('\nConfusion Matrix:')
print(cm)
print('\n')


## Compute and print...
# 1. Number and fraction of gates that are never opened.
ix_conn = np.where(bank_conn)
n_conn = len(ix_conn[0])
conn_gate_status = gate_status[:,ix_conn[0],ix_conn[1]]
ix_never = np.where(~np.any(conn_gate_status, axis=0))[0]
n_never = len(ix_never)
print('%d of %d gates (%0.1f%%) are never opened.' % (n_never, n_conn, 100.*n_never/n_conn))

# 2. Number and fraction of gates that are open on average, for ` single sample.
n_sample_open_avg = np.mean(np.sum(conn_gate_status, axis=1))
print('%d of %d gates (%0.1f%%) are open for individual samples, on average.' % (n_sample_open_avg, n_conn, 100.*n_sample_open_avg/n_conn))

# 3. Of gates that are not never opened, number/fraction that are open on average, for a single sample.
ix_sometimes = np.where(np.any(conn_gate_status, axis=0))[0]
n_sample_open_avg2 = np.mean(np.sum(conn_gate_status[ix_sometimes], axis=1))
print('Excluding gates that are always closed, %d of %d gates (%0.1f%%) are open for individual samples, on average.' % \
    (n_sample_open_avg2, n_conn-n_never, 100.*n_sample_open_avg2/(n_conn-n_never)))


## Plot the fraction of open gates, grouped by target labels
targets_unique = np.sort(np.unique(target))
plt.figure(fn)
fn = fn + 1
plt.clf()
for i, targ in enumerate(targets_unique):
    idx = np.where(target==targ)[0]
    mn = np.mean(gate_status[idx,:,:], axis=0)
    plt.subplot(2,5,i+1)
    plt.imshow(mn)
    plt.clim(0,1)
    plt.title(targ)


## Plot the fraction of open gates, grouped by target labels
## For 2D models...
targets_unique = np.sort(np.unique(target))
plt.figure(fn)
fn = fn + 1
plt.clf()
for i_targ, targ in enumerate(targets_unique):
    idx = np.where(target==targ)[0]
    mn = np.mean(gate_status[idx,:,:], axis=0)
    for i_src_layer in range(n_layers-1):
        plt.subplot(10, n_layers-1, i_targ*(n_layers-1)+i_src_layer+1)
        k = n_banks_per_layer_per_dim**2
        j_src = np.arange(i_src_layer*k, (i_src_layer+1)*k)
        j_targ = np.arange((i_src_layer+1)*k, (i_src_layer+2)*k)
        plt.imshow(mn[j_src, j_targ])
        plt.clim(0,1)
        plt.ylabel(targ, fontsize=12)
        plt.xticks()    # turn off tick labels
        plt.yticks()    # turn off tick labels


## Plot the fraction of test set open gates in connectivity map,
## grouped by target labels.
print('\nPlotting test set connectivity maps. This may take a minute...')
targets_unique = np.sort(np.unique(target))
plt.figure(fn)
fn = fn + 1
plt.clf()
layer_num = np.zeros((0))
node_num = np.zeros((0))
for i_layer in range(len(banks_per_layer)):
    layer_num = np.append(layer_num, np.full((banks_per_layer[i_layer]), i_layer))
    node_num = np.append(node_num, np.arange(banks_per_layer[i_layer]))
for i, targ in enumerate(targets_unique):
    idx = np.where(target==targ)[0]
    mn = np.mean(gate_status[idx,:,:], axis=0)
    plt.subplot(2,5,i+1)
    plt.scatter(layer_num+1, node_num, s=10, facecolors='none', edgecolors='k')
    for i_source in range(np.sum(banks_per_layer)):
        for i_target in range(np.sum(banks_per_layer)):
            alpha = mn[i_source, i_target]
            if alpha > 0.0:
                plt.plot((layer_num[i_source]+1, layer_num[i_target]+1), (node_num[i_source], node_num[i_target]), 'k-', alpha=alpha)
                # plt.plot((layer_num[i_source], layer_num[i_target]), (node_num[i_source], node_num[i_target]), 'k-')
    plt.title('"%d"' % (targ), fontsize=18)
    plt.xticks(np.arange(len(banks_per_layer))+1)
    plt.yticks(np.arange(0, banks_per_layer[0], 5))
    plt.xticks(fontsize=12)
    plt.yticks(fontsize=12)
    plt.scatter(len(banks_per_layer), banks_per_layer[0]-10+i, s=20, facecolors='r', edgecolors='r')
print('Done.')





sys.exit()


## Get model outputs for a test batch
model.eval()
cnt = 0
for data, target in test_loader:
    cnt += 1
    # data = torch.transpose(data, 2, 3).contiguous()
    # data = torch.zeros_like(data)
    if args.cuda:
        data, target = data.cuda(), target.cuda()
    data, target = Variable(data, volatile=True), Variable(target)
    output, total_gate_act, prob_open_gate, gate_status = model.forward_softgate(data, return_gate_status=True, b_batch_norm = True)

## See if there are any samples for which all the banks
## connected to output nodes are gated off, effectively
## meaning no output was generated.
x = gate_status[:,20:40,50:60]
b_open = np.any(x, axis=(1,2))

sys.exit()



weights2 = []
grads2 = []
for param in model2.parameters():
    weights2.append(param.data)
    grads2.append(param.grad)

## Plot the weights
weights = []
grads = []
for param in model.parameters():
    weights.append(param.data)
    grads.append(param.grad)
plt.figure(fn)
fn = fn + 1
plt.clf()
for i in range(0, len(weights)/2):
    plt.subplot(2,2,i+1)
    plt.imshow(weights[2*i], aspect='auto', interpolation='nearest')
    plt.colorbar()

## Plot activations for batch
model.eval()
test_loss_nll = 0
test_loss_act = 0
correct = 0
for data, target in test_loader:
    if args.cuda:
        data, target = data.cuda(), target.cuda()
    data, target = Variable(data, volatile=True), Variable(target)
    act_fc1, act_fc2, act_fc3, output = model(data)

    mx = torch.max(act_fc3).data[0]
    mx = max(mx, torch.max(act_fc2).data[0])
    mx = max(mx, torch.max(act_fc1).data[0])

    plt.figure(fn)
    fn = fn + 1
    n_samps_display = 30
    plt.clf()
    plt.subplot(2,2,1)
    plt.imshow(act_fc1.data.cpu().numpy()[:n_samps_display,:], aspect='auto', interpolation='nearest')
    plt.clim(0, mx)
    plt.subplot(2,2,2)
    plt.imshow(act_fc2.data.cpu().numpy()[:n_samps_display,:], aspect='auto', interpolation='nearest')
    plt.clim(0, mx)
    plt.subplot(2,2,3)
    plt.imshow(act_fc3.data.cpu().numpy()[:n_samps_display,:], aspect='auto', interpolation='nearest')
    plt.clim(0, mx)
    plt.colorbar()
    plt.subplot(2,2,4)
    plt.imshow(output.data.cpu().numpy()[:n_samps_display,:], aspect='auto', interpolation='nearest')
    # plt.clim(0, 1)

    plt.figure(fn)
    fn = fn + 1
    n_samps_display = 30
    plt.clf()
    plt.subplot(2,2,1)
    plt.imshow(act_fc1.data.cpu().numpy()[:n_samps_display,:]>0, aspect='auto', interpolation='nearest')
    plt.subplot(2,2,2)
    plt.imshow(act_fc2.data.cpu().numpy()[:n_samps_display,:]>0, aspect='auto', interpolation='nearest')
    plt.subplot(2,2,3)
    plt.imshow(act_fc3.data.cpu().numpy()[:n_samps_display,:]>0, aspect='auto', interpolation='nearest')
    plt.subplot(2,2,4)
    plt.imshow(output.data.cpu().numpy()[:n_samps_display,:], aspect='auto', interpolation='nearest')

    ## Plot activations for some individual classes: E.g., 0, 1, 2 ,3
    ## Do we see trends in the activation patterns at higher layers?
    for i_class in range(3):
        ix = np.where(target.data.cpu().numpy()==i_class)[0]
        ix = ix[0:n_samps_display]
        plt.figure(fn)
        fn = fn + 1
        plt.clf()
        plt.subplot(2,2,1)
        plt.imshow(act_fc1.data.cpu().numpy()[ix,:]>0, aspect='auto', interpolation='nearest')
        plt.title('Class #%d' % (i_class))
        plt.subplot(2,2,2)
        plt.imshow(act_fc2.data.cpu().numpy()[ix,:]>0, aspect='auto', interpolation='nearest')
        plt.subplot(2,2,3)
        plt.imshow(act_fc3.data.cpu().numpy()[ix,:]>0, aspect='auto', interpolation='nearest')
        plt.subplot(2,2,4)
        plt.imshow(output.data.cpu().numpy()[ix,:], aspect='auto', interpolation='nearest')

    sys.exit()