def export_model(saver: tf.train.Saver, input_names: list, output_name: str, model_name: str):
"""
You can find node names by using debugger: just connect it right after model is created and look for nodes in the inspec
:param saver:
:param input_names:
:param output_name:
:param model_name:
:return:
"""
os.makedirs("./out", exist_ok=True)
tf.train.write_graph(K.get_session().graph_def, 'out',
model_name + '_graph.pbtxt')
saver.save(K.get_session(), 'out/' + model_name + '.chkp')
# pbtxt is human readable representation of the graph
freeze_graph.freeze_graph('out/' + model_name + '_graph.pbtxt', None,
False, 'out/' + model_name + '.chkp', output_name,
"save/restore_all", "save/Const:0",
'out/frozen_' + model_name + '.pb', True, "")
input_graph_def = tf.GraphDef()
with tf.gfile.Open('out/frozen_' + model_name + '.pb', "rb") as f:
input_graph_def.ParseFromString(f.read())
# optimization of the graph so we can use it in the android app
output_graph_def = optimize_for_inference_lib.optimize_for_inference(
input_graph_def, input_names, [output_name],
tf.float32.as_datatype_enum)
# This is archived optimal graph in the protobuf format we'll use in our android App.
with tf.gfile.FastGFile('out/opt_' + model_name + '.pb', "wb") as f:
f.write(output_graph_def.SerializeToString())
print("graph saved!")
def save_model(self,
sess: tf.Session,
model_params: dict,
saver: tf.train.Saver,
epochs: int) -> None:
saver.save(sess, model_params['model_path'] + '/' + self.name, global_step=epochs)
with open(model_params['model_path'] + '/model_params.json', 'w+') as f:
json.dump(model_params, f)
def model_save(saver: tf.train.Saver, sess: tf.Session, model_path: str,
model_prefix: str, batch_id: int):
try:
saver.save(sess, f"{model_path}/{model_prefix}", global_step=batch_id)
print(f"Successful saving model[{batch_id}] ...")
return True
except Exception as error:
print(f"Failed saving model[{batch_id}] : {error}")
return False
def save_model(sess: tf.Session, path: str, saver: tf.train.Saver = None) -> tf.train.Saver:
"""
Saves a tensorflow session to the given path.
NOTE: This currently saves *all* variables in the session, unless one passes in a custom Saver object.
:param sess: The tensorflow session to save from
:param path: The path to store the saved data
:param saver: A custom saver object to use. This can be used to only save certain variables. If None,
creates a saver object that saves all variables.
:return: The saver object used.
"""
if saver is None:
saver = tf.train.Saver(tf.all_variables())
saver.save(sess, path)
return saver
def saver_fn(i, session: tf.Session, saver: tf.train.Saver):
print('Saving Session {}. . .'.format(i))
saver.save(session, '/tmp/new_save/checkpoint', global_step=i)
def save(saver: tf.train.Saver, sess: tf.Session, checkpoint_dir, step):
if not os.path.exists(checkpoint_dir):
os.makedirs(checkpoint_dir)
saver.save(sess, checkpoint_dir, global_step=step, write_meta_graph=False)
def save_graph(self, epoch, logs, session: tf.Session,
saver: tf.train.Saver, checkpoint_path: str,
save_summary_writer: tf.summary.FileWriter):
save_summary_writer.add_graph(session.graph)
saver.save(session, save_path=f"{checkpoint_path}/{epoch}")
def eval_model(is_training: tf.Variable, sess: tf.Session, best_iou: float,
val_loss: tf.Tensor, val_acc: tf.Tensor,
val_iou_update: tf.Operation, val_iou: tf.Tensor,
val_iou_reset: tf.Operation, val_writer: tf.summary.FileWriter,
epoch: int, saver: tf.train.Saver) -> float:
"""
evaluates model with one pass over validation set
:param is_training: tf var which indicates if model is training
:param sess: tf sess
:param best_iou: best validation iou until now
:param val_loss: val loss tensor
:param val_acc: val accuracy tensor
:param val_iou_update: val iou update operation
:param val_iou: val iou tensor
:param val_iou_reset: val iou reset operation
:param val_writer: val summary writer
:param epoch: index of current epoch
:param saver: tf model saver
:return: new best iou
"""
acc_sum, loss_sum = 0, 0
# toggle training off
assign_op = is_training.assign(False)
sess.run(assign_op)
val_batches = N_VAL_SAMPLES // BATCH_SIZE
print(f"starting evaluation {val_batches} batches")
for j in range(val_batches):
loss_val, acc_val, _, val_iou_val = sess.run(
[val_loss, val_acc, val_iou_update, val_iou])
print(
f"\tevaluation epoch: {epoch:03d}\tbatch {j:03d} eval:"
f"\tloss: {loss_val:.4f}\taccuracy: {acc_val:.4f}\taccumulated iou {val_iou_val:.4f}"
)
acc_sum += acc_val
loss_sum += loss_val
# validation summary
loss = loss_sum / val_batches
acc = acc_sum / val_batches
iou = val_iou_val
summary = get_tf_summary(loss, acc, iou)
val_writer.add_summary(summary, epoch)
print(
f"evaluation:\tmean loss: {loss:.4f}\tmean acc: {acc:.4f}\tmean iou {iou:.4f}\n"
)
# save model if it is better
if iou > best_iou:
best_iou = iou
save_path = saver.save(
sess,
os.path.join(LOG_DIR + "_train",
f"best_model_epoch_{epoch:03d}.ckpt"))
print(f"Model saved in file: {save_path}\n")
# reset accumulator
sess.run(val_iou_reset)
# toggle training on
assign_op = is_training.assign(True)
sess.run(assign_op)
return best_iou
def train(parameters: dict,
path: str = os.path.join('./model', 'model'),
saver: tf.train.Saver = None) -> None:
weights = {
'h1':
tf.Variable(
tf.random_normal(
[parameters['image_size'], parameters['n_hidden_1']])),
'h2':
tf.Variable(
tf.random_normal(
[parameters['n_hidden_1'], parameters['n_hidden_2']])),
'out':
tf.Variable(
tf.random_normal(
[parameters['n_hidden_2'], parameters['n_classes']]))
}
biases = {
'b1': tf.Variable(tf.random_normal(parameters['n_hidden_1'])),
'b2': tf.Variable(tf.random_normal([parameters['n_hidden_2']])),
'out': tf.Variable(tf.random_normal([parameters['n_classes']]))
}
x, y = init_x_y(parameters)
pred = multilayer_perceptron(x, weights, biases)
# Define loss and optimizer
cost = tf.reduce_mean(
tf.nn.softmax_cross_entropy_with_logits(logits=pred, labels=y))
optimizer = tf.train.AdamOptimizer(
learning_rate=parameters['learning_rate']).minimize(cost)
# Initializing the variables
init = tf.global_variables_initializer()
batch_size = parameters['batch_size']
# Launch the graph
with tf.Session() as sess:
sess.run(init)
# Training cycle
for epoch in range(parameters['training_epochs']):
avg_cost = 0.
total_batch = int(mnist.train.num_examples / batch_size)
# Loop over all batches
for i in range(total_batch):
batch_x, batch_y = mnist.train.next_batch(batch_size)
# Run optimization op (backprop) and cost op (to get loss value)
_, c = sess.run([optimizer, cost],
feed_dict={
x: batch_x,
y: batch_y
})
# Compute average loss
avg_cost += c / total_batch
# Display logs per epoch step
if epoch % parameters['display_step'] == 0:
print("Epoch:", '%04d' % (epoch + 1), "cost=",
"{:.9f}".format(avg_cost))
print("Optimization Finished!")
if saver:
saver.save(sess, save_path=path)
# Test model
correct_prediction = tf.equal(tf.argmax(pred, 1), tf.argmax(y, 1))
# Calculate accuracy
accuracy = tf.reduce_mean(tf.cast(correct_prediction, "float"))
print("Accuracy:",
accuracy.eval({
x: mnist.test.images,
y: mnist.test.labels
}))
def save(self, sess: tf.Session, saver: tf.train.Saver, ckpt_dir,
global_step):
self.logger.info("saving model ...")
saver.save(sess,
"{}/model.ckpt".format(ckpt_dir),
global_step=global_step)
