diff --git a/prepare.py b/prepare.py
index c3908c3..703f4dd 100644
--- a/prepare.py
+++ b/prepare.py
@@ -22,9 +22,13 @@ def load_data(path):
     return data
 
 
-def tokenize_data(sdata, mxlen):
+def init_tokenizer(sdata):
     texts = [t[1] for t in sdata]
     tokenizer.fit_on_texts(texts)
+
+
+def tokenize_data(sdata, mxlen):
+    texts = [t[1] for t in sdata]
     seqs = tokenizer.texts_to_sequences(texts)
     seqs = pad_sequences(seqs, mxlen)
     data = {}
diff --git a/train.py b/train.py
index 8a9a3e5..053ee70 100644
--- a/train.py
+++ b/train.py
@@ -1,40 +1,51 @@
 import numpy as np
 import keras
 from keras.models import Sequential, Model
-from keras.layers import Dense, Dropout, Activation, Flatten, Input, Embedding,\
+from keras.layers import Dense, Dropout, Activation, Flatten, Input, Embedding, \
     LSTM, Bidirectional, Lambda, Concatenate, Add
 from keras.layers.convolutional import Conv2D, MaxPooling2D, AveragePooling2D
 from keras.layers.normalization import BatchNormalization
-from keras.optimizers import Adam
+from keras.optimizers import Adam, RMSprop
 import gc
 import prepare
 import subprocess
+import pickle
 
 mxlen = 32
 embedding_dim = 50
 lstm_unit = 128
-MLP_unit = 128
+MLP_unit = 256
 epochs = 100
+batch_size = 128
 
 train_json = 'nlvr\\train\\train.json'
 train_img_folder = 'nlvr\\train\\images'
+test_json = 'nlvr\\test\\test.json'
+test_img_folder = 'nlvr\\test\\images'
 data = prepare.load_data(train_json)
+prepare.init_tokenizer(data)
 data = prepare.tokenize_data(data, mxlen)
-imgs, ws, labels = prepare.load_images(train_img_folder, data, debug=True)
+imgs, ws, labels = prepare.load_images(train_img_folder, data)
 data.clear()
+
+test_data = prepare.load_data(test_json)
+test_data = prepare.tokenize_data(test_data, mxlen)
+test_imgs, test_ws, test_labels = prepare.load_images(test_img_folder, test_data)
+test_data.clear()
+
 imgs_mean = np.mean(imgs)
 imgs_std = np.std(imgs - imgs_mean)
 imgs = (imgs - imgs_mean) / imgs_std
 
-epochs = 100
-batch_size = 64
+test_imgs = (test_imgs - imgs_mean) / imgs_std
 
 
 def bn_layer(x, conv_unit):
     def f(inputs):
-        md = Conv2D(x, (conv_unit, conv_unit), padding='same')(inputs)
+        md = Conv2D(x, (conv_unit, conv_unit), padding='same', kernel_initializer='he_normal')(inputs)
         md = BatchNormalization()(md)
         return Activation('relu')(md)
+
     return f
 
 
@@ -54,25 +65,31 @@ def conv_net(inputs):
 input1 = Input((50, 200, 3))
 input2 = Input((mxlen,))
 cnn_features = conv_net(input1)
-embedding_layer = prepare.embedding_layer(prepare.tokenizer.word_index, prepare.get_embeddings_index(), mxlen)
-embedding = embedding_layer(input2)
+# embedding_layer = prepare.embedding_layer(prepare.tokenizer.word_index, prepare.get_embeddings_index(), mxlen)
+# embedding = embedding_layer(input2)
+embedding = Embedding(mxlen, embedding_dim)(input2)
 bi_lstm = Bidirectional(LSTM(lstm_unit, implementation=2, return_sequences=False))
 lstm_encode = bi_lstm(embedding)
 shapes = cnn_features.shape
 w, h = shapes[1], shapes[2]
 
+
 def slice_1(t):
     return t[:, 0, :, :]
 
+
 def slice_2(t):
     return t[:, 1:, :, :]
 
+
 def slice_3(t):
     return t[:, 0, :]
 
+
 def slice_4(t):
     return t[:, 1:, :]
 
+
 slice_layer1 = Lambda(slice_1)
 slice_layer2 = Lambda(slice_2)
 slice_layer3 = Lambda(slice_3)
@@ -107,31 +124,43 @@ def get_MLP(n, denses):
         for k in range(n):
             d = denses[k](d)
         return d
+
     return g
 
 
-def dropout_dense(x):
+def bn_dense(x):
     y = Dense(MLP_unit)(x)
-    y = Dropout(0.5)(y)
+    y = BatchNormalization()(y)
     y = Activation('relu')(y)
     return y
 
+
 g_MLP = get_MLP(4, get_dense(4))
-f_MLP = get_MLP(2, get_dense(2))
 
 mid_relations = []
 for r in relations:
     mid_relations.append(g_MLP(r))
 combined_relation = Add()(mid_relations)
 
-rn = dropout_dense(combined_relation)
-rn = dropout_dense(rn)
+rn = bn_dense(combined_relation)
+rn = bn_dense(rn)
 pred = Dense(1, activation='sigmoid')(rn)
 
 model = Model(inputs=[input1, input2], outputs=pred)
 optimizer = Adam(lr=3e-5)
 model.compile(optimizer=optimizer, loss='binary_crossentropy', metrics=['accuracy'])
-model.fit([imgs, ws], labels, validation_split=0.1, epochs=epochs)
+for epoch in range(epochs):
+    model.fit([imgs, ws], labels, epochs=1, batch_size=batch_size)
+    p = model.predict([test_imgs, test_ws], batch_size=batch_size)
+    p = np.array([t[0] for t in p])
+    acc = np.sum((p >= 0.5) == (test_labels >= 0.5)) / len(p)
+    avg = np.sum(p) / len(p)
+    print('epoch: ', epoch, ", acc: ", acc, ", avg = ", avg)
+    for k in range(100):
+        print(p[k], test_labels[k])
 model.save('model')
+tokenizer_file = open('tokenizer', 'wb')
+pickle.dump(prepare.tokenizer, tokenizer_file)
+tokenizer_file.close()
 gc.collect()
 subprocess.Popen("rundll32.exe powrprof.dll,SetSuspendState 0,1,0")