Fix model saving/loading

Author: michaelmior

.gitignore

@@ -9,6 +9,7 @@ patterns.json
 *.npy
 *.json
 *.h5
+*.keras
 *.png
 *.parquet
 regex101/

README.md

@@ -32,7 +32,7 @@ This will generate `preprocessed_train.txt` which contains all the feature vecto
 
 `pipenv run python train.py`
 
-The model architecture will be stored in `nn_model_sherlock.json` with the weights in `nn_model_sherlock.weights.h5`.
+The model architecture will be stored in `nn_model_sherlock.json` with the weights in `nn_model_sherlock.weights.keras`.
 
 ## Evaluation
 

explain.py

@@ -46,8 +46,7 @@ def update_sample(samples, N, sample):
 )
 
 # Load the trained model
-model = tf.keras.models.model_from_json(open("nn_model_sherlock.json").read())
-model.load_weights("nn_model_sherlock.weights.h5")
+model = tf.keras.models.load_model(os.path.join(args.input_dir, "nn_model_sherlock.keras"))
 
 # Produce a randomly sample of background from the training data
 background = []

test.py

@@ -7,7 +7,7 @@
 from pyarrow.parquet import ParquetFile
 from sklearn.metrics import classification_report
 from sklearn.preprocessing import LabelEncoder
-from tensorflow.keras.models import model_from_json
+from tensorflow.keras.models import load_model
 from tqdm import tqdm
 
 BATCH_SIZE = 1000
@@ -27,10 +27,7 @@
 # labels = le.transform(labels.values.ravel())
 num_examples = len(labels)
 
-model = model_from_json(
-    open(os.path.join(args.input_dir, "nn_model_sherlock.json"), "r").read()
-)
-model.load_weights(os.path.join(args.input_dir, "nn_model_sherlock.weights.h5"))
+model = load_model(os.path.join(args.input_dir, "nn_model_sherlock.keras"))
 
 sys.stderr.write("Evaluating...\n")
 labels_pred = [""] * len(labels)

train.py

@@ -51,31 +51,31 @@
 regex_model1 = BatchNormalization(axis=1)(regex_model_input)
 regex_model2 = Dense(
     1000,
-    activation=tf.nn.relu,
+    activation='relu',
     kernel_regularizer=tf.keras.regularizers.l2(0.0001),
 )(regex_model1)
 regex_model3 = Dropout(0.35)(regex_model2)
 regex_model4 = Dense(
     1000,
-    activation=tf.nn.relu,
+    activation='relu',
     kernel_regularizer=tf.keras.regularizers.l2(0.0001),
 )(regex_model3)
 
 merged_model2 = BatchNormalization(axis=1)(regex_model4)
 merged_model3 = Dense(
     500,
-    activation=tf.nn.relu,
+    activation='relu',
     kernel_regularizer=tf.keras.regularizers.l2(0.0001),
 )(merged_model2)
 merged_model4 = Dropout(0.35)(merged_model3)
 merged_model5 = Dense(
     500,
-    activation=tf.nn.relu,
+    activation='relu',
     kernel_regularizer=tf.keras.regularizers.l2(0.0001),
 )(merged_model4)
 merged_model_output = Dense(
     len(le.classes_),
-    activation=tf.nn.softmax,
+    activation='softmax',
     kernel_regularizer=tf.keras.regularizers.l2(0.0001),
 )(merged_model5)
 
@@ -86,9 +86,6 @@
     loss="categorical_crossentropy",
     metrics=["categorical_accuracy"],
 )
-open(os.path.join(args.output_dir, "nn_model_sherlock.json"), "w").write(
-    model.to_json()
-)
 
 preprocessed = open(os.path.join(args.input_dir, "preprocessed_train.txt"), "r")
 i = 0
@@ -115,4 +112,4 @@
         pbar.update(len(matrix))
 
 # Save the trained model weights
-model.save_weights(os.path.join(args.output_dir, "nn_model_sherlock.weights.h5"))
+model.save(os.path.join(args.output_dir, "nn_model_sherlock.keras"))