def _fn(features, labels, params, mode, config):
"""model_fn implementation."""
local_model = copy.deepcopy(model)
if mode == tf.estimator.ModeKeys.TRAIN:
features_shards = dispatcher.shard(features)
labels_shards = dispatcher.shard(labels)
losses_shards = dispatcher(_loss_op, local_model, features_shards,
labels_shards, params, mode)
loss = _extract_loss(losses_shards)
train_op = local_model.optimize_loss(loss, params=params, hvd=hvd)
extra_variables = []
if isinstance(train_op, tuple):
train_op, extra_variables = train_op
training_hooks = []
if extra_variables:
training_hooks.append(
hooks.VariablesInitializerHook(extra_variables))
if config is not None:
local_model.examples_inputter.visualize(config.model_dir)
features_length = local_model.features_inputter.get_length(
features)
labels_length = (local_model.labels_inputter.get_length(labels)
if not model.unsupervised else None)
num_words = {}
if features_length is not None:
num_words["source"] = tf.reduce_sum(features_length)
if labels_length is not None:
num_words["target"] = tf.reduce_sum(labels_length)
training_hooks.append(
hooks.LogWordsPerSecondHook(
num_words,
every_n_steps=config.save_summary_steps,
output_dir=config.model_dir))
return tf.estimator.EstimatorSpec(mode,
loss=loss,
train_op=train_op,
training_hooks=training_hooks)
elif mode == tf.estimator.ModeKeys.EVAL:
logits, predictions = local_model(features, labels, params, mode)
loss = local_model.compute_loss(logits,
labels,
training=False,
params=params)
loss = _extract_loss(loss)
eval_metric_ops = local_model.compute_metrics(predictions, labels)
evaluation_hooks = []
if predictions is not None and eval_prediction_hooks_fn is not None:
evaluation_hooks.extend(eval_prediction_hooks_fn(predictions))
return tf.estimator.EstimatorSpec(
mode,
loss=loss,
eval_metric_ops=eval_metric_ops,
evaluation_hooks=evaluation_hooks)
elif mode == tf.estimator.ModeKeys.PREDICT:
_, predictions = local_model(features, labels, params, mode)
# Forward example index for reordering predictions.
if "index" in features:
predictions["index"] = features["index"]
export_outputs = {}
export_outputs[tf.saved_model.signature_constants.
DEFAULT_SERVING_SIGNATURE_DEF_KEY] = (
tf.estimator.export.PredictOutput(predictions))
return tf.estimator.EstimatorSpec(mode,
predictions=predictions,
export_outputs=export_outputs)
else:
raise ValueError("Invalid mode")
def _fn(features, labels, params, mode, config):
"""model_fn implementation."""
local_model_trans = copy.deepcopy(trans_model)
local_model_ae = copy.deepcopy(ae_model)
decode_two = True
if mode == tf.estimator.ModeKeys.TRAIN:
features_shards = dispatcher.shard(features)
labels_shards = dispatcher.shard(labels)
losses_shards = dispatcher(_loss_op, local_model_trans,
features_shards, labels_shards, params,
mode)
loss = _extract_loss(losses_shards)
train_op = local_model_trans.optimize_loss(loss,
params=params,
hvd=hvd)
extra_variables = []
if isinstance(train_op, tuple):
train_op, extra_variables = train_op
training_hooks = []
if extra_variables:
training_hooks.append(
hooks.VariablesInitializerHook(extra_variables))
if config is not None:
local_model_trans.examples_inputter.visualize(config.model_dir)
features_length = local_model_trans.features_inputter.get_length(
features)
labels_length = (
local_model_trans.labels_inputter.get_length(labels)
if not trans_model.unsupervised else None)
num_words = {}
if features_length is not None:
num_words["source"] = tf.reduce_sum(features_length)
if labels_length is not None:
num_words["target"] = tf.reduce_sum(labels_length)
training_hooks.append(
hooks.LogWordsPerSecondHook(
num_words,
every_n_steps=config.save_summary_steps,
output_dir=config.model_dir))
return tf.estimator.EstimatorSpec(mode,
loss=loss,
train_op=train_op,
training_hooks=training_hooks)
elif mode == tf.estimator.ModeKeys.EVAL:
logits, predictions = local_model_trans(features, labels, params,
mode)
loss = local_model_trans.compute_loss(logits,
labels,
training=False,
params=params)
loss = _extract_loss(loss)
eval_metric_ops = local_model_trans.compute_metrics(
predictions, labels)
evaluation_hooks = []
if predictions is not None and eval_prediction_hooks_fn is not None:
evaluation_hooks.extend(eval_prediction_hooks_fn(predictions))
return tf.estimator.EstimatorSpec(
mode,
loss=loss,
eval_metric_ops=eval_metric_ops,
evaluation_hooks=evaluation_hooks)
elif mode == tf.estimator.ModeKeys.PREDICT and not infering:
_, predictions = local_model_trans(features, labels, params, mode)
if "index" in features:
predictions["index"] = features["index"]
export_outputs = {}
export_outputs[tf.saved_model.signature_constants.
DEFAULT_SERVING_SIGNATURE_DEF_KEY] = (
tf.estimator.export.PredictOutput(predictions))
return tf.estimator.EstimatorSpec(mode,
predictions=predictions,
export_outputs=export_outputs)
elif mode == tf.estimator.ModeKeys.PREDICT:
print("==========================================")
#print (local_model.name)
print(json.dumps(params, indent=2))
#print (constants.END_OF_SENTENCE_ID)
#print (constants.END_OF_SENTENCE_TOKEN)
#ae_features = {
# "ids": features["ids_ae"],
# "length": features["length_ae"],
# "tokens": features["tokens_ae"]
# }
#ae_metadata = {
# "source_words_vocabulary":"../vocab/vocab.zh",
# "target_words_vocabulary":"../vocab/vocab.zh"
# }
#local_model = load_model("../model/AutoEncoder", "auto_encoder.py", model_name=None, serialize_model=True)
#local_model.initialize(ae_metadata)
#_, predictions = local_model(ae_features, labels, params, mode)
#if "index" in features:
# predictions["index"] = features["index"]
#export_outputs = {}
#export_outputs[tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY] = (
# tf.estimator.export.PredictOutput(predictions))
#return tf.estimator.EstimatorSpec(mode,predictions=predictions,export_outputs=export_outputs)
#local_model = local_model_ae
#local_model = local_model_trans
beam_size = params.get("beam_width", 1)
def get_info(local_model, features):
with tf.variable_scope(
local_model.name,
initializer=local_model._initializer(params)):
if not compat.reuse():
print("REUSING")
local_model._build()
inputs = local_model.features_inputter.make_inputs(
features)
length = local_model.features_inputter.get_length(features)
#inputs = local_model.features_inputter.make_inputs({"ids":features["ids"]})
#length = local_model.features_inputter.get_length({"length":features["length"]})
#_, predictions = local_model._call(features, labels, params, mode)
with tf.variable_scope("encoder"):
encoder_outputs, encoder_state, encoder_sequence_length = local_model.encoder.encode(
inputs, sequence_length=length, mode=mode)
with tf.variable_scope("decoder"):
batch_size = tf.shape(
tf.contrib.framework.nest.flatten(encoder_outputs)
[0])[0]
initial_state = tf.contrib.seq2seq.tile_batch(
encoder_state, multiplier=beam_size)
memory = tf.contrib.seq2seq.tile_batch(
encoder_outputs, multiplier=beam_size)
memory_sequence_length = tf.contrib.seq2seq.tile_batch(
encoder_sequence_length, multiplier=beam_size)
output_layer = local_model.output_layer
if output_layer == None:
output_layer = tf.layers.Dense(
local_model.labels_inputter.vocabulary_size,
use_bias=True,
dtype=local_model.labels_inputter.dtype)
output_layer.build(
[None, local_model.decoder.output_size])
embedding_fn = lambda ids: tf.nn.embedding_lookup(
local_model.labels_inputter.embedding, ids)
step_fn, initial_state = local_model.decoder.step_fn(
tf.estimator.ModeKeys.PREDICT,
batch_size * beam_size,
initial_state=initial_state,
memory=memory,
memory_sequence_length=memory_sequence_length,
dtype=local_model.labels_inputter.dtype)
def symbols_to_logits_fn(ids, step, state):
inputs = embedding_fn(ids)
returned_values = step_fn(
step, inputs, state,
tf.estimator.ModeKeys.PREDICT)
outputs, state, attention = returned_values
logits = output_layer(outputs)
return logits, state, attention
return symbols_to_logits_fn, initial_state, batch_size
trans_features = {
"ids": features["ids"],
"length": features["length"]
}
ae_features = {
"ids": features["ids_ae"],
"length": features["length_ae"]
}
trans_symbols_to_logits_fn, trans_initial_state, batch_size = get_info(
local_model_trans, trans_features)
ae_symbols_to_logits_fn, ae_initial_state, _ = get_info(
local_model_ae, ae_features)
start_tokens = tf.fill([batch_size],
constants.START_OF_SENTENCE_ID)
decoding_strategy = decoding_infer.BeamSearch(beam_size,
length_penalty=0,
coverage_penalty=0)
sampler = decoding.BestSampler()
#sampler = decoding.RandomSampler(from_top_k=1000)
low_rate = 0.05
ids_ori = tf.contrib.seq2seq.tile_batch(features["ids_ori"],
beam_size)
pass_spe_token = tf.where(ids_ori > 2)
appear_pos = tf.stack([
pass_spe_token[:, 0],
tf.cast(tf.gather_nd(ids_ori, pass_spe_token), tf.int64)
],
axis=1)
low_prob = tf.sparse_to_dense(appear_pos, [
batch_size * beam_size,
local_model_trans.labels_inputter.vocabulary_size
], low_rate, 0, False)
print(
"==============================================================="
)
print(low_prob)
low_prob = None
sampled_ids, sampled_length = decoding_infer.dynamic_decode(
trans_symbols_to_logits_fn,
ae_symbols_to_logits_fn,
local_model_trans.name,
local_model_ae.name,
start_tokens,
end_id=constants.END_OF_SENTENCE_ID,
initial_state_trans=trans_initial_state,
initial_state_ae=ae_initial_state,
decoding_strategy=decoding_strategy,
sampler=sampler,
maximum_iterations=150,
minimum_iterations=0,
attention_history=False,
attention_size=None,
low_prob=low_prob)
sampled_length = tf.minimum(sampled_length + 1,
tf.shape(sampled_ids)[2])
target_vocab_rev = local_model_trans.labels_inputter.vocabulary_lookup_reverse(
)
target_tokens = target_vocab_rev.lookup(
tf.cast(sampled_ids, tf.int64))
predictions = {
"tokens": target_tokens,
"length": sampled_length,
}
num_hypotheses = params.get("num_hypotheses", 1)
if num_hypotheses > 0:
if num_hypotheses > beam_size:
raise ValueError(
"n_best cannot be greater than beam_width")
for key, value in six.iteritems(predictions):
predictions[key] = value[:, :num_hypotheses]
if ae_model != None:
ae_init_checkpoint = params["ae_init_checkpoint"]
ae_variables_map = get_pretrained_variables_map(
ae_init_checkpoint)
tf.contrib.framework.init_from_checkpoint(
ae_init_checkpoint, ae_variables_map)
trans_init_checkpoint = params["trans_init_checkpoint"]
trans_variables_map = get_pretrained_variables_map(
trans_init_checkpoint)
tf.contrib.framework.init_from_checkpoint(trans_init_checkpoint,
trans_variables_map)
if "index" in features:
predictions["index"] = features["index"]
export_outputs = {}
export_outputs[tf.saved_model.signature_constants.
DEFAULT_SERVING_SIGNATURE_DEF_KEY] = (
tf.estimator.export.PredictOutput(predictions))
return tf.estimator.EstimatorSpec(mode,
predictions=predictions,
export_outputs=export_outputs)
else:
raise ValueError("Invalid mode")
def _model_fn(features, labels, params, mode, config):
"""model_fn implementation."""
if mode == tf.estimator.ModeKeys.TRAIN:
counters = self._register_word_counters(features, labels)
training_hooks = []
if config is not None:
training_hooks.append(
hooks.CountersHook(
every_n_steps=config.save_summary_steps,
output_dir=config.model_dir,
counters=counters))
features_shards = dispatcher.shard(features)
labels_shards = dispatcher.shard(labels)
with tf.variable_scope(self.name,
initializer=self._initializer(params)):
losses_shards = dispatcher(_loss_op, features_shards,
labels_shards, params, mode,
config)
loss = _extract_loss(losses_shards)
train_op, extra_variables = optimize_loss(
loss, params, mixed_precision=(self.dtype == tf.float16))
if extra_variables:
training_hooks.append(
hooks.VariablesInitializerHook(extra_variables))
return tf.estimator.EstimatorSpec(
mode,
loss=loss,
train_op=train_op,
training_hooks=training_hooks)
elif mode == tf.estimator.ModeKeys.EVAL:
with tf.variable_scope(self.name):
logits, predictions = self._build(features,
labels,
params,
mode,
config=config)
loss = self._compute_loss(features, labels, logits, params,
mode)
loss = _extract_loss(loss)
eval_metric_ops = self._compute_metrics(
features, labels, predictions) # pylint: disable=assignment-from-none
evaluation_hooks = []
if predictions is not None and eval_prediction_hooks_fn is not None:
evaluation_hooks.extend(
eval_prediction_hooks_fn(predictions))
return tf.estimator.EstimatorSpec(
mode,
loss=loss,
eval_metric_ops=eval_metric_ops,
evaluation_hooks=evaluation_hooks)
elif mode == tf.estimator.ModeKeys.PREDICT:
with tf.variable_scope(self.name):
_, predictions = self._build(features,
labels,
params,
mode,
config=config)
# Forward example index for reordering predictions.
if "index" in features:
predictions["index"] = features["index"]
export_outputs = {}
export_outputs[tf.saved_model.signature_constants.
DEFAULT_SERVING_SIGNATURE_DEF_KEY] = (
tf.estimator.export.PredictOutput(
predictions))
return tf.estimator.EstimatorSpec(
mode,
predictions=predictions,
export_outputs=export_outputs)
else:
raise RuntimeError("Invalid mode")
def _model_fn(features, labels, params, mode, config):
"""model_fn implementation."""
if mode == tf.estimator.ModeKeys.TRAIN:
counters = self._register_word_counters(features, labels)
training_hooks = []
if config is not None:
training_hooks.append(hooks.CountersHook(
every_n_steps=config.save_summary_steps,
output_dir=config.model_dir,
counters=counters))
features_shards = dispatcher.shard(features)
labels_shards = dispatcher.shard(labels)
with tf.variable_scope(self.name, initializer=self._initializer(params)):
losses_shards = dispatcher(
_loss_op, features_shards, labels_shards, params, mode, config)
loss = _extract_loss(losses_shards)
#train_op, extra_variables = optimize_loss(
# loss, params, mixed_precision=(self.dtype == tf.float16))
#TODO: freeze_update, get config, get var_list, pass var_list, move >>> DIFFERENT VARLIST FOR LSTM AND TRANSFORMER, THUS, BEFORE SLECTING WHAT VARS, CHECK IF IT EXISTS IN CKPT >> MAKE IT ARCHI DEPENDENT
freeze_params = params.get("freeze")
if False: #freeze_params is not None:
tf.logging.info("Optimizing selected network components:", freeze_params)
#json.dumps(freeze_params, indent=2, sort_keys=True))
var_list = self._get_variables(freeze_params)
train_op, extra_variables = optimize_loss(
loss, params, mixed_precision=(self.dtype == tf.float16), var_list=var_list)
else:
train_op, extra_variables = optimize_loss(
loss, params, mixed_precision=(self.dtype == tf.float16))
#TODO: simpler way of doing it, if freeze is None, pass var_list=None, else select vars to pass.
#train_op, extra_variables = optimize_loss(
# loss, params, mixed_precision=(self.dtype == tf.float16), var_list=self._get_variables(params))
if extra_variables:
training_hooks.append(hooks.VariablesInitializerHook(extra_variables))
return tf.estimator.EstimatorSpec(
mode,
loss=loss,
train_op=train_op,
training_hooks=training_hooks)
elif mode == tf.estimator.ModeKeys.EVAL:
with tf.variable_scope(self.name):
logits, predictions = self._build(features, labels, params, mode, config=config)
loss = self._compute_loss(features, labels, logits, params, mode)
loss = _extract_loss\
(loss)
eval_metric_ops = self._compute_metrics(features, labels, predictions) # pylint: disable=assignment-from-none
evaluation_hooks = []
if predictions is not None and eval_prediction_hooks_fn is not None:
evaluation_hooks.extend(eval_prediction_hooks_fn(predictions))
return tf.estimator.EstimatorSpec(
mode,
loss=loss,
eval_metric_ops=eval_metric_ops,
evaluation_hooks=evaluation_hooks)
elif mode == tf.estimator.ModeKeys.PREDICT:
with tf.variable_scope(self.name):
_, predictions = self._build(features, labels, params, mode, config=config)
# Forward example index for reordering predictions.
if "index" in features:
predictions["index"] = features["index"]
export_outputs = {}
export_outputs[tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY] = (
tf.estimator.export.PredictOutput(predictions))
return tf.estimator.EstimatorSpec(
mode,
predictions=predictions,
export_outputs=export_outputs)
else:
raise RuntimeError("Invalid mode")