Maximum Noise Entropy implementation using TensorFlow

maximum_noise_entropy.py

"""Implementation of maximum noise entropy using TensorFlow.

Paper: http://journals.plos.org/ploscompbiol/article?id=10.1371/journal.pcbi.1002249
"""

# For Python 3 compatibility.
from __future__ import print_function

# For building the algorithm.
import tensorflow as tf
import numpy as np

# Loads the MNIST dataset.
from tensorflow.examples.tutorials.mnist import input_data
mnist = input_data.read_data_sets("MNIST_data/", one_hot=False)

nb_train = 1000  # Number of training steps for each layer.
batch_size = 1000  # Number of samples in a batch.
learning_rate = 0.001 # Gradient descent learning rate.
print_every = 10  # Print loss every `print_every` steps.
good_num = 1  # The number to look for.

# Returns Numpy arrays with the initial weights.
W_init = lambda shape, dtype, partition_info: np.random.normal(scale=0.05, size=shape)
b_init = lambda shape, dtype, partition_info: np.zeros(shape)

# This is a placeholder for the input data.
data_pl = tf.placeholder('float32', shape=(batch_size, 28 * 28))
cat_pl = tf.placeholder('int32', shape=(batch_size,))

# These will store the layer weights and outputs.
y_hats, layer_weights = [], []

# This is the optimizer that the training part will use.
opt = tf.train.AdamOptimizer(learning_rate=learning_rate)

with tf.Session() as sess:

  J = tf.get_variable('J',
      shape=(28 * 28, 28 * 28),
      dtype='float32',
      initializer=W_init)

  h = tf.get_variable('h',
      shape=(1, 28 * 28),
      dtype='float32',
      initializer=W_init)

  a = tf.get_variable('a',
      shape=(),
      dtype='float32',
      initializer=b_init)

  h_m = tf.reduce_sum(data_pl * h, axis=1)
  j_m = tf.reduce_sum(tf.matmul(data_pl, J) * data_pl, axis=1)

  is_num = tf.cast(tf.equal(cat_pl, good_num), 'float32')
  loss = tf.nn.sigmoid_cross_entropy_with_logits(a + h_m + j_m, is_num)
  loss = tf.reduce_mean(loss)

  min_op = opt.minimize(loss, var_list=[J, h, a])

  # Initialize all variables.
  sess.run(tf.global_variables_initializer())

  for i in range(nb_train):
    X_data, y_data = mnist.train.next_batch(batch_size)
    # X_data = X_data * 2 - 1
    if not i % print_every:
      loss_v = sess.run(loss, feed_dict={data_pl: X_data, cat_pl: y_data})
      print('loss on step %d: %f' % (i, loss_v))
    sess.run(min_op, feed_dict={data_pl: X_data, cat_pl: y_data})

  np.save('J.npy', sess.run(J))