You can follow along with the code in the Jupyter notebook ch-09a_CNN_MNIST_TF_and_Keras.
Prepare the MNIST data into test and train sets:
from tensorflow.examples.tutorials.mnist import input_data
mnist = input_data.read_data_sets(os.path.join('.','mnist'), one_hot=True)
X_train = mnist.train.images
X_test = mnist.test.images
Y_train = mnist.train.labels
Y_test = mnist.test.labels