diff --git a/lenet5.py b/lenet5.py
index 7460014..8fd5d31 100644
--- a/lenet5.py
+++ b/lenet5.py
@@ -2,11 +2,25 @@
 LeNet-5 example
 """
 import gzip
+from time import time
 import numpy as np
 from requests import get
 import matplotlib.pyplot as plt
 import pandas as pd
 from sklearn.model_selection import train_test_split
+import keras
+from keras.models import Sequential
+import keras.layers as layers
+from keras.preprocessing.image import ImageDataGenerator
+from keras.utils.np_utils import to_categorical
+from keras.callbacks import TensorBoard
+
+EPOCHS = 10
+BATCH_SIZE = 128
+
+train = {}
+test = {}
+validation = {}
 
 def download_file(url, file_name):
     """
@@ -32,12 +46,12 @@ def read_mnist(images_path: str, labels_path: str):
 
     return features, labels
 
-def display_image(dataset, position):
+def display_image(position):
     """
     Display image at position of the given dataset.
     """
-    image = dataset['features'][position].squeeze()
-    plt.title('Example %d. Label: %d' % (position, dataset['labels'][position]))
+    image = train['features'][position].squeeze()
+    plt.title('Example %d. Label: %d' % (position, train['labels'][position]))
     plt.imshow(image, cmap=plt.get_cmap('gray_r'))
     plt.show()
 
@@ -62,9 +76,6 @@ def main():
     download_file('http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz',
                   't10k-labels-idx1-ubyte.gz')
 
-    train = {}
-    test = {}
-
     # Step 2:
     # Read MNIST dataset (training and testing)
     train['features'], train['labels'] = read_mnist('train-images-idx3-ubyte.gz',
@@ -78,10 +89,10 @@ def main():
     print('Number of test images:', test['features'].shape[0])
 
     # Step 4:
-    # Dispan some images
-    # display_image(train, 0)
-    # display_image(train, 1)
-    # display_image(train, 2)
+    # Display some images
+    display_image(0)
+    display_image(1)
+    display_image(2)
 
     # Step 5:
     # Plot information about the training data
@@ -92,7 +103,6 @@ def main():
 
     # Step 5:
     # Split training data into training and validation
-    validation = {}
     train['features'], validation['features'], train['labels'], validation['labels'] \
         = train_test_split(train['features'], train['labels'], test_size=0.2, random_state=0)
 
@@ -101,13 +111,75 @@ def main():
 
     # Step 6:
     # Prepare our input features.
-    # The LeNet architecture accepts 32x32 pixel images as input, but MNIST data is 28x28 pixels.
+    # The LeNet-5 architecture accepts 32x32 pixel images as input, but MNIST data is 28x28 pixels.
     # We simply pad the imges with zeros to overcome that.
     train['features']      = np.pad(train['features'], ((0,0),(2,2),(2,2),(0,0)), 'constant')
-    validation['features'] = np.pad(validation['features'], ((0,0),(2,2),(2,2),(0,0)), 'constant')
     test['features']       = np.pad(test['features'], ((0,0),(2,2),(2,2),(0,0)), 'constant')
+    validation['features'] = np.pad(validation['features'], ((0,0),(2,2),(2,2),(0,0)), 'constant')
 
     print("Updated Image Shape: {}".format(train['features'][0].shape))
 
+    # Step 7:
+    # Create and compile the model
+    model = Sequential()
+
+    # C1: (None,32,32,1) -> (None,28,28,6).
+    model.add(layers.Conv2D(6, kernel_size=(5, 5), strides=(1, 1), activation='tanh',
+                            input_shape=(32,32,1), padding='valid'))
+
+    # S2: (None,28,28,6) -> (None,14,14,6).
+    model.add(layers.AveragePooling2D(pool_size=(2, 2), strides=(2, 2), padding='valid'))
+
+    # C3: (None,14,14,6) -> (None,10,10,16).
+    model.add(layers.Conv2D(16, kernel_size=(5, 5), strides=(1, 1), activation='tanh',
+                            padding='valid'))
+
+    # S4: (None,10,10,16) -> (None,5,5,16).
+    model.add(layers.AveragePooling2D(pool_size=(2, 2), strides=(2, 2), padding='valid'))
+
+    # Flatten: (None,5,5,16) -> (None, 400).
+    model.add(layers.Flatten())
+
+    # C5: (None, 400) -> (None,120).
+    model.add(layers.Dense(120, activation='tanh'))
+
+    # F6: (None,120) -> (None,84).
+    model.add(layers.Dense(84, activation='tanh'))
+
+    # Output: (None,84) -> (None,10).
+    model.add(layers.Dense(10, activation='softmax'))
+
+    # Compile the model
+    # model.compile(loss='sparse_categorical_crossentropy', optimizer='sgd', metrics=['accuracy'])
+    model.compile(loss=keras.losses.categorical_crossentropy,
+                  optimizer=keras.optimizers.Adam(), metrics=['accuracy'])
+
+    # Step 8:
+    # Train the model
+    x_train, y_train = train['features'], to_categorical(train['labels'])
+    x_validation, y_validation = validation['features'], to_categorical(validation['labels'])
+
+    train_generator = ImageDataGenerator().flow(x_train, y_train, batch_size=BATCH_SIZE)
+    validation_generator = ImageDataGenerator().flow(x_validation, y_validation,
+                                                     batch_size=BATCH_SIZE)
+
+    steps_per_epoch = x_train.shape[0] // BATCH_SIZE
+    validation_steps = x_validation.shape[0] // BATCH_SIZE
+
+    print('Number of training images:', train['features'].shape[0])
+    print('Number of validation images:', validation['features'].shape[0])
+
+    tensorboard = TensorBoard(log_dir="logs/{}".format(time()))
+
+    model.fit_generator(train_generator, steps_per_epoch=steps_per_epoch, epochs=EPOCHS,
+                    validation_data=validation_generator, validation_steps=validation_steps,
+                    shuffle=True, callbacks=[tensorboard])
+
+    # Step 9:
+    # Print results
+    score = model.evaluate(test['features'], to_categorical(test['labels']))
+    print('Test loss:', score[0])
+    print('Test accuracy:', score[1])
+
 if __name__ == "__main__":
     main()