diff options
Diffstat (limited to 'cli/app/utils/tf_utils.py')
| -rw-r--r-- | cli/app/utils/tf_utils.py | 10 |
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) |