forked from tf-encrypted/tf-encrypted
-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
172 lines (121 loc) · 5.11 KB
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
import tensorflow as tf
from tensorflow.python.framework import graph_io
from tensorflow.python.framework import graph_util
from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras import backend as K
import argparse
from tensorflow.python import errors_impl as errors
from utils import read_one_column
import os
import time
from utils import data_prep
parser = argparse.ArgumentParser()
parser.add_argument(
'--predict', type=int, default='-1',
help='Do a prediction on specified row of input file')
parser.add_argument('--input_file', type=str, default='final_data_with_feature_engineered.csv',
help='Location of input defaults to final_data_with_feature_engineered.csv')
parser.add_argument('--bench_prediction', type=int, default=-1,
help='Benchmark prediction doing n iterations and taking the average')
config = parser.parse_args()
epochs = 20
batch_size = 256
checkpoint_path = "./saved_models/train"
predict_row = config.predict
input_file = config.input_file
bench_prediction = config.bench_prediction
def export_to_pb(sess, x, filename):
pred_names = ['output']
tf.identity(x, name=pred_names[0])
graph = graph_util.convert_variables_to_constants(sess, sess.graph.as_graph_def(), pred_names)
graph = graph_util.remove_training_nodes(graph)
path = graph_io.write_graph(graph, ".", filename, as_text=False)
print('saved the frozen graph (ready for inference) at: ', path)
def print_nodes(graph):
print([n.name for n in graph.as_graph_def().node])
def build_model(input_shape):
model = keras.Sequential([
layers.Dense(1, use_bias=False, activation='sigmoid', input_shape=[input_shape]),
])
model.compile(loss='binary_crossentropy',
optimizer=tf.train.AdamOptimizer(),
metrics=['accuracy'])
return model
def train(train_x_df, train_y_df):
x = list(train_x_df.columns.values)
model = build_model(len(x))
# Display training progress by printing a single dot for each completed epoch
class PrintDot(keras.callbacks.Callback):
def on_epoch_end(self, epoch, logs):
print('')
print('.', end='')
os.makedirs('./saved_models', exist_ok=True)
cp_callback = keras.callbacks.ModelCheckpoint(checkpoint_path,
save_weights_only=True,
save_best_only=True,
verbose=1)
# first 80 percent for training
train_x = train_x_df[1:246005]
train_y = train_y_df[1:246005]
# other 20 percent for evaluating
eval_x = train_x_df[246006:len(train_x_df) - 1]
eval_y = train_y_df[246006:len(train_y_df) - 1]
# train model
model.fit(train_x, train_y, epochs=epochs,
validation_split=0.2, verbose=0, batch_size=batch_size,
callbacks=[cp_callback])
print('done training')
# get the default session and graph for exporting and calculating the AUC
sess = K.get_session()
graph = K.get_session().graph
# export the graph to a protobuf file for loading in tfe and secure enclave
export_to_pb(K.get_session(),
graph.get_tensor_by_name('dense/Sigmoid:0'),
'house_credit_default.pb')
# evaluate the model and get the AUC, the metric used in the kaggle competition
loss = model.evaluate(eval_x, eval_y, batch_size=batch_size)
predictions = model.predict(eval_x, batch_size=batch_size)
auc = tf.metrics.auc(eval_y, predictions)
print("Evaluation Loss:", loss[0])
print("Accuracy:", loss[1])
sess.run(tf.global_variables_initializer())
sess.run(tf.local_variables_initializer())
print("AUC: ", sess.run([auc])[0][1])
def predict_preamble(train_x_df, train_y_df):
x = list(train_x_df.columns.values)
model = build_model(len(x))
try:
model.load_weights(checkpoint_path)
except errors.InvalidArgumentError:
print("Weights couldn't be found, training before predicting")
train(train_x_df, train_y_df)
model = build_model(len(x))
x = read_one_column(predict_row, train_x_df)
return model, x
def predict(train_x_df, train_y_df):
model, x = predict_preamble(train_x_df, train_y_df)
print("Prediction:", model.predict(x)[0][0])
def benchmark(train_x_df, train_y_df):
model, x = predict_preamble(train_x_df, train_y_df)
total_duration = 0
for i in range(0, bench_prediction):
start = time.time()
model.predict(x)
end = time.time()
duration = end - start
total_duration = total_duration + duration
print("Total Duration:", total_duration)
print("Avg Runtime:", total_duration / bench_prediction * 1000, "ms")
def main():
print('Home Credit Default!')
# TODO only load all data when training
train_x_df, train_y_df = data_prep(input_file)
if predict_row != -1:
predict(train_x_df, train_y_df)
elif bench_prediction != -1:
benchmark(train_x_df, train_y_df)
else:
train(train_x_df, train_y_df)
if __name__ == '__main__':
main()