tf utils to read vectors

1 files changed, 10 insertions, 0 deletions

diff --git a/cli/app/utils/tf_utils.py b/cli/app/utils/tf_utils.py
new file mode 100644
index 0000000..69a38e4
--- /dev/null
+++ b/cli/app/utils/tf_utils.py
@@ -0,0 +1,10 @@
+from tensorflow.python.training import py_checkpoint_reader
+
+def read_checkpoint(fp_ckpt, key):
+  reader = py_checkpoint_reader.NewCheckpointReader(fp_ckpt)
+  var_to_shape_map = reader.get_variable_to_shape_map()
+  var_to_dtype_map = reader.get_variable_to_dtype_map()
+  # for key, value in sorted(var_to_shape_map.items()):
+  #   print("tensor: %s (%s) %s" % (key, var_to_dtype_map[key].name, value))
+  #   print(reader.get_tensor(key))
+  return reader.get_tensor(key)