Restricted Boltzmann Machine for the MNIST dataset implemented in pure NumPy

rbm.py

import numpy as np
import matplotlib.pyplot as plt
from tensorflow.examples.tutorials.mnist import input_data


def sample_vector(v, w, b, bernoulli=True):
    h = 1 / (1 + np.exp(-(v.dot(w) + b)))
    if bernoulli:
        h = np.random.binomial(1, h)
    return h


def train():
    mnist = input_data.read_data_sets('MNIST_data', one_hot=True)
    X = mnist.train.images.round()
    W = 0.01 * np.random.randn(256, 81)
    bv = np.zeros(256)
    bh = np.zeros(81)
    # W = np.load('dbn_params/W2.npy')
    # bv = np.load('dbn_params/bv2.npy')
    # bh = np.load('dbn_params/bh2.npy')
    Es = []
    W0 = np.load('dbn_params/W0.npy')
    W1 = np.load('dbn_params/W1.npy')
    bh0 = np.load('dbn_params/bh0.npy')
    bh1 = np.load('dbn_params/bh1.npy')
    for i in xrange(5000):
        id = np.random.choice(55000, 128, replace=False)
        x = X[id, :]
        x = sample_vector(x, W0, bh0, bernoulli=False)
        # x = sample_vector(x, W1, bh1, bernoulli=False)
        h0 = sample_vector(x, W, bh)
        xh0 = x.T.dot(h0) / 128
        v1 = sample_vector(h0, W.T, bv)
        h1 = sample_vector(v1, W, bh, bernoulli=False)
        xh1 = v1.T.dot(h1) / 128
        E = -np.mean(x.dot(bv) + h0.dot(bh) + np.diag(x.dot(W).dot(h0.T)))
        Es.append(E)
        if i % 100 == 0:
            print E
        W += 0.0075 * (xh0 - xh1)
        bv += 0.0075 * (x - v1).mean(axis=0)
        bh += 0.0075 * (h0 - h1).mean(axis=0)
    plt.plot(Es)
    plt.show()
    update = raw_input('Update parameters? [Y/n]')
    if update == 'Y':
        np.save('dbn_params/W2.npy', W)
        np.save('dbn_params/bv2.npy', bv)
        np.save('dbn_params/bh2.npy', bh)


def test():
    mnist = input_data.read_data_sets('MNIST_data', one_hot=True)
    X = mnist.test.images
    x = X[np.random.randint(10000), :].round()
    del X
    plt.imshow(x.reshape(28, 28), cmap='Greys')
    plt.xticks([])
    plt.yticks([])
    plt.show()
    l = 3
    W, bh, bv = [], [], []
    for i in xrange(l):
        W.append(np.load('dbn_params/W{0}.npy'.format(i)))
        bh.append(np.load('dbn_params/bh{0}.npy'.format(i)))
        bv.append(np.load('dbn_params/bv{0}.npy'.format(i)))
    for i in xrange(l):
        x = sample_vector(x, W[i], bh[i])
    fig, ax = plt.subplots(ncols=5, nrows=5, sharex=True, sharey=True)
    ax = ax.flatten()
    for i in xrange(25):
        img = np.copy(x)
        img[i] = np.abs(img[i] - 1)
        for k in reversed(xrange(l)):
            img = sample_vector(img, W[k].T, bv[k], bernoulli=False)
        ax[i].imshow(img.reshape(28, 28), cmap='Greys', interpolation='nearest')
    ax[0].set_xticks([])
    ax[0].set_yticks([])
    plt.show()


if __name__ == '__main__':
    test()