diff options
Diffstat (limited to 'cli')
| -rw-r--r-- | cli/app/utils/tf_utils.py | 8 |
1 files changed, 4 insertions, 4 deletions
diff --git a/cli/app/utils/tf_utils.py b/cli/app/utils/tf_utils.py index 0453644..9c12dfc 100644 --- a/cli/app/utils/tf_utils.py +++ b/cli/app/utils/tf_utils.py @@ -1,9 +1,9 @@ -from tensorflow.train import py_checkpoint_reader +from tensorflow.train import NewCheckpointReader 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() + 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)) |