def model_fn(features, labels, mode, params):
"""Defines how to train, evaluate and predict from the transformer model."""
with tf.variable_scope("model"):
inputs, targets = features, labels
# Create model and get output logits.
model = transformer.Transformer(params,
mode == tf.estimator.ModeKeys.TRAIN)
logits = model(inputs, targets)
# When in prediction mode, the labels/targets is None. The model output
# is the prediction
if mode == tf.estimator.ModeKeys.PREDICT:
return tf.estimator.EstimatorSpec(
tf.estimator.ModeKeys.PREDICT,
predictions=logits,
export_outputs={
"translate": tf.estimator.export.PredictOutput(logits)
})
# Explicitly set the shape of the logits for XLA (TPU). This is needed
# because the logits are passed back to the host VM CPU for metric
# evaluation, and the shape of [?, ?, vocab_size] is too vague. However
# it is known from Transformer that the first two dimensions of logits
# are the dimensions of targets. Note that the ambiguous shape of logits is
# not a problem when computing xentropy, because padded_cross_entropy_loss
# resolves the shape on the TPU.
logits.set_shape(targets.shape.as_list() + logits.shape.as_list()[2:])
# Calculate model loss.
# xentropy contains the cross entropy loss of every nonpadding token in the
# targets.
xentropy, weights = metrics.padded_cross_entropy_loss(
logits, targets, params["label_smoothing"],
params["targets_vocab_size"])
loss = tf.reduce_sum(xentropy) / tf.reduce_sum(weights)
# Save loss as named tensor that will be logged with the logging hook.
tf.identity(loss, "cross_entropy")
if mode == tf.estimator.ModeKeys.EVAL:
return tf.estimator.EstimatorSpec(
mode=mode,
loss=loss,
predictions={"predictions": logits},
eval_metric_ops=metrics.get_eval_metrics(
logits, labels, params))
else:
train_op, metric_dict = get_train_op_and_metrics(loss, params)
# Epochs can be quite long. This gives some intermediate information
# in TensorBoard.
metric_dict["minibatch_loss"] = loss
record_scalars(metric_dict)
return tf.estimator.EstimatorSpec(mode=mode,
loss=loss,
train_op=train_op)
def model_fn(features, labels, mode, params):
"""Defines how to train, evaluate and predict from the transformer model.
labels: liste de 4 tensors shape [batch_size, seq_len, 1], un tenseur pour chaque ligne
"""
with tf.variable_scope("model"):
inputs, targets = features, labels
# Create model and get output logits.
conv = cnn_layer.CNNNetwork(4, 2, [5, 5], 0.4,
mode == tf.estimator.ModeKeys.TRAIN)
model = transformer.Transformer(params,
mode == tf.estimator.ModeKeys.TRAIN)
conv_output = conv(inputs)
log = [
model(conv_output[0], targets["line1"]),
model(conv_output[1], targets["line2"]),
model(conv_output[2], targets["line3"]),
model(conv_output[3], targets["line4"])
]
# When in prediction mode, the labels/targets is None. The model output
# is the prediction
if mode == tf.estimator.ModeKeys.PREDICT:
if params["use_tpu"]:
raise NotImplementedError(
"Prediction is not yet supported on TPUs.")
return tf.estimator.EstimatorSpec(
tf.estimator.ModeKeys.PREDICT,
predictions=log,
# export_outputs={
# "translate": tf.estimator.export.PredictOutput(logits)
# })
)
# Explicitly set the shape of the logits for XLA (TPU). This is needed
# because the logits are passed back to the host VM CPU for metric
# evaluation, and the shape of [?, ?, vocab_size] is too vague. However
# it is known from Transformer that the first two dimensions of logits
# are the dimensions of targets. Note that the ambiguous shape of logits is
# not a problem when computing xentropy, because padded_cross_entropy_loss
# resolves the shape on the TPU.
for line, logits in enumerate(log):
logits.set_shape(targets.shape.as_list() +
logits.shape.as_list()[2:])
# Calculate model loss.
# xentropy contains the cross entropy loss of every nonpadding token in the
# targets.
xentropy, weights = metrics.padded_cross_entropy_loss(
logits, targets, params["label_smoothing"],
params["vocab_size"])
loss = tf.reduce_sum(xentropy) / tf.reduce_sum(weights)
if mode == tf.estimator.ModeKeys.EVAL:
if params["use_tpu"]:
# host call functions should only have tensors as arguments.
# This lambda pre-populates params so that metric_fn is
# TPUEstimator compliant.
raise NotImplementedError(
"Prediction is not yet supported on TPUs.")
return tf.estimator.EstimatorSpec(
mode=mode,
loss=loss,
predictions={"predictions": logits},
eval_metric_ops=metrics.get_eval_metrics(
logits,
list(labels.keys())[line], params))
else:
train_op, metric_dict = get_train_op_and_metrics(loss, params)
record_scalars(metric_dict)
return tf.estimator.EstimatorSpec(mode=mode,
loss=loss,
train_op=train_op)
def model_fn(features, labels, mode, params):
"""Defines how to train, evaluate and predict from the transformer model."""
with tf.variable_scope("model"):
inputs = features[0]
segments = features[1]
masks = features[2]
targets = labels
# Create model and get output logits.
model = transformer.Transformer(params,
mode == tf.estimator.ModeKeys.TRAIN)
logits = model(inputs, segments, masks, targets)
# When in prediction mode, the labels/targets is None. The model output
# is the prediction
if mode == tf.estimator.ModeKeys.PREDICT:
if params["use_tpu"]:
raise NotImplementedError(
"Prediction is not yet supported on TPUs.")
return tf.estimator.EstimatorSpec(
tf.estimator.ModeKeys.PREDICT,
predictions=logits,
export_outputs={
"translate": tf.estimator.export.PredictOutput(logits)
})
# Explicitly set the shape of the logits for XLA (TPU). This is needed
# because the logits are passed back to the host VM CPU for metric
# evaluation, and the shape of [?, ?, vocab_size] is too vague. However
# it is known from Transformer that the first two dimensions of logits
# are the dimensions of targets. Note that the ambiguous shape of logits is
# not a problem when computing xentropy, because padded_cross_entropy_loss
# resolves the shape on the TPU.
logits.set_shape(targets.shape.as_list() + logits.shape.as_list()[2:])
# Calculate model loss.
# xentropy contains the cross entropy loss of every nonpadding token in the
# targets.
xentropy, weights = metrics.padded_cross_entropy_loss(
logits, targets, params["label_smoothing"], params["vocab_size"])
loss = tf.reduce_sum(xentropy) / tf.reduce_sum(weights)
# Save loss as named tensor that will be logged with the logging hook.
tf.identity(loss, "cross_entropy")
if mode == tf.estimator.ModeKeys.EVAL:
if params["use_tpu"]:
# host call functions should only have tensors as arguments.
# This lambda pre-populates params so that metric_fn is
# TPUEstimator compliant.
metric_fn = lambda logits, labels: (metrics.get_eval_metrics(
logits, labels, params=params))
eval_metrics = (metric_fn, [logits, targets])
return tf.contrib.tpu.TPUEstimatorSpec(
mode=mode,
loss=loss,
predictions={"predictions": logits},
eval_metrics=eval_metrics)
return tf.estimator.EstimatorSpec(
mode=mode,
loss=loss,
predictions={"predictions": logits},
eval_metric_ops=metrics.get_eval_metrics(
logits, targets, params))
else:
train_op, metric_dict = get_train_op_and_metrics(loss, params)
# Epochs can be quite long. This gives some intermediate information
# in TensorBoard.
metric_dict["minibatch_loss"] = loss
if params["use_tpu"]:
return tf.contrib.tpu.TPUEstimatorSpec(
mode=mode,
loss=loss,
train_op=train_op,
host_call=tpu_util.construct_scalar_host_call(
metric_dict=metric_dict,
model_dir=params["model_dir"],
prefix="training/"))
record_scalars(metric_dict)
return tf.estimator.EstimatorSpec(mode=mode,
loss=loss,
train_op=train_op)