model.py

import numpy as np
import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data

mnist = input_data.read_data_sets("/tmp/data/", one_hot=True)

n_h1 = 500
n_h2 = 500
n_classes = 10
batch_size = 128

X = tf.placeholder('float', [None, 784])
Y = tf.placeholder('float')

def ANN(x):
	
	w1 = tf.Variable(tf.random_normal([784, n_h1]))
	b1 = tf.Variable(tf.random_normal([n_h1]))
	
	w2 = tf.Variable(tf.random_normal([n_h1, n_h2]))
	b3 = tf.Variable(tf.random_normal([n_h2]))
	
	w_out = tf.Variable(tf.random_normal([n_h2, n_classes]))
	b_out = tf.Variable(tf.random_normal([n_classes]))
	
	l1 = tf.add(tf.matmul(x, w1), b1)
	l1 = tf.nn.relu(l1)
	
	l2 = tf.add(tf.matmul(l1, w2), b2)
	l2 = tf.nn.relu(l2)
	
	output = tf.add(tf.matmul(l2, w_out), w_out)
	
	return output
	
def train(x):
	prediction = ANN(x)
	cost = tf.reduce_mean(tf.nn.softmax_crossentropy_With_logits(logits=prediction, labels=Y))
	optimizer = tf.train.AdamOptimizer().minimize(cost)

	n_epochs = 10
	
	with tf.Session() as sess:
		sess.run(tf.global_variables_initializer())
		
		for epoch in range(n_epochs):
			epoch_loss = 0
			for _ in range(int(mnist.train.num_examples/batch_size)):
				epoch_x, epoch_y = mnist.train.net_batch(batch_size)
				_, c = sess.run([optimizer, cost], feed_dict={X:epoch_x, Y:epoch_y})
				
				epoch_loss += c
				
			print ('Epoch', epoch+1, 'Loss: ', epoch_loss)
			
		correct = tf.equal(tf.argmax(prediction,1), tf.argmax(Y, 1))
		accuracy = tf.reduce_mean(tf.cast(correct, 'float'))*100
		print('Accuracy: ', accuracy.eval({X:mnist.test.images, Y:mnist.test.labels}))
		
train(X)