diff --git a/lenet5.py b/lenet5.py
index 68b3878..8455d58 100644
--- a/lenet5.py
+++ b/lenet5.py
@@ -56,6 +56,71 @@ def display_image(dataset, position):
     plt.imshow(image, cmap=plt.get_cmap('gray_r'))
     plt.show()
 
+def create_lenet5():
+    """
+    Creates the LeNet-5 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'])
+
+    print(model.summary())
+
+    return model
+
+def train_lenet5(model):
+    """
+    Trains the LeNet-5 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])
+
 def main():
     """
     Defined starting point of source code.
@@ -121,64 +186,14 @@ def main():
     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'])
+    # Create and compile the LeNet-5 model
+    model = create_lenet5()
 
     # 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])
+    # Train the LeNet-5 model
+    train_lenet5(model)
 
     model.save('lenet5.h5')
-    """
 
     model = load_model('lenet5.h5')