TensorFlow MLP loss increasing When I train my model the loss increases over each epoch. I feel like this is a simple solve and I am missing something obvious but I cannot figure out what is it. Any help would be greatly appreciated.

The neural network:

def neural_network(data):
    hidden_L1 = {'weights': tf.Variable(tf.random_normal([784, neurons_L1])),
                'biases': tf.Variable(tf.random_normal([neurons_L1]))}

    hidden_L2 = {'weights': tf.Variable(tf.random_normal([neurons_L1, neurons_L2])),
                'biases': tf.Variable(tf.random_normal([neurons_L2]))}

    output_L = {'weights': tf.Variable(tf.random_normal([neurons_L2, num_of_classes])),
                'biases': tf.Variable(tf.random_normal([num_of_classes]))}

    L1 = tf.add(tf.matmul(data, hidden_L1['weights']), hidden_L1['biases']) #matrix multiplication
    L1 = tf.nn.relu(L1)

    L2 = tf.add(tf.matmul(L1, hidden_L2['weights']), hidden_L2['biases']) #matrix multiplication
    L2 = tf.nn.relu(L2)

    output = tf.add(tf.matmul(L2, output_L['weights']), output_L['biases']) #matrix multiplication
    output = tf.nn.softmax(output)

    return output


My loss, optimiser and loop for each epoch:

output = neural_network(x)
loss = tf.reduce_mean( tf.nn.softmax_cross_entropy_with_logits(logits=output, labels=y) )
optimiser = tf.train.AdamOptimizer().minimize(loss)

init = tf.global_variables_initializer()

epochs = 5
total_batch_count = 60000//batch_size

with tf.Session() as sess:
    sess.run(init)

    for epoch in range(epochs):

        avg_loss = 0 

        for i in range(total_batch_count):

            batch_x, batch_y = next_batch(batch_size, x_train, y_train)

            _, c = sess.run([optimiser, loss], feed_dict = {x:batch_x, y:batch_y})

            avg_loss +=c/total_batch_count

            print("epoch = ", epoch + 1, "loss =", avg_loss)

    sess.close()


I have a feeling my problems lies in the either the loss function or the loop I wrote for each epoch, however I am new to TensorFlow and cannot figure this out.
You are using the function softmax_cross_entropy_with_logits which, according to Tensorflow's documentation, has the following specification for logits,


logits: Per-label activations, typically a linear output. These activation energies are interpreted as unnormalized log probabilities.


Hence, you should pass the activations before the non-linearity application (in your case, softmax). You can fix it by doing the following,

def neural_network(data):
    hidden_L1 = {'weights': tf.Variable(tf.random_normal([784, neurons_L1])),
                'biases': tf.Variable(tf.random_normal([neurons_L1]))}

    hidden_L2 = {'weights': tf.Variable(tf.random_normal([neurons_L1, neurons_L2])),
                'biases': tf.Variable(tf.random_normal([neurons_L2]))}

    output_L = {'weights': tf.Variable(tf.random_normal([neurons_L2, num_of_classes])),
                'biases': tf.Variable(tf.random_normal([num_of_classes]))}

    L1 = tf.add(tf.matmul(data, hidden_L1['weights']), hidden_L1['biases']) #matrix multiplication
    L1 = tf.nn.relu(L1)

    L2 = tf.add(tf.matmul(L1, hidden_L2['weights']), hidden_L2['biases']) #matrix multiplication
    L2 = tf.nn.relu(L2)

    logits = tf.add(tf.matmul(L2, output_L['weights']), output_L['biases']) #matrix multiplication
    output = tf.nn.softmax(logits)

    return output, logits


Then, outside your function, you can retrieve the logits, and pass it to your loss function, as in the example bellow,

output, logits = neural_network(x)
loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=logits,
                                                             labels=y))


I remark that you may still be interested in the outputs tensor, for calculating your network's accuracy. If this substitution doesn't work, you should also experiment with the learning rate parameter on your AdamOptimizer (see the documentation here).