def _model_fn(features, labels, mode, config):
"""Model function."""
assert labels is None, labels
(loss, scores, model_predictions, training_op, init_op,
is_initialized) = gmm_ops.gmm(
self._parse_tensor_or_dict(features),
self._training_initial_clusters, self._num_clusters,
self._random_seed, self._covariance_type, self._params)
incr_step = state_ops.assign_add(training_util.get_global_step(),
1)
training_op = with_dependencies([training_op, incr_step], loss)
training_hooks = [
_InitializeClustersHook(init_op, is_initialized,
config.is_chief)
]
predictions = {
GMM.ASSIGNMENTS: model_predictions[0][0],
}
eval_metric_ops = {
GMM.SCORES: scores,
GMM.LOG_LIKELIHOOD: _streaming_sum(loss),
}
return model_fn_lib.ModelFnOps(mode=mode,
predictions=predictions,
eval_metric_ops=eval_metric_ops,
loss=loss,
train_op=training_op,
training_hooks=training_hooks)
def estimator_spec_to_model_fn_ops(estimator_spec, export_alternatives=False):
if export_alternatives:
alternatives = _export_outputs_to_output_alternatives(
estimator_spec.export_outputs)
else:
alternatives = []
return model_fn.ModelFnOps(mode=_core_mode_to_contrib_mode(
estimator_spec.mode),
predictions=estimator_spec.predictions,
loss=estimator_spec.loss,
train_op=estimator_spec.train_op,
eval_metric_ops=estimator_spec.eval_metric_ops,
output_alternatives=alternatives)
def _model_fn(features, labels, mode):
"""Function that returns predictions, training loss, and training op."""
model_fn_ops = []
for i in range(len(model_fns)):
with variable_scope.variable_scope('label_{0}'.format(i)):
sliced_labels = array_ops.slice(labels, [0, i], [-1, 1])
model_fn_ops.append(model_fns[i](features, sliced_labels,
mode))
training_hooks = []
for mops in model_fn_ops:
training_hooks += mops.training_hooks
predictions = {}
if (mode == model_fn_lib.ModeKeys.EVAL
or mode == model_fn_lib.ModeKeys.INFER):
# Flatten the probabilities into one dimension.
predictions[eval_metrics.INFERENCE_PROB_NAME] = array_ops.concat(
[
mops.predictions[eval_metrics.INFERENCE_PROB_NAME]
for mops in model_fn_ops
],
axis=1)
predictions[eval_metrics.INFERENCE_PRED_NAME] = array_ops.stack(
[
mops.predictions[eval_metrics.INFERENCE_PRED_NAME]
for mops in model_fn_ops
],
axis=1)
loss = None
if (mode == model_fn_lib.ModeKeys.EVAL
or mode == model_fn_lib.ModeKeys.TRAIN):
loss = math_ops.reduce_sum(
array_ops.stack([mops.loss for mops in model_fn_ops
])) / len(model_fn_ops)
train_op = None
if mode == model_fn_lib.ModeKeys.TRAIN:
train_op = control_flow_ops.group(
*[mops.train_op for mops in model_fn_ops])
return model_fn_lib.ModelFnOps(mode=mode,
predictions=predictions,
loss=loss,
train_op=train_op,
training_hooks=training_hooks,
scaffold=None,
output_alternatives=None)
def create_model_fn_ops(self,
predictions,
output_alternatives,
mode=model_fn.ModeKeys.INFER):
return model_fn.ModelFnOps(
model_fn.ModeKeys.INFER,
predictions=predictions,
loss=constant_op.constant([1]),
train_op=control_flow_ops.no_op(),
eval_metric_ops={
"metric_key":
(constant_op.constant(1.), control_flow_ops.no_op()),
"loss": (constant_op.constant(1.), control_flow_ops.no_op()),
},
training_chief_hooks=[basic_session_run_hooks.StepCounterHook()],
training_hooks=[basic_session_run_hooks.StepCounterHook()],
output_alternatives=output_alternatives,
scaffold=monitored_session.Scaffold())
def _model_fn(features, labels, mode, params):
"""Model function that appends logistic evaluation metrics."""
thresholds = params.get('thresholds') or [.5]
predictions, loss, train_op = model_fn(features, labels, mode)
if mode == model_fn_lib.ModeKeys.EVAL:
eval_metric_ops = _make_logistic_eval_metric_ops(
labels=labels, predictions=predictions, thresholds=thresholds)
else:
eval_metric_ops = None
return model_fn_lib.ModelFnOps(
mode=mode,
predictions=predictions,
loss=loss,
train_op=train_op,
eval_metric_ops=eval_metric_ops,
output_alternatives={
'head': (constants.ProblemType.LOGISTIC_REGRESSION, {
'predictions': predictions
})
})
def _model_fn(features, labels, mode):
_ = labels
x = features['x']
y = features['y']
with ops.name_scope('outputs'):
predictions = {
'sum': math_ops.add(x, y, name='sum'),
'product': math_ops.multiply(x, y, name='product'),
'difference': math_ops.subtract(x, y, name='difference')
}
if core:
export_outputs = {
k: export_output.PredictOutput({k: v})
for k, v in predictions.items()
}
export_outputs[
signature_constants.
DEFAULT_SERVING_SIGNATURE_DEF_KEY] = export_outputs['sum']
return model_fn.EstimatorSpec(mode=mode,
predictions=predictions,
export_outputs=export_outputs,
loss=constant_op.constant(0),
train_op=control_flow_ops.no_op())
else:
output_alternatives = {
k: (constants.ProblemType.UNSPECIFIED, {
k: v
})
for k, v in predictions.items()
}
return contrib_model_fn.ModelFnOps(
mode=mode,
predictions=predictions,
output_alternatives=output_alternatives,
loss=constant_op.constant(0),
train_op=control_flow_ops.no_op())
def _wals_factorization_model_function(features, labels, mode, params):
"""Model function for the WALSFactorization estimator.
Args:
features: Dictionary of features. See WALSMatrixFactorization.
labels: Must be None.
mode: A model_fn.ModeKeys object.
params: Dictionary of parameters containing arguments passed to the
WALSMatrixFactorization constructor.
Returns:
A ModelFnOps object.
Raises:
ValueError: If `mode` is not recognized.
"""
assert labels is None
use_factors_weights_cache = (
params["use_factors_weights_cache_for_training"]
and mode == model_fn.ModeKeys.TRAIN)
use_gramian_cache = (params["use_gramian_cache_for_training"]
and mode == model_fn.ModeKeys.TRAIN)
max_sweeps = params["max_sweeps"]
model = factorization_ops.WALSModel(
params["num_rows"],
params["num_cols"],
params["embedding_dimension"],
unobserved_weight=params["unobserved_weight"],
regularization=params["regularization_coeff"],
row_init=params["row_init"],
col_init=params["col_init"],
num_row_shards=params["num_row_shards"],
num_col_shards=params["num_col_shards"],
row_weights=params["row_weights"],
col_weights=params["col_weights"],
use_factors_weights_cache=use_factors_weights_cache,
use_gramian_cache=use_gramian_cache)
# Get input rows and cols. We either update rows or columns depending on
# the value of row_sweep, which is maintained using a session hook.
input_rows = features[WALSMatrixFactorization.INPUT_ROWS]
input_cols = features[WALSMatrixFactorization.INPUT_COLS]
# TRAIN mode:
if mode == model_fn.ModeKeys.TRAIN:
# Training consists of the following ops (controlled using a SweepHook).
# Before a row sweep:
# row_update_prep_gramian_op
# initialize_row_update_op
# During a row sweep:
# update_row_factors_op
# Before a col sweep:
# col_update_prep_gramian_op
# initialize_col_update_op
# During a col sweep:
# update_col_factors_op
is_row_sweep_var = variable_scope.variable(
True,
trainable=False,
name="is_row_sweep",
collections=[ops.GraphKeys.GLOBAL_VARIABLES])
is_sweep_done_var = variable_scope.variable(
False,
trainable=False,
name="is_sweep_done",
collections=[ops.GraphKeys.GLOBAL_VARIABLES])
completed_sweeps_var = variable_scope.variable(
0,
trainable=False,
name=WALSMatrixFactorization.COMPLETED_SWEEPS,
collections=[ops.GraphKeys.GLOBAL_VARIABLES])
loss_var = variable_scope.variable(
0.,
trainable=False,
name=WALSMatrixFactorization.LOSS,
collections=[ops.GraphKeys.GLOBAL_VARIABLES])
# The root weighted squared error =
# \\(\sqrt( \sum_{i,j} w_ij * (a_ij - r_ij)^2 / \sum_{i,j} w_ij )\\)
rwse_var = variable_scope.variable(
0.,
trainable=False,
name=WALSMatrixFactorization.RWSE,
collections=[ops.GraphKeys.GLOBAL_VARIABLES])
summary.scalar("loss", loss_var)
summary.scalar("root_weighted_squared_error", rwse_var)
summary.scalar("completed_sweeps", completed_sweeps_var)
def create_axis_ops(sp_input, num_items, update_fn, axis_name):
"""Creates book-keeping and training ops for a given axis.
Args:
sp_input: A SparseTensor corresponding to the row or column batch.
num_items: An integer, the total number of items of this axis.
update_fn: A function that takes one argument (`sp_input`), and that
returns a tuple of
* new_factors: A float Tensor of the factor values after update.
* update_op: a TensorFlow op which updates the factors.
* loss: A float Tensor, the unregularized loss.
* reg_loss: A float Tensor, the regularization loss.
* sum_weights: A float Tensor, the sum of factor weights.
axis_name: A string that specifies the name of the axis.
Returns:
A tuple consisting of:
* reset_processed_items_op: A TensorFlow op, to be run before the
beginning of any sweep. It marks all items as not-processed.
* axis_train_op: A Tensorflow op, to be run during this axis' sweeps.
"""
processed_items_init = array_ops.fill(dims=[num_items],
value=False)
with ops.colocate_with(processed_items_init):
processed_items = variable_scope.variable(
processed_items_init,
collections=[ops.GraphKeys.GLOBAL_VARIABLES],
trainable=False,
name="processed_" + axis_name)
_, update_op, loss, reg, sum_weights = update_fn(sp_input)
input_indices = sp_input.indices[:, 0]
with ops.control_dependencies([
update_op,
state_ops.assign(loss_var, loss + reg),
state_ops.assign(rwse_var,
math_ops.sqrt(loss / sum_weights))
]):
with ops.colocate_with(processed_items):
update_processed_items = state_ops.scatter_update(
processed_items,
input_indices,
array_ops.ones_like(input_indices, dtype=dtypes.bool),
name="update_processed_{}_indices".format(axis_name))
with ops.control_dependencies([update_processed_items]):
is_sweep_done = math_ops.reduce_all(processed_items)
axis_train_op = control_flow_ops.group(
state_ops.assign(is_sweep_done_var, is_sweep_done),
state_ops.assign_add(
completed_sweeps_var,
math_ops.cast(is_sweep_done, dtypes.int32)),
name="{}_sweep_train_op".format(axis_name))
return processed_items.initializer, axis_train_op
reset_processed_rows_op, row_train_op = create_axis_ops(
input_rows, params["num_rows"],
lambda x: model.update_row_factors(sp_input=x,
transpose_input=False), "rows")
reset_processed_cols_op, col_train_op = create_axis_ops(
input_cols, params["num_cols"],
lambda x: model.update_col_factors(sp_input=x,
transpose_input=True), "cols")
switch_op = control_flow_ops.group(state_ops.assign(
is_row_sweep_var, math_ops.logical_not(is_row_sweep_var)),
reset_processed_rows_op,
reset_processed_cols_op,
name="sweep_switch_op")
row_prep_ops = [
model.row_update_prep_gramian_op, model.initialize_row_update_op
]
col_prep_ops = [
model.col_update_prep_gramian_op, model.initialize_col_update_op
]
init_op = model.worker_init
sweep_hook = _SweepHook(is_row_sweep_var, is_sweep_done_var, init_op,
row_prep_ops, col_prep_ops, row_train_op,
col_train_op, switch_op)
global_step_hook = _IncrementGlobalStepHook()
training_hooks = [sweep_hook, global_step_hook]
if max_sweeps is not None:
training_hooks.append(_StopAtSweepHook(max_sweeps))
return model_fn.ModelFnOps(mode=model_fn.ModeKeys.TRAIN,
predictions={},
loss=loss_var,
eval_metric_ops={},
train_op=control_flow_ops.no_op(),
training_hooks=training_hooks)
# INFER mode
elif mode == model_fn.ModeKeys.INFER:
projection_weights = features.get(
WALSMatrixFactorization.PROJECTION_WEIGHTS)
def get_row_projection():
return model.project_row_factors(
sp_input=input_rows,
projection_weights=projection_weights,
transpose_input=False)
def get_col_projection():
return model.project_col_factors(
sp_input=input_cols,
projection_weights=projection_weights,
transpose_input=True)
predictions = {
WALSMatrixFactorization.PROJECTION_RESULT:
control_flow_ops.cond(
features[WALSMatrixFactorization.PROJECT_ROW],
get_row_projection, get_col_projection)
}
return model_fn.ModelFnOps(mode=model_fn.ModeKeys.INFER,
predictions=predictions,
loss=None,
eval_metric_ops={},
train_op=control_flow_ops.no_op(),
training_hooks=[])
# EVAL mode
elif mode == model_fn.ModeKeys.EVAL:
def get_row_loss():
_, _, loss, reg, _ = model.update_row_factors(
sp_input=input_rows, transpose_input=False)
return loss + reg
def get_col_loss():
_, _, loss, reg, _ = model.update_col_factors(sp_input=input_cols,
transpose_input=True)
return loss + reg
loss = control_flow_ops.cond(
features[WALSMatrixFactorization.PROJECT_ROW], get_row_loss,
get_col_loss)
return model_fn.ModelFnOps(mode=model_fn.ModeKeys.EVAL,
predictions={},
loss=loss,
eval_metric_ops={},
train_op=control_flow_ops.no_op(),
training_hooks=[])
else:
raise ValueError("mode=%s is not recognized." % str(mode))
def _dynamic_rnn_model_fn(features, labels, mode):
"""The model to be passed to an `Estimator`."""
with ops.name_scope(name):
sequence_length = features.get(sequence_length_key)
sequence_input = build_sequence_input(features,
sequence_feature_columns,
context_feature_columns)
dropout = (dropout_keep_probabilities
if mode == model_fn.ModeKeys.TRAIN else None)
# This class promises to use the cell type selected by that function.
cell = rnn_common.construct_rnn_cell(num_units, cell_type, dropout)
initial_state = dict_to_state_tuple(features, cell)
rnn_activations, final_state = construct_rnn(
initial_state,
sequence_input,
cell,
target_column.num_label_columns,
dtype=dtype,
parallel_iterations=parallel_iterations,
swap_memory=swap_memory)
loss = None # Created below for modes TRAIN and EVAL.
if prediction_type == rnn_common.PredictionType.MULTIPLE_VALUE:
prediction_dict = rnn_common.multi_value_predictions(
rnn_activations, target_column, problem_type,
predict_probabilities)
if mode != model_fn.ModeKeys.INFER:
loss = _multi_value_loss(rnn_activations, labels,
sequence_length, target_column,
features)
elif prediction_type == rnn_common.PredictionType.SINGLE_VALUE:
prediction_dict = _single_value_predictions(
rnn_activations, sequence_length, target_column,
problem_type, predict_probabilities)
if mode != model_fn.ModeKeys.INFER:
loss = _single_value_loss(rnn_activations, labels,
sequence_length, target_column,
features)
state_dict = state_tuple_to_dict(final_state)
prediction_dict.update(state_dict)
eval_metric_ops = None
if mode != model_fn.ModeKeys.INFER:
eval_metric_ops = rnn_common.get_eval_metric_ops(
problem_type, prediction_type, sequence_length,
prediction_dict, labels)
train_op = None
if mode == model_fn.ModeKeys.TRAIN:
train_op = optimizers.optimize_loss(
loss=loss,
global_step=None, # Get it internally.
learning_rate=learning_rate,
optimizer=optimizer,
clip_gradients=gradient_clipping_norm,
summaries=optimizers.OPTIMIZER_SUMMARIES)
output_alternatives = _get_output_alternatives(prediction_type,
problem_type,
prediction_dict)
return model_fn.ModelFnOps(mode=mode,
predictions=prediction_dict,
loss=loss,
train_op=train_op,
eval_metric_ops=eval_metric_ops,
output_alternatives=output_alternatives)
def _rnn_model_fn(features, labels, mode):
"""The model to be passed to an `Estimator`."""
with ops.name_scope(name):
dropout = (dropout_keep_probabilities
if mode == model_fn.ModeKeys.TRAIN else None)
cell = rnn_common.construct_rnn_cell(num_units, cell_type, dropout)
batch = _read_batch(
cell=cell,
features=features,
labels=labels,
mode=mode,
num_unroll=num_unroll,
batch_size=batch_size,
sequence_feature_columns=sequence_feature_columns,
context_feature_columns=context_feature_columns,
num_threads=num_threads,
queue_capacity=queue_capacity,
seed=seed)
sequence_features = batch.sequences
context_features = batch.context
if mode != model_fn.ModeKeys.INFER:
labels = sequence_features.pop(rnn_common.RNNKeys.LABELS_KEY)
inputs = _prepare_inputs_for_rnn(sequence_features,
context_features,
sequence_feature_columns,
num_unroll)
state_name = _get_state_names(cell)
rnn_activations, final_state = construct_state_saving_rnn(
cell=cell,
inputs=inputs,
num_label_columns=target_column.num_label_columns,
state_saver=batch,
state_name=state_name)
loss = None # Created below for modes TRAIN and EVAL.
prediction_dict = rnn_common.multi_value_predictions(
rnn_activations, target_column, problem_type,
predict_probabilities)
if mode != model_fn.ModeKeys.INFER:
loss = _multi_value_loss(rnn_activations, labels, batch.length,
target_column, features)
eval_metric_ops = None
if mode != model_fn.ModeKeys.INFER:
eval_metric_ops = rnn_common.get_eval_metric_ops(
problem_type, rnn_common.PredictionType.MULTIPLE_VALUE,
batch.length, prediction_dict, labels)
state_dict = state_tuple_to_dict(final_state)
prediction_dict.update(state_dict)
train_op = None
if mode == model_fn.ModeKeys.TRAIN:
train_op = optimizers.optimize_loss(
loss=loss,
global_step=None, # Get it internally.
learning_rate=learning_rate,
optimizer=optimizer,
clip_gradients=gradient_clipping_norm,
summaries=optimizers.OPTIMIZER_SUMMARIES)
return model_fn.ModelFnOps(mode=mode,
predictions=prediction_dict,
loss=loss,
train_op=train_op,
eval_metric_ops=eval_metric_ops)