# This file shows how to implement a single hidden layer neural network for
# performing binary classification on the GPU using cudamat.

from __future__ import division
import pdb
import time
import numpy as np
import cudamat as cm
from cudamat import learn as cl
import util

# initialize CUDA
cm.cublas_init()

# load data
util.load('mnist49.dat', globals())

# Put training data onto the GPU.
dat_train = dat_train/255.
dat_train = dat_train - (np.mean(dat_train, 1)+10**-8)[:, np.newaxis]
dev_train = cm.CUDAMatrix(dat_train)
dev_lbl = cm.CUDAMatrix(lbl_train)

# training parameters
epsilon = 0.01
momentum = 0.9

num_epochs = 30
batch_size = 128
num_batches = dat_train.shape[1]//batch_size

# model parameters
dim_in = dat_train.shape[0]
dim_out = 1
num_hid = 1024

# initialize weights
w_w1 = cm.CUDAMatrix(dim_in ** -0.5 * np.random.randn(dim_in, num_hid))
w_b1 = cm.CUDAMatrix(np.zeros((num_hid, 1)))
w_w2 = cm.CUDAMatrix(num_hid ** -0.5 * np.random.randn(num_hid, dim_out))
w_b2 = cm.CUDAMatrix(np.zeros((dim_out, 1)))

# initialize weight update matrices
wu_w1 = cm.empty(w_w1.shape).assign(0)
wu_b1 = cm.empty(w_b1.shape).assign(0)
wu_w2 = cm.empty(w_w2.shape).assign(0)
wu_b2 = cm.empty(w_b2.shape).assign(0)

# initialize temporary storage
h = cm.empty((num_hid, batch_size))
out = cm.empty((dim_out, batch_size))
delta = cm.empty((num_hid, batch_size))

# Train neural network.
start_time = time.time()
for epoch in range(num_epochs):
    print("Epoch %i" % (epoch + 1))
    err = []

    for batch in range(num_batches):
        # get current minibatch
        inp = dev_train.slice(batch*batch_size,(batch + 1)*batch_size)
        target = dev_lbl.slice(batch*batch_size,(batch + 1)*batch_size)

        # forward pass
        cm.dot(w_w1.T, inp, target = h)

        h.add_col_vec(w_b1)
        h.apply_sigmoid()

        cm.dot(w_w2.T, h, target = out)

        out.add_col_vec(w_b2)
        out.apply_sigmoid()

        # back prop errors
        out.subtract(target) # compute error

        # gradients for w_w2 and w_b2
        wu_w2.add_dot(h, out.T, beta = momentum)
        wu_b2.add_sums(out, axis = 1, beta = momentum)

        # compute delta
        cm.dot(w_w2, out, target = delta)

        # delta = delta * h * (1 - h)
        cl.mult_by_sigmoid_deriv(delta, h)

        # gradients for w_w1 and w_b1
        wu_w1.add_dot(inp, delta.T, beta = momentum)
        wu_b1.add_sums(delta, axis = 1, beta = momentum)

        # update weights
        w_w1.subtract_mult(wu_w1, epsilon/batch_size)
        w_b1.subtract_mult(wu_b1, epsilon/batch_size)
        w_w2.subtract_mult(wu_w2, epsilon/batch_size)
        w_b2.subtract_mult(wu_b2, epsilon/batch_size)

        # calculate error on current minibatch 
        err.append(np.abs(out.asarray())>0.5)

    print("Training misclassification rate: %f" % np.mean(err))
    print("Time: %f" % (time.time() - start_time))

# Evaluate neural network on test data.

# Load test data onto the GPU.
dat_test = dat_test/255.
dat_test = dat_test - np.mean(dat_test, 1)[:, np.newaxis]
dev_test = cm.CUDAMatrix(dat_test)
dev_lbl = cm.CUDAMatrix(lbl_test)

# Initalize temporary storage.
h = cm.empty((num_hid, dat_test.shape[1]))
out = cm.empty((dim_out, dat_test.shape[1]))

# forward pass
cm.dot(w_w1.T, dev_test, target = h)

h.add_col_vec(w_b1)
h.apply_sigmoid()

cm.dot(w_w2.T, h, target = out)

out.add_col_vec(w_b2)
out.apply_sigmoid()

# compute error
out.subtract(dev_lbl)

print("Testing misclassification rate: %f" % np.mean(np.abs(out.asarray())>0.5))

cm.cublas_shutdown()

[(13, 'cudamat.cublas_init', 'cm.cublas_init', True, 'import cudamat as cm\n'), (21, 'cudamat.CUDAMatrix', 'cm.CUDAMatrix', True, 'import cudamat as cm\n'), (22, 'cudamat.CUDAMatrix', 'cm.CUDAMatrix', True, 'import cudamat as cm\n'), (50, 'cudamat.empty', 'cm.empty', True, 'import cudamat as cm\n'), (51, 'cudamat.empty', 'cm.empty', True, 'import cudamat as cm\n'), (52, 'cudamat.empty', 'cm.empty', True, 'import cudamat as cm\n'), (55, 'time.time', 'time.time', False, 'import time\n'), (110, 'cudamat.CUDAMatrix', 'cm.CUDAMatrix', True, 'import cudamat as cm\n'), (111, 'cudamat.CUDAMatrix', 'cm.CUDAMatrix', True, 'import cudamat as cm\n'), (114, 'cudamat.empty', 'cm.empty', True, 'import cudamat as cm\n'), (115, 'cudamat.empty', 'cm.empty', True, 'import cudamat as cm\n'), (118, 'cudamat.dot', 'cm.dot', True, 'import cudamat as cm\n'), (123, 'cudamat.dot', 'cm.dot', True, 'import cudamat as cm\n'), (133, 'cudamat.cublas_shutdown', 'cm.cublas_shutdown', True, 'import cudamat as cm\n'), (39, 'numpy.zeros', 'np.zeros', True, 'import numpy as np\n'), (41, 'numpy.zeros', 'np.zeros', True, 'import numpy as np\n'), (38, 'numpy.random.randn', 'np.random.randn', True, 'import numpy as np\n'), (40, 'numpy.random.randn', 'np.random.randn', True, 'import numpy as np\n'), (44, 'cudamat.empty', 'cm.empty', True, 'import cudamat as cm\n'), (45, 'cudamat.empty', 'cm.empty', True, 'import cudamat as cm\n'), (46, 'cudamat.empty', 'cm.empty', True, 'import cudamat as cm\n'), (47, 'cudamat.empty', 'cm.empty', True, 'import cudamat as cm\n'), (66, 'cudamat.dot', 'cm.dot', True, 'import cudamat as cm\n'), (71, 'cudamat.dot', 'cm.dot', True, 'import cudamat as cm\n'), (84, 'cudamat.dot', 'cm.dot', True, 'import cudamat as cm\n'), (87, 'cudamat.learn.mult_by_sigmoid_deriv', 'cl.mult_by_sigmoid_deriv', True, 'from cudamat import learn as cl\n'), (109, 'numpy.mean', 'np.mean', True, 'import numpy as np\n'), (20, 'numpy.mean', 'np.mean', True, 'import numpy as np\n'), (102, 'numpy.mean', 'np.mean', True, 'import numpy as np\n'), (103, 'time.time', 'time.time', False, 'import time\n')]