COMPMID-1839: Add script from extracting tf frozen models.

Change-Id: I9a61b9005ea829cd9ecae5bebf8985fe72e28b8e
Reviewed-on: https://review.mlplatform.org/448
Reviewed-by: Manuel Bottini <manuel.bottini@arm.com>
Tested-by: Arm Jenkins <bsgcomp@arm.com>
Reviewed-by: Michele Di Giorgio <michele.digiorgio@arm.com>

1 files changed, 62 insertions, 0 deletions

diff --git a/scripts/tf_frozen_model_extractor.py b/scripts/tf_frozen_model_extractor.py
new file mode 100644
index 0000000000..fbd5b38019
--- /dev/null
+++ b/scripts/tf_frozen_model_extractor.py
@@ -0,0 +1,62 @@
+#!/usr/bin/env python
+""" Extract trainable parameters from a frozen model and stores them in numpy arrays.
+Usage:
+    python tf_frozen_model_extractor -m path_to_frozem_model -d path_to_store_the_parameters
+
+Saves each variable to a {variable_name}.npy binary file.
+
+Note that the script permutes the trainable parameters to NCHW format. This is a pretty manual step thus it's not thoroughly tested.
+"""
+import argparse
+import os
+import numpy as np
+import tensorflow as tf
+from tensorflow.python.platform import gfile
+
+strings_to_remove=["read", "/:0"]
+permutations = { 1 : [0], 2 : [1, 0], 3 : [2, 1, 0], 4 : [3, 2, 0, 1]}
+
+if __name__ == "__main__":
+    # Parse arguments
+    parser = argparse.ArgumentParser('Extract TensorFlow net parameters')
+    parser.add_argument('-m', dest='modelFile', type=str, required=True, help='Path to TensorFlow frozen graph file (.pb)')
+    parser.add_argument('-d', dest='dumpPath', type=str, required=False, default='./', help='Path to store the resulting files.')
+    parser.add_argument('--nostore', dest='storeRes', action='store_false', help='Specify if files should not be stored. Used for debugging.')
+    parser.set_defaults(storeRes=True)
+    args = parser.parse_args()
+
+    # Create directory if not present
+    if not os.path.exists(args.dumpPath):
+        os.makedirs(args.dumpPath)
+
+    # Extract parameters
+    with tf.Graph().as_default() as graph:
+        with tf.Session() as sess:
+            print("Loading model.")
+            with gfile.FastGFile(args.modelFile, 'rb') as f:
+                graph_def = tf.GraphDef()
+                graph_def.ParseFromString(f.read())
+                sess.graph.as_default()
+
+                tf.import_graph_def(graph_def, input_map=None, return_elements=None, name="", op_dict=None, producer_op_list=None)
+
+                for op in graph.get_operations():
+                    for op_val in op.values():
+                        varname = op_val.name
+
+                        # Skip non-const values
+                        if "read" in varname:
+                            t  = op_val.eval()
+                            tT = t.transpose(permutations[len(t.shape)])
+                            t  = np.ascontiguousarray(tT)
+
+                            for s in strings_to_remove:
+                                varname = varname.replace(s, "")
+                            if os.path.sep in varname:
+                                varname = varname.replace(os.path.sep, '_')
+                                print("Renaming variable {0} to {1}".format(op_val.name, varname))
+
+                            # Store files
+                            if args.storeRes:
+                                print("Saving variable {0} with shape {1} ...".format(varname, t.shape))
+                                np.save(os.path.join(args.dumpPath, varname), t)