Simple LSTM with tensorflow eager mode on MNIST

LSTM.py

import numpy as np
import tensorflow as tf
import tensorflow.contrib.eager as tfe
tfe.enable_eager_execution()

from tensorflow.examples.tutorials.mnist import input_data
mnist = input_data.read_data_sets("/tmp/data/", one_hot=True)

# Training Parameters
learning_rate = 0.001
training_steps = 100
batch_size = 128
display_step = 200

# Network Parameters
num_input = 28 # MNIST data input (img shape: 28*28)
timesteps = 28 # timesteps
num_hidden = 128 # hidden layer num of features
num_classes = 10 # MNIST total classes (0-9 digits)


class LSTM(tfe.Network):
    def __init__(self,num_hidden,dim):
        super(LSTM, self).__init__(name="")
        self.dim=dim
        self.lstm_cell=self.track_layer(tf.nn.rnn_cell.BasicLSTMCell(num_hidden, forget_bias=1.0))
        self.dense=self.track_layer(tf.layers.Dense(num_classes))

    def call(self,X):
        X = tf.unstack(X,axis= 1)
        outputs, states = tf.nn.static_rnn(self.lstm_cell,X,dtype=tf.float32)
        return self.dense(outputs[-1])

def compute_loss(logits, Y):
    # Define loss and optimizer
    loss_op = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(
        logits=logits, labels=Y))
    return loss_op

def accuracy(logits,Y):
    prediction = tf.nn.softmax(logits)
    correct_pred = tf.equal(tf.argmax(prediction, 1), tf.argmax(Y, 1))
    accuracy = tf.reduce_mean(tf.cast(correct_pred, tf.float32))
    return accuracy

def model_loss(model,x,y):
    logits = model(tf.convert_to_tensor(x, np.float32))
    loss_value = compute_loss(logits, y)
    print("Train Accuracy:{}".format(accuracy(logits,y)))
    return loss_value

model = LSTM(num_hidden,num_input)
optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)
grad_fn = tfe.implicit_value_and_gradients(model_loss)

test_len = 128
test_data = mnist.test.images[:test_len].reshape((-1, timesteps, num_input))
test_label = mnist.test.labels[:test_len]


for step in range(1, training_steps+1):
        batch_x, batch_y = mnist.train.next_batch(batch_size)
        # Reshape data to get 28 seq of 28 elements
        batch_x = batch_x.reshape((batch_size, timesteps, num_input))
        loss_batch, grads = grad_fn(model,batch_x,batch_y)
        optimizer.apply_gradients(grads)
        test_logits=model(tf.convert_to_tensor(test_data, np.float32))
        print("Test Accuracy:{}".format(accuracy(test_logits, test_label)))

enable_eager_execution(): Yes, but not with tensorflow.contrib.eager instead with tensorflow.