remove Identity block

Author: dskkato

examples/mnist/expected_values.txt

@@ -1 +1 @@
-1.8487935e-05, 5.188386e-08, 0.0003525585, 0.0014435184, 3.3050244e-07, 3.180701e-05, 2.4361995e-08, 0.9980293, 1.7698758e-05, 0.00010610428
+4.0609893e-06, 1.1228256e-07, 4.821958e-05, 0.00013738014, 4.2588137e-09, 2.3615916e-05, 1.6721445e-10, 0.99975353, 7.9972335e-07, 3.221138e-05

examples/mnist/mnist.py

@@ -58,9 +58,15 @@ def probability_model(x):
 assert_array_equal(expected, tf_model_output)
 assert_array_equal(expected, frozen_model_output[0])
 
+# Remove `Identity` operation
+# # TODO Check whether the last identity operation is actualy required
+frozen_model_graph_def = frozen_model.graph.as_graph_def()
+optimized_graph_def = tf.Graph().as_graph_def()
+optimized_graph_def.node.extend(frozen_model_graph_def.node[:-1])
+
 # write the computation graph and weights
 directory = "examples/mnist"
-tf.io.write_graph(frozen_model.graph, directory, "model.pb", as_text=False)
+tf.io.write_graph(optimized_graph_def, directory, "model.pb", as_text=False)
 
 # dump expected values to compare Rust's outputs
 with open("examples/mnist/expected_values.txt", "w") as f:
@@ -71,4 +77,4 @@ def probability_model(x):
 logdir = "logs/mnist"
 writer = tf.summary.create_file_writer(logdir)
 with writer.as_default():
-    tf.summary.graph(frozen_model.graph)
+    tf.summary.graph(optimized_graph_def)