Commit 1fa5ba1c authored by Yuxin Wu's avatar Yuxin Wu

don't use --config in dump-model-params

parent 3f52e871
...@@ -5,31 +5,18 @@ ...@@ -5,31 +5,18 @@
import argparse import argparse
import tensorflow as tf import tensorflow as tf
import imp
from tensorpack import TowerContext, logger, PlaceholderInput from tensorpack import logger
from tensorpack.tfutils import varmanip, get_model_loader from tensorpack.tfutils import varmanip, get_model_loader
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
parser.add_argument('--config', help='config file') parser.add_argument('--meta', help='metagraph file', required=True)
parser.add_argument('--meta', help='metagraph file')
parser.add_argument(dest='model') parser.add_argument(dest='model')
parser.add_argument(dest='output') parser.add_argument(dest='output')
args = parser.parse_args() args = parser.parse_args()
assert args.config or args.meta, "Either config or metagraph must be present!"
with tf.Graph().as_default() as G: with tf.Graph().as_default() as G:
if args.config: tf.train.import_meta_graph(args.meta)
logger.warn("Using a config script is not reliable. Please use metagraph.")
MODEL = imp.load_source('config_script', args.config).Model
M = MODEL()
with TowerContext('', is_training=False):
input = PlaceholderInput()
input.setup(M.get_inputs_desc())
M.build_graph(*input.get_input_tensors())
else:
tf.train.import_meta_graph(args.meta)
# loading... # loading...
init = get_model_loader(args.model) init = get_model_loader(args.model)
......
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