train.py

import numpy as np
import argparse
import sys
from vis import *
import data
from model import *
sys.setrecursionlimit(2**20)

import sys
sys.path.append('/home/zombori//k-arm')
sys.path.append('/Users/zsoltzombori/git/k-arm')
from arm import ArmLayer


parser = argparse.ArgumentParser(description="Image classifer using sparsifying arm layers embedded into convolution.")
parser.add_argument('--iteration', dest="iteration", type=int, default=6, help="Number of iterations in k-arm approximation")
parser.add_argument('--threshold', dest="threshold", type=float, default=0.02, help="Sparsity coefficient")
parser.add_argument('--reconsCoef', dest="reconsCoef", type=float, default=0, help="Reconstruction coefficient of the arm layers")
parser.add_argument('--dict', dest="dict_size", type=int, default=100, help="Size of the feature dictionary")
parser.add_argument('--epoch', dest="epoch", type=int, default=20, help="Number of epochs")
parser.add_argument('--batch', dest="batchSize", type=int, default=8, help="Batch size")
parser.add_argument('--dataset', dest="dataset", default="mnist", help="mnist/cifar10")
parser.add_argument('--testSize', dest="testSize", type=int, default=10000, help="test size")
parser.add_argument('--trainSize', dest="trainSize", type=int, default=60000, help="train size")
parser.add_argument('--modelType', dest="modelType", default="conv", help="arm/conv")

args = parser.parse_args()
kernel_size = (9,9)
pool_size = (2,2)

print "dataset: {}, modelType: {}, iteration: {}, threshold: {}, reconsCoef: {}, dict_size: {}, batch: {}, epoch: {}, trainSize:{}, testSize: {}".format(args.dataset,args.modelType,args.iteration,args.threshold,args.reconsCoef,args.dict_size,args.batchSize,args.epoch,args.trainSize, args.testSize)
(X_train, Y_train), (X_test, Y_test), datagen, test_datagen, nb_classes = data.load_data(args.dataset)
X_train = X_train[:args.trainSize]
Y_train = Y_train[:args.trainSize]
X_test = X_test[:args.testSize]
Y_test = Y_test[:args.testSize]
 
firstLayers = [Flatten()]
secondLayers = [Flatten()]
if args.modelType == "arm":
    firstLayers.append(ArmLayer(dict_size=args.dict_size, iteration=args.iteration, threshold=args.threshold, reconsCoef=args.reconsCoef, name="conv1"))
    secondLayers.append(ArmLayer(dict_size=args.dict_size, iteration=args.iteration, threshold=args.threshold, reconsCoef=args.reconsCoef, name="conv2"))
elif args.modelType == "conv":
    firstLayers.append(Dense(args.dict_size, input_shape=[np.prod(pool_size)], name="conv1"))
    secondLayers.append(Dense(args.dict_size, input_shape=[np.prod(pool_size)], name="conv2"))
else:
    raise Exception("Not Yet Implemented model type: " + modelType)

model = build_model(X_train.shape, nb_classes, args.batchSize, kernel_size, pool_size, firstLayers, secondLayers)
model.summary()

# fit the model on the batches generated by datagen.flow()
model.fit_generator(datagen.flow(X_train, Y_train, batch_size=args.batchSize, shuffle=True),
                        samples_per_epoch=X_train.shape[0],
                        nb_epoch=args.epoch,
                        validation_data=test_datagen.flow(X_test, Y_test, batch_size=args.batchSize),
                        nb_val_samples=X_test.shape[0]
                        )
score = model.evaluate(X_test, Y_test, verbose=0, batch_size=args.batchSize)
print('Test score:', score[0])
print('Test accuracy:', score[1])

for i in (1,2):
    layer = model.get_layer(name="conv{}".format(i))
    W_learned = layer.get_weights()[0]

    W_scaled = W_learned - np.min(W_learned)
    W_scaled /= np.max(W_scaled)
    W_scaled *= 255
    W_scaled = np.reshape(W_scaled,[args.dict_size, kernel_size[0], -1])
    fileName = "pictures/{}_{}_dict{}_kernel{}_acc{:.2f}_{}.png".format(args.dataset, args.modelType, args.dict_size, kernel_size[0], 100*score[1], i);
    vis(W_scaled, fileName)