Commit 69b68b26 authored by Yuxin Wu's avatar Yuxin Wu

fix #967

parent fded4f40
...@@ -4,7 +4,7 @@ ...@@ -4,7 +4,7 @@
import numpy as np import numpy as np
from tensorpack.tfutils.varmanip import dump_chkpt_vars from tensorpack.tfutils.varmanip import load_chkpt_vars
from tensorpack.utils import logger from tensorpack.utils import logger
import argparse import argparse
...@@ -21,13 +21,13 @@ if __name__ == '__main__': ...@@ -21,13 +21,13 @@ if __name__ == '__main__':
elif args.model.endswith('.npz'): elif args.model.endswith('.npz'):
params = dict(np.load(args.model)) params = dict(np.load(args.model))
else: else:
params = dump_chkpt_vars(args.model) params = load_chkpt_vars(args.model)
logger.info("Variables in the model:") logger.info("Variables in the model:")
logger.info(str(params.keys())) logger.info(str(params.keys()))
if args.dump: if args.dump:
assert args.dump.endswith('.npz'), args.dump assert args.dump.endswith('.npz'), args.dump
np.save(args.dump, **params) np.savez(args.dump, **params)
if args.shell: if args.shell:
# params is a dict. play with it # params is a dict. play with it
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment