diff --git a/__pycache__/constant.cpython-37.pyc b/__pycache__/constant.cpython-37.pyc
new file mode 100644
index 0000000..cd5a462
Binary files /dev/null and b/__pycache__/constant.cpython-37.pyc differ
diff --git a/__pycache__/model.cpython-37.pyc b/__pycache__/model.cpython-37.pyc
new file mode 100644
index 0000000..1e3b808
Binary files /dev/null and b/__pycache__/model.cpython-37.pyc differ
diff --git a/data.py b/data.py
index ee10f76..c522ff9 100644
--- a/data.py
+++ b/data.py
@@ -1,3 +1,14 @@
+from __future__ import division, print_function, absolute_import
+
+import tensorflow as tf
+from sklearn.model_selection import train_test_split
+from constant import *
+import os
+from six.moves import urllib, xrange
+
+
+SOURCE_URL = 'http://yann.lecun.com/exdb/mnist/'
+
 class Dataset:
     def __init__(self):
         pass
diff --git a/model.py b/model.py
index 9beb974..7567c37 100644
--- a/model.py
+++ b/model.py
@@ -1,53 +1,72 @@
-# TODO 2:  
+# TODO 2:
+import os
 
+os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
 import tensorflow as tf
-
-class HighwayMLP(tf.keras.Model):
-  """
-  Highway MLP Layer
-  """
-  def __init__(self, input_size, t_bias=-2, acti_h = tf.nn.relu, acti_t = tf.nn.tanh):
-    super(HighwayMLP, self).__init__()
-    self.acti_h = acti_h
-    self.acti_t = acti_t
-
-    # TODO 
-    # Dense H
-    self.h = None
-
-    # Dense T
-
-    self.t = None
-
-
-    pass
-
-  def call(self, x):
-    # Do Highway: y = H(x,WH)· T(x,WT) + x · C(x,WC).
-
-    y = None
-
-
-class HighwayNetwork(tf.keras.Model): 
-  """
-  Highway Network with several layers
-  """
-  def __init__(self, input_size, output_size):
-    super(HighwayNetwork, self).__init__()
-    self.mlplayers = [
-      # TO DO
-    ]
-
-    # Classification layer
-
-  def call(self, x):
-    # Run input on these mlp layers
-
-    # pass output to classification layer
-    pass
-
-
-model = HighwayNetwork(784, 10)
-
-
-
+from tensorflow import keras
+from tensorflow.keras import layers
+from tensorflow.python.keras import activations
+
+
+class HighwayBlock(layers.Layer):
+    def __init__(self, units, t_bias, acti_h, acti_t):
+        super(HighwayBlock, self).__init__()
+        self.units = units
+        self.t_bias = t_bias
+        self.acti_t = acti_t
+        self.acti_h = acti_h
+
+    def build(self, input_shape):
+        self.W = self.add_weight(
+            name="w",
+            shape=(input_shape[-1], self.units),
+            initializer="random_normal",
+            trainable=True,
+        )
+        self.W_T = self.add_weight(
+            name="w_T",
+            shape=(input_shape[-1], self.units),
+            initializer="random_normal",
+            trainable=True,
+        )
+
+        self.b = self.add_weight(
+            name="b", shape=(self.units,), initializer="random_normal", trainable=True,
+        )
+
+        self.b_T = tf.Variable(tf.constant(self.t_bias, shape=self.units), name='bias', trainable=True)
+        # assert self.b_T.shape == (50,), 'b.shape: {}'.format(self.b.shape)
+
+    def call(self, inputs):
+        h = self.acti_h(tf.matmul(inputs, self.W) + self.b)
+        t = self.acti_t(tf.matmul(inputs, self.W_T) + self.b_T)
+        y = tf.add(tf.multiply(h, t), tf.multiply(inputs, (1 - t)))
+        return y
+
+
+class HighwayNetwork(tf.keras.Model):
+    """
+    Highway Network with several layers
+    """
+
+    def __init__(self, t_bias=-9.0, acti_h=tf.nn.relu, acti_t=tf.nn.sigmoid, num_classes=10, num_of_layers=3):
+        super(HighwayNetwork, self).__init__()
+        self.projection = keras.Sequential([
+            layers.Conv2D(16, 3, padding='same', activation='relu', input_shape=(28, 28, 1)),
+            layers.MaxPooling2D(),
+            layers.Conv2D(32, 3, padding='same', activation='relu'),
+            layers.MaxPooling2D(),
+            layers.Conv2D(64, 3, padding='same', activation='relu'),
+            layers.MaxPooling2D(),
+            layers.Flatten(),
+            layers.Dense(50)])
+        self.mlplayers = [HighwayBlock(50, t_bias=t_bias, acti_h=acti_h, acti_t=acti_t) for _ in range(num_of_layers)]
+
+        self.classifier = layers.Dense(num_classes, activation='softmax')
+
+    def call(self, x):
+        X = self.projection(x)
+        for layer in self.mlplayers:
+            X = layer(X)
+        y = self.classifier(X)
+        return y
\ No newline at end of file
diff --git a/model/highway_network/saved_model.pb b/model/highway_network/saved_model.pb
new file mode 100644
index 0000000..00ea850
Binary files /dev/null and b/model/highway_network/saved_model.pb differ
diff --git a/model/highway_network/variables/variables.data-00000-of-00001 b/model/highway_network/variables/variables.data-00000-of-00001
new file mode 100644
index 0000000..d657894
Binary files /dev/null and b/model/highway_network/variables/variables.data-00000-of-00001 differ
diff --git a/model/highway_network/variables/variables.index b/model/highway_network/variables/variables.index
new file mode 100644
index 0000000..7d699ea
Binary files /dev/null and b/model/highway_network/variables/variables.index differ
diff --git a/predict.py b/predict.py
index 8ab8fd8..e003108 100644
--- a/predict.py
+++ b/predict.py
@@ -1,10 +1,18 @@
-import os
+import tensorflow as tf
+from model import HighwayNetwork
 from argparse import ArgumentParser
+import os
+import matplotlib.pyplot as plt
+import numpy as np
+
 
 if __name__ == "__main__":
+    home_dir = os.getcwd()
     parser = ArgumentParser()
-    parser.add_argument("--batch-size", default=64, type=int)
-    parser.add_argument("--epochs", default=1000, type=int)
+    #parser.add_argument("--test-file-path", default='{}/data/test'.format(home_dir), type=str, required=True)
+    parser.add_argument("--image-size", default=28, type=int)
+    parser.add_argument("--image-index", default=0, type=int)
+    parser.add_argument("--model-folder", default='{}/model/highway_network/'.format(home_dir), type=str)
 
     # FIXME
     args = parser.parse_args()
@@ -12,13 +20,38 @@
     # FIXME
     # Project Description
 
-    print('---------------------Welcome to ${name}-------------------')
+    print('---------------------Welcome to Highway Network-------------------')
     print('Github: ${accout}')
     print('Email: ${email}')
     print('---------------------------------------------------------------------')
-    print('Training ${name} model with hyper-params:') # FIXME
+    print('Predict using Highway Network for image')
     print('===========================')
 
-    # FIXME
-    # Do Training
+    # Loading Model
+    highway_network = tf.keras.models.load_model(args.model_folder)
+
+
+    # Load data mnist
+    mnist = tf.keras.datasets.mnist
+    (train_images, train_labels), (test_images, test_labels) = mnist.load_data()
+
+    # Load test images from folder
+    # image = tf.keras.preprocessing.image.load_img(args.test_file_path, target_size=(args.image_size, args.image_size))
+    # image = tf.image.rgb_to_grayscale(
+    #     image, name=None
+    # )
+    # input_arr = tf.keras.preprocessing.image.img_to_array(image)
+    # img = input_arr.reshape(-1,28,28,1).astype("float32") / 255
+
+    # Normalize data
+
+    # x_train = train_images.reshape(60000, 784).astype("float32") / 255
+    img = test_images.reshape(-1, 28, 28, 1).astype("float32") / 255
+
+    predictions = highway_network.predict(img)
+    print('---------------------Prediction Result: -------------------')
+    print('Output Softmax: {}'.format(predictions[args.image_index]))
+    print('This image belongs to class: {}'.format(np.argmax(predictions[args.image_index]), axis=1))
 
+    plt.imshow(img[args.image_index])
+    plt.show()
diff --git a/train.py b/train.py
index d11d13a..e9f29e6 100644
--- a/train.py
+++ b/train.py
@@ -1,58 +1,85 @@
 import os
 from argparse import ArgumentParser
 from model import HighwayNetwork
+from tensorflow.keras import backend as K
 
 import tensorflow as tf
 
-def accuracy_function(real, pred):
-    pass
-
-
-def train_step(image, label):
-    pass
-
-
-def val_step():
-    pass
-
-
-@tf.function
-def train():
-    pass
     
-
 if __name__ == "__main__":
     parser = ArgumentParser()
-    
+    home_dir = os.getcwd()
+
     # FIXME
     # Arguments users used when running command lines
-    parser.add_argument("--batch-size", default=64, type=int)
-    parser.add_argument("--epochs", default=1000, type=int)
+    parser.add_argument("--t-bias", default=-9.0, type=float)
+    parser.add_argument("--acti-h", default=tf.nn.relu)
+    parser.add_argument("--acti-t", default=tf.nn.sigmoid)
+    parser.add_argument("--batch-size", default=128, type=int)
+    parser.add_argument("--epochs", default=10, type=int)
+    parser.add_argument("--number-of-layers", default=15, type=int)
+    parser.add_argument("--model-folder", default='{}/model/highway_network/'.format(home_dir), type=str)
+    parser.add_argument("--image-size", default=28, type=int)
+    parser.add_argument("--image-channels", default=1, type=int)
+    parser.add_argument("--learning-rate", default=0.01, type=float)
 
-    home_dir = os.getcwd()
     args = parser.parse_args()
 
+
     # FIXME
     # Project Description
 
     print('---------------------Welcome to ${name}-------------------')
     print('Github: ${accout}')
-    print('Email: ${email}')
+    print('Email: tranquan030894@gmail.com')
     print('---------------------------------------------------------------------')
-    print('Training ${name} model with hyper-params:') # FIXME
+    print('Training Highway Network model with hyper-params:') # FIXME
+    print('===========================')
+    for i, arg in enumerate(vars(args)):
+        print('{}.{}: {}'.format(i, arg, vars(args)[arg]))
     print('===========================')
-    
-    # FIXME
-    # Do Prediction 
-
-    batch_size = 64
-    highway_number = 5
-    epochs = 10
-    log_interval = 10
 
+    # FIXME
+    # Do Prediction
 
     # TODO 1: Load MNIST
 
+    number_class = 10
+    mnist = tf.keras.datasets.mnist
+    (train_images, train_labels), (test_images, test_labels) = mnist.load_data()
+
+
+    # Normalize data
+
+    x_train = train_images.astype("float32")/ 255
+    x_test = test_images.astype("float32") /255
+
+    x_train = x_train.reshape(-1, 28, 28, 1)
+    x_test = x_test.reshape(-1, 28, 28, 1)
+
+    model = HighwayNetwork(t_bias=args.t_bias,
+                           acti_h=args.acti_h,
+                           acti_t=args.acti_t,
+                           num_of_layers=args.number_of_layers)
+
+    # Set up loss function
+    loss_object = tf.keras.losses.SparseCategoricalCrossentropy()
+    # Optimizer Definition
+    # All networks were optimized using SGD with momentum # Paper
+    sgd = tf.keras.optimizers.SGD(learning_rate=args.learning_rate, momentum=0.09, nesterov=False, name="SGD")
+    # Compile optimizer and loss function into model
+    model.compile(optimizer=sgd, loss=loss_object, metrics=['acc'])
+
+    # Do training model
+    model.fit(
+        x_train, train_labels,
+        epochs=args.epochs,
+        batch_size=args.batch_size,
+        validation_data= (x_test, test_labels),
+    )
+    #
+    #saving model
+    model.save(args.model_folder)
 
     # TODO 3: Custom training