def _get_train_op_and_ensemble(self, head, config, is_classification,
train_in_memory):
"""Calls bt_model_fn() and returns the train_op and ensemble_serialzed."""
features, labels = _make_train_input_fn(is_classification)()
estimator_spec = boosted_trees._bt_model_fn( # pylint:disable=protected-access
features=features,
labels=labels,
mode=model_fn.ModeKeys.TRAIN,
head=head,
feature_columns=self._feature_columns,
tree_hparams=self._tree_hparams,
example_id_column_name=EXAMPLE_ID_COLUMN,
n_batches_per_layer=1,
config=config,
train_in_memory=train_in_memory)
resources.initialize_resources(resources.shared_resources()).run()
variables.global_variables_initializer().run()
variables.local_variables_initializer().run()
# Gets the train_op and serialized proto of the ensemble.
shared_resources = resources.shared_resources()
self.assertEqual(1, len(shared_resources))
train_op = estimator_spec.train_op
with ops.control_dependencies([train_op]):
_, ensemble_serialized = (
gen_boosted_trees_ops.boosted_trees_serialize_ensemble(
shared_resources[0].handle))
return train_op, ensemble_serialized
def _get_train_op_and_ensemble(self, head, config, is_classification,
train_in_memory):
"""Calls bt_model_fn() and returns the train_op and ensemble_serialzed."""
features, labels = _make_train_input_fn(is_classification)()
estimator_spec = boosted_trees._bt_model_fn( # pylint:disable=protected-access
features=features,
labels=labels,
mode=model_fn.ModeKeys.TRAIN,
head=head,
feature_columns=self._feature_columns,
tree_hparams=self._tree_hparams,
example_id_column_name=EXAMPLE_ID_COLUMN,
n_batches_per_layer=1,
config=config,
train_in_memory=train_in_memory)
resources.initialize_resources(resources.shared_resources()).run()
variables.global_variables_initializer().run()
variables.local_variables_initializer().run()
# Gets the train_op and serialized proto of the ensemble.
shared_resources = resources.shared_resources()
self.assertEqual(1, len(shared_resources))
train_op = estimator_spec.train_op
with ops.control_dependencies([train_op]):
_, ensemble_serialized = (
gen_boosted_trees_ops.boosted_trees_serialize_ensemble(
shared_resources[0].handle))
return train_op, ensemble_serialized
def _model_fn(features, labels, mode, config):
return canned_boosted_trees._bt_model_fn(features,
labels,
mode,
head,
feature_columns,
tree_hparams,
n_batches_per_layer,
config=config)
def _model_fn(features, labels, mode, config):
return canned_boosted_trees._bt_model_fn(
features,
labels,
mode,
head,
feature_columns,
tree_hparams,
n_batches_per_layer,
config=config)
def _model_fn(features, labels, mode, config):
return canned_boosted_trees._bt_model_fn(
features,
labels,
mode,
head,
feature_columns,
tree_hparams,
n_batches_per_layer=1,
config=config,
closed_form_grad_and_hess_fn=closed_form,
train_in_memory=True)
def _model_fn(features, labels, mode, config):
return canned_boosted_trees._bt_model_fn(
features,
labels,
mode,
head,
feature_columns,
tree_hparams,
n_batches_per_layer=1,
config=config,
closed_form_grad_and_hess_fn=closed_form,
train_in_memory=True)
def _model_fn(features, labels, mode, config):
return canned_boosted_trees._bt_model_fn( # pylint: disable=protected-access
features, labels, mode, head, feature_columns, tree_hparams,
n_batches_per_layer, config)
