def _get_default_head(params, weights_name, output_type, name=None):
"""Creates a default head based on a type of a problem."""
if output_type == ModelBuilderOutputType.MODEL_FN_OPS:
if params.regression:
return head_lib.regression_head(
weight_column_name=weights_name,
label_dimension=params.num_outputs,
enable_centered_bias=False,
head_name=name)
else:
return head_lib.multi_class_head(
params.num_classes,
weight_column_name=weights_name,
enable_centered_bias=False,
head_name=name)
else:
if params.regression:
return core_head_lib.regression_head(
weight_column=weights_name,
label_dimension=params.num_outputs,
name=name,
loss_reduction=losses.Reduction.SUM_OVER_BATCH_SIZE)
else:
if params.num_classes == 2:
return core_head_lib.binary_classification_head(
weight_column=weights_name,
name=name,
loss_reduction=losses.Reduction.SUM_OVER_BATCH_SIZE)
else:
return core_head_lib.multi_class_head(
n_classes=params.num_classes,
weight_column=weights_name,
name=name,
loss_reduction=losses.Reduction.SUM_OVER_BATCH_SIZE)
def _get_default_head(params, weights_name, output_type, name=None):
"""Creates a default head based on a type of a problem."""
if output_type == ModelBuilderOutputType.MODEL_FN_OPS:
if params.regression:
return head_lib.regression_head(weight_column_name=weights_name,
label_dimension=params.num_outputs,
enable_centered_bias=False,
head_name=name)
else:
return head_lib.multi_class_head(params.num_classes,
weight_column_name=weights_name,
enable_centered_bias=False,
head_name=name)
else:
if params.regression:
return core_head_lib.regression_head(
weight_column=weights_name,
label_dimension=params.num_outputs,
name=name,
loss_reduction=losses.Reduction.SUM_OVER_BATCH_SIZE)
else:
if params.num_classes == 2:
return core_head_lib.binary_classification_head(
weight_column=weights_name,
name=name,
loss_reduction=losses.Reduction.SUM_OVER_BATCH_SIZE)
else:
return core_head_lib.multi_class_head(
n_classes=params.num_classes,
weight_column=weights_name,
name=name,
loss_reduction=losses.Reduction.SUM_OVER_BATCH_SIZE)
def _rnn_estimator_fn(feature_columns, n_classes, cell_units, model_dir):
return rnn.RNNEstimator(
head=head_lib.multi_class_head(n_classes=n_classes),
num_units=cell_units,
sequence_feature_columns=feature_columns,
model_dir=model_dir)
def _rnn_estimator_fn(feature_columns, n_classes, cell_units, model_dir):
return rnn.RNNEstimator(
head=head_lib.multi_class_head(n_classes=n_classes),
num_units=cell_units,
sequence_feature_columns=feature_columns,
model_dir=model_dir)
def _dnn_estimator_classifier_fn(n_classes=3, *args, **kwargs): # pylint: disable=keyword-arg-before-vararg
"""Returns a DNNEstimator that uses multi_class_head."""
return dnn.DNNEstimator(head=head_lib.multi_class_head(n_classes=n_classes),
*args, **kwargs)
def _dnn_estimator_classifier_fn(n_classes=3, *args, **kwargs): # pylint: disable=keyword-arg-before-vararg
"""Returns a DNNEstimator that uses multi_class_head."""
return dnn.DNNEstimator(
head=head_lib.multi_class_head(n_classes=n_classes), *args, **kwargs)