def load_pretrain_model(model_name, pretrain_model_dir, problem_name):
""" Loads pretrained model.
Args:
model_name: The name of the model.
pretrain_model_dir: The pretrained model dir.
problem_name: The problem name.
Returns:
A list of assign ops.
"""
top_scope_name = get_model_top_scope_name(model_name, problem_name)
pt_model_configs = ModelConfigs.load(pretrain_model_dir)
pt_model_top_scope_name = get_model_top_scope_name(
pt_model_configs["model"], pt_model_configs["problem_name"])
tf.logging.info("loading variables from {}".format(pretrain_model_dir))
assign_op = []
for var_name, _ in tf.contrib.framework.list_variables(pretrain_model_dir):
if var_name.startswith("OptimizeLoss"):
continue
if tf.GraphKeys.GLOBAL_STEP in var_name or "learning_rate" in var_name or "lr" in var_name:
tf.logging.info("Pretrain: ignore {}".format(var_name))
continue
tf.logging.info("Pretrain: reload {}".format(var_name))
var = tf.contrib.framework.load_variable(pretrain_model_dir, var_name)
with tf.variable_scope(top_scope_name, reuse=True):
v = tf.get_variable(name=var_name[len(pt_model_top_scope_name) +
1:],
shape=var.shape,
dtype=var.dtype)
assign_op.append(v.assign(var))
return assign_op
def load_pretrain_model(model_name, pretrain_model_dir, problem_name):
""" Loads pretrained model.
Args:
model_name: The name of the model.
pretrain_model_dir: The pretrained model dir.
problem_name: The problem name.
Returns:
A list of assign ops.
"""
top_scope_name = get_model_top_scope_name(model_name, problem_name)
pt_model_configs = ModelConfigs.load(pretrain_model_dir)
pt_model_top_scope_name = get_model_top_scope_name(pt_model_configs["model"], pt_model_configs["problem_name"])
tf.logging.info("loading variables from {}".format(pretrain_model_dir))
assign_op = []
for var_name, _ in tf.contrib.framework.list_variables(pretrain_model_dir):
if var_name.startswith("OptimizeLoss"):
continue
if tf.GraphKeys.GLOBAL_STEP in var_name or "learning_rate" in var_name or "lr" in var_name:
tf.logging.info("Pretrain: ignore {}".format(var_name))
continue
tf.logging.info("Pretrain: reload {}".format(var_name))
var = tf.contrib.framework.load_variable(pretrain_model_dir, var_name)
with tf.variable_scope(top_scope_name, reuse=True):
v = tf.get_variable(name=var_name[len(pt_model_top_scope_name) + 1:],
shape=var.shape, dtype=var.dtype)
assign_op.append(v.assign(var))
return assign_op
def model_fn_ensemble(model_dirs,
dataset,
weight_scheme,
inference_options,
verbose=True):
""" Reloads NMT models from checkpoints and builds the ensemble
model inference.
Args:
model_dirs: A list of model directories (checkpoints).
dataset: A `Dataset` object.
weight_scheme: A string, the ensemble weights. See
`EnsembleModel.get_ensemble_weights()` for more details.
inference_options: Contains beam_size, length_penalty and
maximum_labels_length.
verbose: Print logging info if set True.
Returns: A `EstimatorSpec` object.
"""
# load variable, rename (add prefix to varname), build model
models = []
for index, model_dir in enumerate(model_dirs):
if verbose:
tf.logging.info("loading variables from {}".format(model_dir))
# load variables
model_name = None
for var_name, _ in tf.contrib.framework.list_variables(model_dir):
if var_name.startswith("OptimizeLoss"):
continue
if model_name is None:
model_name = _inspect_varname_prefix(var_name)
var = tf.contrib.framework.load_variable(model_dir, var_name)
with tf.variable_scope(GlobalNames.ENSEMBLE_VARNAME_PREFIX +
str(index)):
var = tf.get_variable(name=var_name,
shape=var.shape,
dtype=tf.float32,
initializer=tf.constant_initializer(var))
# load model configs
assert model_name, ("Fail to fetch model name")
model_configs = ModelConfigs.load(model_dir)
if verbose:
tf.logging.info("Create model: {}.".format(model_configs["model"]))
model = eval(model_configs["model"])(
params=model_configs["model_params"],
mode=tf.contrib.learn.ModeKeys.INFER,
vocab_source=dataset.vocab_source,
vocab_target=dataset.vocab_target,
name=model_name,
verbose=False)
models.append(model)
ensemble_model = EnsembleModel(weight_scheme=weight_scheme,
inference_options=inference_options)
with tf.variable_scope("", reuse=True):
predictions = ensemble_model.build(input_fields=dataset.input_fields,
base_models=models,
vocab_target=dataset.vocab_target)
return EstimatorSpec(tf.contrib.learn.ModeKeys.INFER,
predictions=predictions)
def __init__(self, model_configs):
""" Initializes the training experiment.
Args:
model_configs: A dictionary of all configurations.
"""
super(TrainingExperiment, self).__init__()
# training options
training_options = parse_params(
params=model_configs["train"],
default_params=self.default_training_options())
# for datasets
datasets_params = parse_params(
params=model_configs["data"],
default_params=self.default_datasets_params())
self._model_configs = model_configs
self._model_configs["train"] = training_options
self._model_configs["data"] = datasets_params
print_params("Datasets: ", self._model_configs["data"])
print_params("Training parameters: ", self._model_configs["train"])
ModelConfigs.dump(self._model_configs, self._model_configs["model_dir"])
def main(_argv):
model_configs = maybe_load_yaml(DEFAULT_EVAL_CONFIGS)
# load flags from config file
model_configs = load_from_config_path(FLAGS.config_paths, model_configs)
# replace parameters in configs_file with tf FLAGS
model_configs = update_eval_model_configs(model_configs, FLAGS)
model_configs = deep_merge_dict(model_configs, ModelConfigs.load(FLAGS.model_dir))
model_configs = update_eval_model_configs(model_configs, FLAGS)
runner = EvalExperiment(model_configs=model_configs)
runner.run()
def __init__(self, model_configs):
""" Initializes the training experiment.
Args:
model_configs: A dictionary of all configurations.
"""
super(TrainingExperiment, self).__init__()
# training options
training_options = parse_params(
params=model_configs["train"],
default_params=self.default_training_options())
# for datasets
datasets_params = parse_params(
params=model_configs["data"],
default_params=self.default_datasets_params())
self._model_configs = model_configs
self._model_configs["train"] = training_options
self._model_configs["data"] = datasets_params
print_params("Datasets: ", self._model_configs["data"])
print_params("Training parameters: ", self._model_configs["train"])
ModelConfigs.dump(self._model_configs, self._model_configs["model_dir"])
def main(_argv):
model_configs = maybe_load_yaml(DEFAULT_EVAL_CONFIGS)
# load flags from config file
model_configs = load_from_config_path(FLAGS.config_paths, model_configs)
# replace parameters in configs_file with tf FLAGS
model_configs = update_eval_model_configs(model_configs, FLAGS)
model_configs = deep_merge_dict(model_configs,
ModelConfigs.load(FLAGS.model_dir))
model_configs = update_eval_model_configs(model_configs, FLAGS)
runner = EvalExperiment(model_configs=model_configs)
runner.run()
def main(_argv):
# load flags from config file
model_configs = load_from_config_path(FLAGS.config_paths)
# replace parameters in configs_file with tf FLAGS
model_configs = update_configs_from_flags(model_configs, FLAGS,
EVAL_ARGS.keys())
model_configs = deep_merge_dict(model_configs,
ModelConfigs.load(FLAGS.model_dir))
model_configs = update_configs_from_flags(model_configs, FLAGS,
EVAL_ARGS.keys())
runner = EvalExperiment(model_configs=model_configs)
runner.run()
def main(_argv):
model_configs = maybe_load_yaml(DEFAULT_INFER_CONFIGS)
# load flags from config file
model_configs = load_from_config_path(FLAGS.config_paths, model_configs)
# replace parameters in configs_file with tf FLAGS
model_configs = update_infer_model_configs(model_configs, FLAGS)
model_dirs = FLAGS.model_dir.strip().split(",")
if len(model_dirs) == 1:
model_configs = deep_merge_dict(model_configs, ModelConfigs.load(model_dirs[0]))
model_configs = update_infer_model_configs(model_configs, FLAGS)
runner = InferExperiment(model_configs=model_configs)
else:
runner = EnsembleExperiment(model_configs=model_configs, model_dirs=model_dirs,
weight_scheme=FLAGS.weight_scheme)
runner.run()
def main(_argv):
model_configs = maybe_load_yaml(DEFAULT_INFER_CONFIGS)
# load flags from config file
model_configs = load_from_config_path(FLAGS.config_paths, model_configs)
# replace parameters in configs_file with tf FLAGS
model_configs = update_infer_model_configs(model_configs, FLAGS)
model_dirs = FLAGS.model_dir.strip().split(",")
if len(model_dirs) == 1:
model_configs = deep_merge_dict(model_configs,
ModelConfigs.load(model_dirs[0]))
runner = InferExperiment(model_configs=model_configs)
else:
runner = EnsembleExperiment(model_configs=model_configs,
model_dirs=model_dirs,
weight_scheme=FLAGS.weight_scheme)
runner.run()
def model_fn_ensemble_bak(
model_dirs,
dataset,
weight_scheme,
inference_options,
verbose=True):
""" Reloads NMT models from checkpoints and builds the ensemble
model inference.
Args:
model_dirs: A list of model directories (checkpoints).
dataset: A `Dataset` object.
weight_scheme: A string, the ensemble weights. See
`EnsembleModel.get_ensemble_weights()` for more details.
inference_options: Contains beam_size, length_penalty and
maximum_labels_length.
verbose: Print logging info if set True.
Returns: A `EstimatorSpec` object.
"""
# load variable, rename (add prefix to varname), build model
models = []
input_fields = None
for index, model_dir in enumerate(model_dirs):
if verbose:
tf.logging.info("loading variables from {}".format(model_dir))
# load variables
model_name = None
ensemble_scope_prefix = None
for var_name, _ in tf.contrib.framework.list_variables(model_dir):
if var_name.startswith("OptimizeLoss"):
continue
if model_name is None:
model_name = inspect_varname_prefix(var_name)
var = tf.contrib.framework.load_variable(model_dir, var_name)
with tf.variable_scope(Constants.ENSEMBLE_VARNAME_PREFIX + str(index)):
if ensemble_scope_prefix is None:
ensemble_scope_prefix = tf.get_variable_scope().name
var = tf.get_variable(
name=var_name, shape=var.shape, dtype=tf.float32,
initializer=tf.constant_initializer(var))
# load model configs
assert model_name, (
"Fail to fetch model name")
model_configs = ModelConfigs.load(model_dir)
if verbose:
tf.logging.info("Create model: {}.".format(
model_configs["model"]))
model = eval(model_configs["model"])(
params=model_configs["model_params"],
mode=ModeKeys.INFER,
vocab_source=dataset.vocab_source,
vocab_target=dataset.vocab_target,
name=os.path.join(ensemble_scope_prefix, model_name),
verbose=False)
models.append(model)
if input_fields is None:
input_fields = eval(model_configs["model"]).create_input_fields(ModeKeys.INFER)
ensemble_model = EnsembleModel(
vocab_target=dataset.vocab_target,
base_models=models,
weight_scheme=weight_scheme,
inference_options=inference_options)
with tf.variable_scope("", reuse=True):
predictions = ensemble_model.build(input_fields=input_fields)
return EstimatorSpec(
ModeKeys.INFER,
input_fields=input_fields,
predictions=predictions)