from tensorflow.python.tools import freeze_graph

optimize_me_parent_path = '/root/models/optimize_me/linear/cpu'

fully_optimized_model_graph_path = '%s/fully_optimized_cpu.pb' % optimize_me_parent_path
fully_optimized_frozen_model_graph_path = '%s/fully_optimized_frozen_cpu.pb' % optimize_me_parent_path

model_checkpoint_path = '%s/model.ckpt' % optimize_me_parent_path

freeze_graph.freeze_graph(input_graph=fully_optimized_model_graph_path, 
                          input_saver="",
                          input_binary=True, 
                          input_checkpoint='/root/models/optimize_me/linear/cpu/model.ckpt',
                          output_node_names="add",
                          restore_op_name="save/restore_all", 
                          filename_tensor_name="save/Const:0",
                          output_graph=fully_optimized_frozen_model_graph_path, 
                          clear_devices=True, 
                          initializer_nodes="")
print(fully_optimized_frozen_model_graph_path)

"""
Explanation: Deploy Fully Optimized Model to TensorFlow Serving
Freeze Fully Optimized Graph
End of explanation
"""


%%bash

ls -l /root/models/optimize_me/linear/cpu/

"""
Explanation: File Size
End of explanation
"""


%%bash

summarize_graph --in_graph=/root/models/optimize_me/linear/cpu/fully_optimized_frozen_cpu.pb

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import re
from google.protobuf import text_format
from tensorflow.core.framework import graph_pb2

def convert_graph_to_dot(input_graph, output_dot, is_input_graph_binary):
    graph = graph_pb2.GraphDef()
    with open(input_graph, "rb") as fh:
        if is_input_graph_binary:
            graph.ParseFromString(fh.read())
        else:
            text_format.Merge(fh.read(), graph)
    with open(output_dot, "wt") as fh:
        print("digraph graphname {", file=fh)
        for node in graph.node:
            output_name = node.name
            print("  \"" + output_name + "\" [label=\"" + node.op + "\"];", file=fh)
            for input_full_name in node.input:
                parts = input_full_name.split(":")
                input_name = re.sub(r"^\^", "", parts[0])
                print("  \"" + input_name + "\" -> \"" + output_name + "\";", file=fh)
        print("}", file=fh)
        print("Created dot file '%s' for graph '%s'." % (output_dot, input_graph))
        

input_graph='/root/models/optimize_me/linear/cpu/fully_optimized_frozen_cpu.pb'
output_dot='/root/notebooks/fully_optimized_frozen_cpu.dot'
convert_graph_to_dot(input_graph=input_graph, output_dot=output_dot, is_input_graph_binary=True)

%%bash

dot -T png /root/notebooks/fully_optimized_frozen_cpu.dot \
    -o /root/notebooks/fully_optimized_frozen_cpu.png > /tmp/a.out

from IPython.display import Image

Image('/root/notebooks/fully_optimized_frozen_cpu.png')

"""
Explanation: Graph
End of explanation
"""


%%bash

benchmark_model --graph=/root/models/optimize_me/linear/cpu/fully_optimized_frozen_cpu.pb \
    --input_layer=weights,bias,x_observed \
    --input_layer_type=float,float,float \
    --input_layer_shape=:: \
    --output_layer=add

"""
Explanation: Run Standalone Benchmarks
Note:  These benchmarks are running against the standalone models on disk.  We will benchmark the models running within TensorFlow Serving soon.
End of explanation
"""


import tensorflow as tf

tf.reset_default_graph()

"""
Explanation: Save Model for Deployment and Inference
Reset Default Graph
End of explanation
"""


sess = tf.Session()

"""
Explanation: Create New Session
End of explanation
"""


from datetime import datetime 

version = int(datetime.now().strftime("%s"))

"""
Explanation: Generate Version Number
End of explanation
"""


%%bash

inspect_checkpoint --file_name=/root/models/optimize_me/linear/cpu/model.ckpt

saver = tf.train.import_meta_graph('/root/models/optimize_me/linear/cpu/model.ckpt.meta')
saver.restore(sess, '/root/models/optimize_me/linear/cpu/model.ckpt')

optimize_me_parent_path = '/root/models/optimize_me/linear/cpu'
fully_optimized_frozen_model_graph_path = '%s/fully_optimized_frozen_cpu.pb' % optimize_me_parent_path
print(fully_optimized_frozen_model_graph_path)

with tf.gfile.GFile(fully_optimized_frozen_model_graph_path, 'rb') as f:
    graph_def = tf.GraphDef()
    graph_def.ParseFromString(f.read())

tf.import_graph_def(
    graph_def, 
    input_map=None, 
    return_elements=None, 
    name="", 
    op_dict=None, 
    producer_op_list=None
)

print("weights = ", sess.run("weights:0"))
print("bias = ", sess.run("bias:0"))

"""
Explanation: Load Optimized, Frozen Graph
End of explanation
"""


from tensorflow.python.saved_model import utils
from tensorflow.python.saved_model import signature_constants
from tensorflow.python.saved_model import signature_def_utils

graph = tf.get_default_graph()

x_observed = graph.get_tensor_by_name('x_observed:0')
y_pred = graph.get_tensor_by_name('add:0')

tensor_info_x_observed = utils.build_tensor_info(x_observed)
print(tensor_info_x_observed)

tensor_info_y_pred = utils.build_tensor_info(y_pred)
print(tensor_info_y_pred)

prediction_signature = signature_def_utils.build_signature_def(inputs = 
                {'x_observed': tensor_info_x_observed}, 
                outputs = {'y_pred': tensor_info_y_pred}, 
                method_name = signature_constants.PREDICT_METHOD_NAME)

"""
Explanation: Create SignatureDef Asset for TensorFlow Serving
End of explanation
"""


from tensorflow.python.saved_model import builder as saved_model_builder
from tensorflow.python.saved_model import tag_constants

fully_optimized_saved_model_path = '/root/models/linear_fully_optimized/cpu/%s' % version
print(fully_optimized_saved_model_path)

builder = saved_model_builder.SavedModelBuilder(fully_optimized_saved_model_path)
builder.add_meta_graph_and_variables(sess, 
                                     [tag_constants.SERVING],
                                     signature_def_map={'predict':prediction_signature,                                     
signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:prediction_signature}, 
                                     clear_devices=True,
)

builder.save(as_text=False)

import os
print(fully_optimized_saved_model_path)
os.listdir(fully_optimized_saved_model_path)
os.listdir('%s/variables' % fully_optimized_saved_model_path)

sess.close()

"""
Explanation: Save Model with Assets
End of explanation
"""


