spiral.py

import sys
import numpy as np
import matplotlib.pyplot as plt
import theano
import theano.tensor as T

rs = np.random.RandomState(1234)

def gen_spiral(label, dt, n_samples=100, noise=0.1):
    """
    Adopted from:
    https://github.com/tensorflow/playground/blob/master/dataset.ts
    """
    p = np.linspace(0,1,n_samples,endpoint=False)
    r = p * 5
    t = 1.75 * p * 2 * np.pi + dt 
    x = r * np.sin(t) + noise*(2*rs.rand(t.shape[0])-1)
    y = r * np.cos(t) + noise*(2*rs.rand(t.shape[0])-1)
    return np.c_[x,y]

def gen_data(n_samples_per_class=100):
    X = np.vstack([
        gen_spiral(label=0, dt=0,     n_samples=n_samples_per_class),
        gen_spiral(label=1, dt=np.pi, n_samples=n_samples_per_class)
    ]).astype(theano.config.floatX)

    Y = np.hstack([
        np.zeros(n_samples_per_class),
        np.ones( n_samples_per_class),
    ]).astype('int'+('32' if theano.config.floatX.endswith('32') else '64'))

    return X,Y

def train_neural_network(
    inputs,
    outputs,
    n_hidden=10,
    n_samples_per_class=100,
    n_iters=1000,
    learning_rate=0.1
):
    """
    Create neural network model and train it.
    """
    X = theano.shared(value=inputs,  name='X')
    Y = theano.shared(value=outputs, name='Y')

    # Input to hidden weights.
    Wih = theano.shared(
            rs.randn(2, n_hidden).astype(theano.config.floatX),
            name='input to hidden weight matrix' )
    Bih = theano.shared(
            rs.randn(n_hidden).astype(theano.config.floatX),
            name='input to hidden bias vector' )

    # Hidden to output weights.
    Who = theano.shared(
            rs.randn(n_hidden, 2).astype(theano.config.floatX),
            name='hidden to output weight matrix' )
    Bho = theano.shared(
            rs.randn(2).astype(theano.config.floatX),
            name='hidden to output bias vector' )

    H = T.tanh(T.dot(X,Wih) + Bih)
    H.name = 'Hidden layer output'
    O = T.nnet.softmax(T.dot(H,Who) + Bho)
    O.name = 'Output layer'
    loss = -T.log(O)[T.arange(Y.shape[0]), Y].mean()
    loss.name = 'loss'

    params = [Wih, Bih, Who, Bho]

    # Get the gradients.
    gparams = []
    for param in params:
        gparams.append( T.grad(cost=loss, wrt=param) )

    # Gradient descent updates.
    updates = []
    for param,gparam in zip(params,gparams):
        updates.append((param, param - learning_rate*gparam))

    train_net = theano.function(
        inputs = [],
        outputs = loss,
        updates = updates
    )

    loss_record = np.zeros(n_iters)
    for i in range(n_iters):
        loss_record[i] = train_net()
        sys.stdout.write(
            "Iteration %d / %d, Loss: %.6f\r" % (i+1, n_iters, loss_record[i])
        )
        sys.stdout.flush()
    print("")


    params = [p.eval() for p in params]
    def output_func(X):
        H = np.tanh(np.dot(X,params[0]) + params[1])
        O = np.exp(-np.dot(H,params[2]) - params[3])
        O = 1. / (1+O)
        O /= O.sum(axis=1).reshape(O.shape[0], 1)
        return O[:,1]

    return loss_record, output_func
    
def plot_stuff(inputs, outputs, losses, net_func, n_hidden):
    fig,axes = plt.subplots(1,2,figsize=(12,6))

    axes[0].plot(np.arange(losses.shape[0])+1, losses)
    axes[0].set_xlabel('iteration')
    axes[0].set_ylabel('loss')
    axes[0].set_xscale('log')
    axes[0].set_yscale('log')

    x,y = np.mgrid[inputs[:,0].min():inputs[:,0].max():51j, inputs[:,1].min():inputs[:,1].max():51j]
    z = net_func( np.c_[x.flatten(), y.flatten()] ).reshape(x.shape)

    axes[1].contourf(x,y,z, cmap=plt.cm.RdBu, alpha=0.6)
    axes[1].plot(inputs[outputs==0,0], inputs[outputs==0,1], 'or') 
    axes[1].plot(inputs[outputs==1,0], inputs[outputs==1,1], 'sb') 
    axes[1].set_title('Percent missclassified: %0.2f%%' % (((net_func(inputs)>0.5) != outputs.astype(np.bool)).mean()*100))

    fig.suptitle('Shallow net with %d hidden units'%n_hidden)
    plt.show()

if __name__=='__main__':
    n_hidden = 36
    inputs, outputs = gen_data(n_samples_per_class=100)
    losses, net_func = train_neural_network(inputs=inputs, outputs=outputs, n_hidden=n_hidden, n_iters=int(1e5), learning_rate=0.1)
    plot_stuff(inputs, outputs, losses, net_func, n_hidden)