simple-mlp-gradient-check.py

# coding: utf-8

# In[1]:


get_ipython().magic('matplotlib inline')

import matplotlib.pyplot as plt
import numpy as np
from itertools import count
from collections import deque
import visdom
import time
import random
import copy


# ### resources
# 
# * [cs231n - gradient checks](http://cs231n.github.io/neural-networks-3/#gradcheck)
# * [so - post](https://stackoverflow.com/questions/44249350/neural-network-numerical-gradient-check-not-working-with-matrices-using-python-n)
# 

# In[2]:


def eval_numerical_gradient(f, x):
    """ 
    a naive implementation of numerical gradient of f at x 
    - f should be a function that takes a single argument
    - x is the point (numpy array) to evaluate the gradient at
    """ 

    fx = f(x) # evaluate function value at original point
    grad = np.zeros(x.shape)
    h = 0.00001

    # iterate over all indexes in x
    it = np.nditer(x, flags=['multi_index'], op_flags=['readwrite'])
    while not it.finished:

        # evaluate function at x+h
        # ix.multi_index has shape (Int,) if input has shape (IntxInt)
        ix = it.multi_index

        old_value = x[ix]
        x[ix] = old_value + h # increment by h
        fx_plus_h = f(x) # evalute f(x + h)
        x[ix] = old_value # restore to previous value (very important!)
        
        old_value = x[ix]
        x[ix] = old_value - h # decrement by h
        fx_minus_h = f(x) # evalute f(x - h)
        x[ix] = old_value # restore to previous value (very important!)
        grad_ix_lhs = (fx_plus_h - fx_minus_h) / (2 * h)
        # grad_ix_lhs = (fx_plus_h - fx) / h 
        # compute the partial derivative
        # grad[np.array(ix)] = grad_ix_lhs # the slope
        # grad[ix] = (fxh - fx) / h # the slope
        grad[ix] = grad_ix_lhs
        
        it.iternext() # step to next dimension
    return grad

def eval_numerical_gradient_v2(f, x, h=1e-6, diff_tolerance=1e-5):
    """ Gradient check for a function f.

    Arguments:
    f -- a function that takes a single argument (x) and outputs the
         cost (fx) and its gradients grad
    x -- the point (numpy array) to check the gradient at
    
    Usage:
    >> sigmoid_f = lambda x: (np.sum(sigmoid(x)), dsigmoid(x))
    >> x = np.array([2.0, 0.5]) # .reshape((2,1))
    >> gradcheck_naive(sigmoid_f, x) #this works fine
    Gradient check passed!
    """
    rndstate = random.getstate()
    random.setstate(rndstate)
    fx, grad = f(x) # Evaluate function value at original point
    #fx=cost
    #grad=gradient      
    # Iterate over all indexes in x
    it = np.nditer(x, flags=['multi_index'], op_flags=['readwrite'])
    while not it.finished:
        ix = it.multi_index #multi-index number

        random.setstate(rndstate)
        xp = copy.deepcopy(x)
        xp[ix] += h
        fxp, gradp = f(xp) 

        random.setstate(rndstate)
        xn = copy.deepcopy(x)
        xn[ix] -= h
        fxn, gradn = f(xn) 
        numgrad = (fxp-fxn) / (2*h)

        # Compare gradients
        reldiff = abs(numgrad - grad[ix]) / max(1, abs(numgrad), abs(grad[ix]))
        if reldiff > diff_tolerance:
            print ("Gradient check failed.")
            print ("First gradient error found at index %s" % str(ix))
            print ("Your gradient: %f \t Numerical gradient: %f" % (
                grad[ix], numgrad))
            return

        it.iternext() # Step to next dimension

    print ("Gradient check passed!")


# In[3]:



def sigmoid(x):
    return 1 / (1 + np.exp(-x))

dsigmoid = lambda x: (1 - sigmoid(x)) * sigmoid(x)

square = lambda x: x ** 2
dsquare = lambda x: 2 * x


# In[4]:

class WeightGradientViz:
    def __init__(self, max_history=50):
        self.max_history = max_history
        self.x_index = deque([], self.max_history)
        self.grads = { 
            'd_W1': [[self._new_deque(),self._new_deque()], [self._new_deque(),self._new_deque()]],
            'd_W2': [[self._new_deque(),self._new_deque()], [self._new_deque(),self._new_deque()]]
        }
        self.wins = { 
            'd_W1': [[None,None], [None, None]],
            'd_W2': [[None,None], [None, None]]
        }
    
    def _new_deque(self):
        return deque([], self.max_history)
    
    def set_grad_viz(self, dw1, dw2, i):
        self.x_index.append(i)
        self.grads['d_W1'][0][0].append(dw1[0,0])
        self.grads['d_W1'][0][1].append(dw1[0,1])
        self.grads['d_W1'][1][0].append(dw1[1,0])
        self.grads['d_W1'][1][1].append(dw1[1,1])
        self.grads['d_W2'][0][0].append(dw2[0,0])
        self.grads['d_W2'][0][1].append(dw2[0,1])
        self.grads['d_W2'][1][0].append(dw2[1,0])
        self.grads['d_W2'][1][1].append(dw2[1,1])
    
    def do_viz(self, gradname, r, c, vis):
        X = np.array(list(self.x_index))
        Y = np.array(list(self.grads[gradname][r][c]))
        opts = {
            'title': '{}[{},{}]'.format(gradname,r,c),
            'width': 370,
            'height': 170
        }
        if self.wins[gradname][r][c] is None:
            self.wins[gradname][r][c] = vis.line(Y,X,opts=opts)
        else:
            vis.line(Y,X,win=self.wins[gradname][r][c],opts=opts)
            
    def do_viz_all(self, vis):
        self.do_viz('d_W1', 0, 0, vis)
        self.do_viz('d_W1', 0, 1, vis)
        self.do_viz('d_W1', 1, 0, vis)
        self.do_viz('d_W1', 1, 1, vis)
        self.do_viz('d_W2', 0, 0, vis)
        self.do_viz('d_W2', 0, 1, vis)
        self.do_viz('d_W2', 1, 0, vis)
        self.do_viz('d_W2', 1, 1, vis)
        


# In[6]:


raw_lbl_2_class = ['r', 'b']
one_hot_2_class = lambda x: 'b' if x == [0,1] else 'r'

def generate_data(N=100):
    np.random.seed(42)
    x1 = np.random.uniform(-10, 10, size=N)
    x2 = np.random.uniform(-10, 10, size=N)
    f_x1 = (1/3.)*x1**2 - 5
    descisions = x2 <= f_x1    
    one_hot = [x for x in map(lambda x: [1,0] if not x else [0,1], descisions)]
    X = np.stack((x1,x2), axis=1)
    y = np.array(one_hot)
    return X, y

X, t = generate_data(N=500)
print(t.shape)
plt.xlim(-10,10)
plt.ylim(-10,10)
colors = [c for c in map(one_hot_2_class, t.tolist())]
plt.scatter(X[:,0], X[:,1], c=colors)


# In[7]:


vis = visdom.Visdom()



def get_batch(data, batch_size=10):
    idxs = np.arange(len(data)) 
    np.random.shuffle(idxs)
    idxs = idxs[0:batch_size]
    batch = []
    for i in range(batch_size):
        batch.append(data[i])
    return batch

def grad_check_wrt_W_1(W_1, b_1, W_2, b_2, x_in, y_target):
    def go(W):
        # forward pass
        z = np.dot(W, x_in) + b_1
        h = sigmoid(z)
        y = np.dot(W_2,h) + b_2
        loss = (1/2.) * (np.linalg.norm(y - y_target) ** 2)
        return loss
    return go

max_epochs = 1 # 100
stop_criterion = 0.04
batch_size = 50
learning_rate = 0.1
train_data = [z for z in zip(X,t)]

np.random.seed(3)
W_1 = np.random.randn(2,2) * 0.01
b_1 = np.ones((2,1))
W_2 = np.random.randn(2,2) * 0.01
b_2 = np.ones((2,1))

wgrad_viz = WeightGradientViz()
total_iter = 0
check_grads = False

for i in count(1):
    training_batch = get_batch(train_data, batch_size=batch_size)
    batch_loss = 0
    
    for x_in, y_target in training_batch:
        # reshape x_in to column vector
        x_in = x_in.reshape((2,1))
        y_target = y_target.reshape((2,1))
        

        # forward pass
        z = np.dot(W_1, x_in) + b_1
        h = sigmoid(z)
        y = np.dot(W_2,h) + b_2

        # note: y - y_target is flipped in HW
        loss = (1/2.) * (np.linalg.norm(y - y_target) ** 2)
        
        batch_loss += loss
        
        # backprop
        dloss = 1
        dy = dloss * (y - y_target)
        dW_2 = np.dot(dy, h.T)
        db_2 = dy
        dh = np.dot(W_2.T, dy)
        dz = dh * (1-h) * h
        dW_1 = np.dot(dz, x_in.T)
        db_1 = dz
        
        if not check_grads:
            
            ########################################################################
            # GRADIENT CHECK ON h
            error_from_child = np.ones((2,1))
            # same functional form as dh
            grad_h = np.dot(W_2.T, error_from_child)
            func_of_h = lambda x: np.dot(W_2,x) + b_2
            compare_res = lambda x: (np.sum(func_of_h(x)), grad_h)
            eval_numerical_gradient_v2(compare_res, h, h=1e-6, diff_tolerance=1e-7)
            
            ########################################################################
            # GRADIENT CHECK ON z
            error_from_child = np.ones((2,1))
            # same functional form as dz
            grad_z = error_from_child * (1-h) * h
            func_of_z = lambda x: sigmoid(x)
            compare_res = lambda x: (np.sum(func_of_z(x)), grad_z)
            eval_numerical_gradient_v2(compare_res, z, h=1e-6, diff_tolerance=1e-7)
            
            ########################################################################
            # GRADIENT CHECK ON W_2
            error_from_child = np.ones((2,1))
            # same functional form as dW_2
            grad_W_2 = np.dot(error_from_child, h.T)
            func_of_W_2 = lambda x: np.dot(x,h) + b_2
            compare_res = lambda x: (np.sum(func_of_W_2(x)), grad_W_2)
            eval_numerical_gradient_v2(compare_res, W_2, h=1e-6, diff_tolerance=1e-7)
            
            ########################################################################
            # GRADIENT CHECK ON W_1
            func_of_W_1 = grad_check_wrt_W_1(W_1, b_1, W_2, b_2, x_in, y_target)
            
            ########################################################################
            # GRADIENT CHECK ON W_1 test 1
            numerical_gradient = eval_numerical_gradient(func_of_W_1, W_1)
            # learn how to use these
            # print(numerical_gradient)
            
            ########################################################################
            # GRADIENT CHECK ON W_1 test 2
            compare_res = lambda x: (np.sum(func_of_W_1(x)), dW_1)
            eval_numerical_gradient_v2(compare_res, W_1, h=1e-6, diff_tolerance=1e-7)            
            
            # check_grads = True
            print('')
            
            

        # wgrad_viz.set_grad_viz(dW_1, dW_2, total_iter)
        
        W_1 += -learning_rate * dW_1
        b_1 += -learning_rate * db_1
        W_2 += -learning_rate * dW_2
        b_2 += -learning_rate * db_2
        
        # wgrad_viz.do_viz_all(vis)
        
        total_iter += 1
    
    batch_loss = batch_loss / batch_size
    
    if i >= max_epochs or batch_loss <= stop_criterion:
        break
    
    if i % 1000 == 0:
        print('batch_loss: {}'.format(batch_loss))

# In[9]:


def forward_viz(x):
    _x = x.reshape((2,1))

    # forward pass
    _z = np.dot(W_1, _x) + b_1
    _h = sigmoid(_z)
    _y = np.dot(W_2,_h) + b_2
    return np.argmax(_y)

point_size = 25
tx1 = np.linspace(-10,10,point_size)
tx2 = np.linspace(-10,10,point_size)
tx1,tx2 = np.meshgrid(tx1,tx2)
X_test = np.array([z for z in zip(tx1.ravel(),tx2.ravel())])
lbl_accum = []
for i in range(X_test.shape[0]):
    pred = forward_viz(X_test[i,:])
    # print(pred)
    lbl_accum.append(raw_lbl_2_class[pred])

plt.scatter(X_test[:,0], X_test[:,1], c=lbl_accum)