def testMakeStatsSummaryMultipleFeatures(self):
"""Tests that MakeStatsSummary works for multiple features."""
with self.cached_session():
max_splits = 3
num_buckets = 4
node_ids = [1, 1, 2, 2, 1, 1, 2, 0]
gradients = [[.1], [.2], [.3], [-.4], [-.05], [.06], [.07], [.08]]
hessians = [[.2], [.3], [.4], [.5], [.06], [.07], [.08], [.09]]
# Tests multiple features.
# The output from another feature will stored be in 3rd dimension.
bucketized_features = [[3, 1, 2, 0, 1, 2, 0, 1],
[0, 0, 0, 2, 2, 3, 3, 2]]
result = boosted_trees_ops.make_stats_summary(
node_ids, gradients, hessians, bucketized_features, max_splits,
num_buckets
) # shape=[max_splits, num_buckets, num_features, 2]
self.assertAllClose(
[
[
[[0., 0.], [.08, .09], [0., 0.], [0., 0.]], # node 0
[[0., 0.], [.15, .36], [.06, .07], [.1, .2]], # node 1
[[-.33, .58], [0., 0.], [.3, .4], [0., 0.]], # node 2
], # feature 0
[
[[0., 0.], [0., 0.], [.08, .09], [0., 0.]], # node 0
[[.3, .5], [0., 0.], [-.05, .06], [.06, .07]
], # node 1
[[.3, .4], [0., 0.], [-.4, .5], [.07, .08]], # node 2
], # feature 1
],
result.eval())
def testMakeStatsSummaryMultipleFeatures(self):
"""Tests that MakeStatsSummary works for multiple features."""
with self.test_session():
max_splits = 3
num_buckets = 4
node_ids = [1, 1, 2, 2, 1, 1, 2, 0]
gradients = [[.1], [.2], [.3], [-.4], [-.05], [.06], [.07], [.08]]
hessians = [[.2], [.3], [.4], [.5], [.06], [.07], [.08], [.09]]
# Tests multiple features.
# The output from another feature will stored be in 3rd dimension.
bucketized_features = [[3, 1, 2, 0, 1, 2, 0, 1], [0, 0, 0, 2, 2, 3, 3, 2]]
result = boosted_trees_ops.make_stats_summary(
node_ids, gradients, hessians, bucketized_features, max_splits,
num_buckets) # shape=[max_splits, num_buckets, num_features, 2]
self.assertAllClose(
[
[
[[0., 0.], [.08, .09], [0., 0.], [0., 0.]], # node 0
[[0., 0.], [.15, .36], [.06, .07], [.1, .2]], # node 1
[[-.33, .58], [0., 0.], [.3, .4], [0., 0.]], # node 2
], # feature 0
[
[[0., 0.], [0., 0.], [.08, .09], [0., 0.]], # node 0
[[.3, .5], [0., 0.], [-.05, .06], [.06, .07]], # node 1
[[.3, .4], [0., 0.], [-.4, .5], [.07, .08]], # node 2
], # feature 1
],
result.eval())
def testMakeStatsSummarySimple(self):
"""Simple test for MakeStatsSummary."""
with self.test_session():
self.assertAllClose([[[[1., 5.], [2., 6.]], [[3., 7.], [4., 8.]]]],
boosted_trees_ops.make_stats_summary(
node_ids=[0, 0, 1, 1],
gradients=[[1.], [2.], [3.], [4.]],
hessians=[[5.], [6.], [7.], [8.]],
bucketized_features_list=[[0, 1, 0, 1]],
max_splits=2,
num_buckets=2).eval())
def testMakeStatsSummarySimple(self):
"""Simple test for MakeStatsSummary."""
with self.cached_session():
self.assertAllClose([[[[1., 5.], [2., 6.]], [[3., 7.], [4., 8.]]]],
boosted_trees_ops.make_stats_summary(
node_ids=[0, 0, 1, 1],
gradients=[[1.], [2.], [3.], [4.]],
hessians=[[5.], [6.], [7.], [8.]],
bucketized_features_list=[[0, 1, 0, 1]],
max_splits=2,
num_buckets=2).eval())
def testMakeStatsSummarySimple(self):
"""Simple test for MakeStatsSummary."""
expected_stats_summary = np.asarray([1., 5., 2., 6., 3., 7., 4., 8.])
self.assertAllClose(
expected_stats_summary.reshape((1, 2, 2, 2)),
boosted_trees_ops.make_stats_summary(
node_ids=[0, 0, 1, 1],
gradients=[[1.], [2.], [3.], [4.]],
hessians=[[5.], [6.], [7.], [8.]],
bucketized_features_list=[[0, 1, 0, 1]],
max_splits=2,
num_buckets=2))
def testMakeStatsSummarySimple(self):
"""Simple test for MakeStatsSummary."""
expected_stats_summary = np.asarray([1., 5., 2., 6., 3., 7., 4., 8.])
self.assertAllClose(
expected_stats_summary.reshape((1, 2, 2, 2)),
boosted_trees_ops.make_stats_summary(
node_ids=[0, 0, 1, 1],
gradients=[[1.], [2.], [3.], [4.]],
hessians=[[5.], [6.], [7.], [8.]],
bucketized_features_list=[[0, 1, 0, 1]],
max_splits=2,
num_buckets=2))
def _verify_precision(self, length):
with self.cached_session():
max_splits = 1
num_buckets = 1
node_ids = array_ops.fill([length], 0)
gradients = constant_op.constant(
2.0 / length, dtype=dtypes.float32, shape=[length, 1])
hessians = constant_op.constant(
0.2 / length, dtype=dtypes.float32, shape=[length, 1])
bucketized_features = array_ops.zeros([length], dtype=dtypes.int32)
result = boosted_trees_ops.make_stats_summary(
node_ids, gradients, hessians, [bucketized_features], max_splits,
num_buckets) # shape=[max_splits, num_buckets, num_features, 2]
self.assertAllClose([[[[2., 0.2]]]], result.eval())
def _verify_precision(self, length):
with self.cached_session():
max_splits = 1
num_buckets = 1
node_ids = array_ops.fill([length], 0)
gradients = constant_op.constant(
2.0 / length, dtype=dtypes.float32, shape=[length, 1])
hessians = constant_op.constant(
0.2 / length, dtype=dtypes.float32, shape=[length, 1])
bucketized_features = array_ops.zeros([length], dtype=dtypes.int32)
result = boosted_trees_ops.make_stats_summary(
node_ids, gradients, hessians, [bucketized_features], max_splits,
num_buckets) # shape=[max_splits, num_buckets, num_features, 2]
self.assertAllClose([[[[2., 0.2]]]], self.evaluate(result))
def testMakeStatsSummaryAccumulate(self):
"""Tests that Summary actually accumulates."""
with self.test_session():
max_splits = 3
num_buckets = 4
node_ids = [1, 1, 2, 2, 1, 1, 2, 0]
gradients = [[.1], [.2], [.3], [-.4], [-.05], [.06], [.07], [.08]]
hessians = [[.2], [.3], [.4], [.5], [.06], [.07], [.08], [.09]]
# Tests a single feature.
bucketized_features = [[3, 1, 2, 0, 1, 2, 0, 1]]
result = boosted_trees_ops.make_stats_summary(
node_ids, gradients, hessians, bucketized_features, max_splits,
num_buckets) # shape=[max_splits, num_buckets, num_features, 2]
self.assertAllClose(
[[
[[0., 0.], [.08, .09], [0., 0.], [0., 0.]], # node 0
[[0., 0.], [.15, .36], [.06, .07], [.1, .2]], # node 1
[[-.33, .58], [0., 0.], [.3, .4], [0., 0.]], # node 2
]],
result.eval())
def testMakeStatsSummaryAccumulate(self):
"""Tests that Summary actually accumulates."""
with self.cached_session():
max_splits = 3
num_buckets = 4
node_ids = [1, 1, 2, 2, 1, 1, 2, 0]
gradients = [[.1], [.2], [.3], [-.4], [-.05], [.06], [.07], [.08]]
hessians = [[.2], [.3], [.4], [.5], [.06], [.07], [.08], [.09]]
# Tests a single feature.
bucketized_features = [[3, 1, 2, 0, 1, 2, 0, 1]]
result = boosted_trees_ops.make_stats_summary(
node_ids, gradients, hessians, bucketized_features, max_splits,
num_buckets) # shape=[max_splits, num_buckets, num_features, 2]
self.assertAllClose(
[[
[[0., 0.], [.08, .09], [0., 0.], [0., 0.]], # node 0
[[0., 0.], [.15, .36], [.06, .07], [.1, .2]], # node 1
[[-.33, .58], [0., 0.], [.3, .4], [0., 0.]], # node 2
]],
self.evaluate(result))
def _train_op_fn(loss):
"""Run one training iteration."""
if training_state_cache:
train_op.append(training_state_cache.insert(tree_ids, node_ids, logits))
if closed_form_grad_and_hess_fn:
gradients, hessians = closed_form_grad_and_hess_fn(logits, labels)
else:
gradients = gradients_impl.gradients(loss, logits, name='Gradients')[0]
hessians = gradients_impl.gradients(
gradients, logits, name='Hessians')[0]
stats_summaries_list = []
for i, feature_ids in enumerate(feature_ids_list):
num_buckets = bucket_size_list[i]
summaries = [
array_ops.squeeze(
boosted_trees_ops.make_stats_summary(
node_ids=node_ids,
gradients=gradients,
hessians=hessians,
bucketized_features_list=[input_feature_list[f]],
max_splits=max_splits,
num_buckets=num_buckets),
axis=0) for f in feature_ids
]
stats_summaries_list.append(summaries)
accumulators = []
def grow_tree_from_stats_summaries(stats_summaries_list,
feature_ids_list):
"""Updates ensemble based on the best gains from stats summaries."""
node_ids_per_feature = []
gains_list = []
thresholds_list = []
left_node_contribs_list = []
right_node_contribs_list = []
all_feature_ids = []
assert len(stats_summaries_list) == len(feature_ids_list)
for i, feature_ids in enumerate(feature_ids_list):
(numeric_node_ids_per_feature, numeric_gains_list,
numeric_thresholds_list, numeric_left_node_contribs_list,
numeric_right_node_contribs_list) = (
boosted_trees_ops.calculate_best_gains_per_feature(
node_id_range=last_layer_nodes_range,
stats_summary_list=stats_summaries_list[i],
l1=tree_hparams.l1,
l2=tree_hparams.l2,
tree_complexity=tree_hparams.tree_complexity,
min_node_weight=tree_hparams.min_node_weight,
max_splits=max_splits))
all_feature_ids += feature_ids
node_ids_per_feature += numeric_node_ids_per_feature
gains_list += numeric_gains_list
thresholds_list += numeric_thresholds_list
left_node_contribs_list += numeric_left_node_contribs_list
right_node_contribs_list += numeric_right_node_contribs_list
grow_op = boosted_trees_ops.update_ensemble(
# Confirm if local_tree_ensemble or tree_ensemble should be used.
tree_ensemble.resource_handle,
feature_ids=all_feature_ids,
node_ids=node_ids_per_feature,
gains=gains_list,
thresholds=thresholds_list,
left_node_contribs=left_node_contribs_list,
right_node_contribs=right_node_contribs_list,
learning_rate=tree_hparams.learning_rate,
max_depth=tree_hparams.max_depth,
pruning_mode=boosted_trees_ops.PruningMode.NO_PRUNING)
return grow_op
if train_in_memory and is_single_machine:
train_op.append(distribute_lib.increment_var(global_step))
train_op.append(
grow_tree_from_stats_summaries(stats_summaries_list,
feature_ids_list))
else:
dependencies = []
for i, feature_ids in enumerate(feature_ids_list):
stats_summaries = stats_summaries_list[i]
accumulator = data_flow_ops.ConditionalAccumulator(
dtype=dtypes.float32,
# The stats consist of grads and hessians (the last dimension).
shape=[len(feature_ids), max_splits, bucket_size_list[i], 2],
shared_name='numeric_stats_summary_accumulator_' + str(i))
accumulators.append(accumulator)
apply_grad = accumulator.apply_grad(
array_ops.stack(stats_summaries, axis=0), stamp_token)
dependencies.append(apply_grad)
def grow_tree_from_accumulated_summaries_fn():
"""Updates the tree with the best layer from accumulated summaries."""
# Take out the accumulated summaries from the accumulator and grow.
stats_summaries_list = []
stats_summaries_list = [
array_ops.unstack(accumulator.take_grad(1), axis=0)
for accumulator in accumulators
]
grow_op = grow_tree_from_stats_summaries(stats_summaries_list,
feature_ids_list)
return grow_op
with ops.control_dependencies(dependencies):
train_op.append(distribute_lib.increment_var(global_step))
if config.is_chief:
min_accumulated = math_ops.reduce_min(
array_ops.stack(
[acc.num_accumulated() for acc in accumulators]))
train_op.append(
control_flow_ops.cond(
math_ops.greater_equal(min_accumulated,
n_batches_per_layer),
grow_tree_from_accumulated_summaries_fn,
control_flow_ops.no_op,
name='wait_until_n_batches_accumulated'))
return control_flow_ops.group(train_op, name='train_op')
def _train_op_fn(loss):
"""Run one training iteration."""
if training_state_cache:
# Cache logits only after center_bias is complete, if it's in progress.
train_op.append(
control_flow_ops.cond(
center_bias_var, control_flow_ops.no_op,
lambda: training_state_cache.insert(tree_ids, node_ids, logits))
)
if closed_form_grad_and_hess_fn:
gradients, hessians = closed_form_grad_and_hess_fn(logits, labels)
else:
gradients = gradients_impl.gradients(loss, logits, name='Gradients')[0]
hessians = gradients_impl.gradients(
gradients, logits, name='Hessians')[0]
stats_summaries_list = []
for i, feature_ids in enumerate(feature_ids_list):
num_buckets = bucket_size_list[i]
summaries = [
array_ops.squeeze(
boosted_trees_ops.make_stats_summary(
node_ids=node_ids,
gradients=gradients,
hessians=hessians,
bucketized_features_list=[input_feature_list[f]],
max_splits=max_splits,
num_buckets=num_buckets),
axis=0) for f in feature_ids
]
stats_summaries_list.append(summaries)
# ========= Helper methods for both in and not in memory. ==============
def grow_tree_from_stats_summaries(stats_summaries_list,
feature_ids_list):
"""Updates ensemble based on the best gains from stats summaries."""
node_ids_per_feature = []
gains_list = []
thresholds_list = []
left_node_contribs_list = []
right_node_contribs_list = []
all_feature_ids = []
assert len(stats_summaries_list) == len(feature_ids_list)
for i, feature_ids in enumerate(feature_ids_list):
(numeric_node_ids_per_feature, numeric_gains_list,
numeric_thresholds_list, numeric_left_node_contribs_list,
numeric_right_node_contribs_list) = (
boosted_trees_ops.calculate_best_gains_per_feature(
node_id_range=last_layer_nodes_range,
stats_summary_list=stats_summaries_list[i],
l1=tree_hparams.l1,
l2=tree_hparams.l2,
tree_complexity=tree_hparams.tree_complexity,
min_node_weight=tree_hparams.min_node_weight,
max_splits=max_splits))
all_feature_ids += feature_ids
node_ids_per_feature += numeric_node_ids_per_feature
gains_list += numeric_gains_list
thresholds_list += numeric_thresholds_list
left_node_contribs_list += numeric_left_node_contribs_list
right_node_contribs_list += numeric_right_node_contribs_list
grow_op = boosted_trees_ops.update_ensemble(
# Confirm if local_tree_ensemble or tree_ensemble should be used.
tree_ensemble.resource_handle,
feature_ids=all_feature_ids,
node_ids=node_ids_per_feature,
gains=gains_list,
thresholds=thresholds_list,
left_node_contribs=left_node_contribs_list,
right_node_contribs=right_node_contribs_list,
learning_rate=tree_hparams.learning_rate,
max_depth=tree_hparams.max_depth,
pruning_mode=boosted_trees_ops.PruningMode.NO_PRUNING)
return grow_op
def _center_bias_fn(mean_gradients, mean_hessians):
"""Updates the ensembles and cache (if needed) with logits prior."""
continue_centering = boosted_trees_ops.center_bias(
tree_ensemble.resource_handle,
mean_gradients=mean_gradients,
mean_hessians=mean_hessians,
l1=tree_hparams.l1,
l2=tree_hparams.l2
)
return center_bias_var.assign(continue_centering)
# ========= End of helper methods. ==============
if train_in_memory and is_single_machine:
train_op.append(distribute_lib.increment_var(global_step))
mean_gradients = array_ops.expand_dims(
math_ops.reduce_mean(gradients, 0), 0)
mean_heassians = array_ops.expand_dims(
math_ops.reduce_mean(hessians, 0), 0)
train_op.append(
control_flow_ops.cond(
center_bias_var,
lambda: _center_bias_fn(mean_gradients, mean_heassians),
functools.partial(grow_tree_from_stats_summaries,
stats_summaries_list, feature_ids_list)))
else:
def center_bias_not_in_mem():
"""Accumulates the data and updates the logits bias, when ready."""
bias_dependencies = []
bias_accumulator = data_flow_ops.ConditionalAccumulator(
dtype=dtypes.float32,
# The stats consist of grads and hessians means only.
# TODO(nponomareva): this will change for a multiclass
shape=[2, 1],
shared_name='bias_accumulator')
grads_and_hess = array_ops.stack([gradients, hessians], axis=0)
grads_and_hess = math_ops.reduce_mean(grads_and_hess, axis=1)
apply_grad = bias_accumulator.apply_grad(grads_and_hess, stamp_token)
bias_dependencies.append(apply_grad)
def center_bias_from_accumulator():
accumulated = array_ops.unstack(
bias_accumulator.take_grad(1), axis=0)
return _center_bias_fn(
array_ops.expand_dims(accumulated[0], 0),
array_ops.expand_dims(accumulated[1], 0))
with ops.control_dependencies(bias_dependencies):
if config.is_chief:
center_bias_op = control_flow_ops.cond(
math_ops.greater_equal(bias_accumulator.num_accumulated(),
n_batches_per_layer),
center_bias_from_accumulator,
control_flow_ops.no_op,
name='wait_until_n_batches_for_bias_accumulated')
return center_bias_op
else:
return control_flow_ops.no_op()
def grow_not_in_mem():
"""Accumulates the data and grows a layer when ready."""
accumulators = []
dependencies = []
for i, feature_ids in enumerate(feature_ids_list):
stats_summaries = stats_summaries_list[i]
accumulator = data_flow_ops.ConditionalAccumulator(
dtype=dtypes.float32,
# The stats consist of grads and hessians (the last dimension).
shape=[len(feature_ids), max_splits, bucket_size_list[i], 2],
shared_name='numeric_stats_summary_accumulator_' + str(i))
accumulators.append(accumulator)
apply_grad = accumulator.apply_grad(
array_ops.stack(stats_summaries, axis=0), stamp_token)
dependencies.append(apply_grad)
def grow_tree_from_accumulated_summaries_fn():
"""Updates tree with the best layer from accumulated summaries."""
# Take out the accumulated summaries from the accumulator and grow.
stats_summaries_list = []
stats_summaries_list = [
array_ops.unstack(accumulator.take_grad(1), axis=0)
for accumulator in accumulators
]
grow_op = grow_tree_from_stats_summaries(stats_summaries_list,
feature_ids_list)
return grow_op
with ops.control_dependencies(dependencies):
if config.is_chief:
min_accumulated = math_ops.reduce_min(
array_ops.stack(
[acc.num_accumulated() for acc in accumulators]))
grow_model = control_flow_ops.cond(
math_ops.greater_equal(min_accumulated, n_batches_per_layer),
grow_tree_from_accumulated_summaries_fn,
control_flow_ops.no_op,
name='wait_until_n_batches_accumulated')
return grow_model
else:
return control_flow_ops.no_op()
update_model = control_flow_ops.cond(
center_bias_var, center_bias_not_in_mem, grow_not_in_mem)
train_op.append(update_model)
with ops.control_dependencies([update_model]):
increment_global = distribute_lib.increment_var(global_step)
train_op.append(increment_global)
return control_flow_ops.group(train_op, name='train_op')
def _train_op_fn(loss):
"""Run one training iteration."""
if training_state_cache:
train_op.append(training_state_cache.insert(tree_ids, node_ids, logits))
if closed_form_grad_and_hess_fn:
gradients, hessians = closed_form_grad_and_hess_fn(logits, labels)
else:
gradients = gradients_impl.gradients(loss, logits, name='Gradients')[0]
hessians = gradients_impl.gradients(
gradients, logits, name='Hessians')[0]
stats_summaries_list = []
for i, feature_ids in enumerate(feature_ids_list):
num_buckets = bucket_size_list[i]
summaries = [
array_ops.squeeze(
boosted_trees_ops.make_stats_summary(
node_ids=node_ids,
gradients=gradients,
hessians=hessians,
bucketized_features_list=[input_feature_list[f]],
max_splits=max_splits,
num_buckets=num_buckets),
axis=0) for f in feature_ids
]
stats_summaries_list.append(summaries)
accumulators = []
def grow_tree_from_stats_summaries(stats_summaries_list,
feature_ids_list):
"""Updates ensemble based on the best gains from stats summaries."""
node_ids_per_feature = []
gains_list = []
thresholds_list = []
left_node_contribs_list = []
right_node_contribs_list = []
all_feature_ids = []
assert len(stats_summaries_list) == len(feature_ids_list)
for i, feature_ids in enumerate(feature_ids_list):
(numeric_node_ids_per_feature, numeric_gains_list,
numeric_thresholds_list, numeric_left_node_contribs_list,
numeric_right_node_contribs_list) = (
boosted_trees_ops.calculate_best_gains_per_feature(
node_id_range=last_layer_nodes_range,
stats_summary_list=stats_summaries_list[i],
l1=tree_hparams.l1,
l2=tree_hparams.l2,
tree_complexity=tree_hparams.tree_complexity,
min_node_weight=tree_hparams.min_node_weight,
max_splits=max_splits))
all_feature_ids += feature_ids
node_ids_per_feature += numeric_node_ids_per_feature
gains_list += numeric_gains_list
thresholds_list += numeric_thresholds_list
left_node_contribs_list += numeric_left_node_contribs_list
right_node_contribs_list += numeric_right_node_contribs_list
grow_op = boosted_trees_ops.update_ensemble(
# Confirm if local_tree_ensemble or tree_ensemble should be used.
tree_ensemble.resource_handle,
feature_ids=all_feature_ids,
node_ids=node_ids_per_feature,
gains=gains_list,
thresholds=thresholds_list,
left_node_contribs=left_node_contribs_list,
right_node_contribs=right_node_contribs_list,
learning_rate=tree_hparams.learning_rate,
max_depth=tree_hparams.max_depth,
pruning_mode=boosted_trees_ops.PruningMode.NO_PRUNING)
return grow_op
if train_in_memory and is_single_machine:
train_op.append(distribute_lib.increment_var(global_step))
train_op.append(
grow_tree_from_stats_summaries(stats_summaries_list,
feature_ids_list))
else:
dependencies = []
for i, feature_ids in enumerate(feature_ids_list):
stats_summaries = stats_summaries_list[i]
accumulator = data_flow_ops.ConditionalAccumulator(
dtype=dtypes.float32,
# The stats consist of grads and hessians (the last dimension).
shape=[len(feature_ids), max_splits, bucket_size_list[i], 2],
shared_name='numeric_stats_summary_accumulator_' + str(i))
accumulators.append(accumulator)
apply_grad = accumulator.apply_grad(
array_ops.stack(stats_summaries, axis=0), stamp_token)
dependencies.append(apply_grad)
def grow_tree_from_accumulated_summaries_fn():
"""Updates the tree with the best layer from accumulated summaries."""
# Take out the accumulated summaries from the accumulator and grow.
stats_summaries_list = []
stats_summaries_list = [
array_ops.unstack(accumulator.take_grad(1), axis=0)
for accumulator in accumulators
]
grow_op = grow_tree_from_stats_summaries(stats_summaries_list,
feature_ids_list)
return grow_op
with ops.control_dependencies(dependencies):
train_op.append(distribute_lib.increment_var(global_step))
if config.is_chief:
min_accumulated = math_ops.reduce_min(
array_ops.stack(
[acc.num_accumulated() for acc in accumulators]))
train_op.append(
control_flow_ops.cond(
math_ops.greater_equal(min_accumulated,
n_batches_per_layer),
grow_tree_from_accumulated_summaries_fn,
control_flow_ops.no_op,
name='wait_until_n_batches_accumulated'))
return control_flow_ops.group(train_op, name='train_op')
def _train_op_fn(loss):
"""Run one training iteration."""
train_op = []
if cache:
train_op.append(cache.insert(tree_ids, node_ids, logits))
if closed_form_grad_and_hess_fn:
gradients, hessians = closed_form_grad_and_hess_fn(
logits, labels)
else:
gradients = gradients_impl.gradients(loss,
logits,
name='Gradients')[0]
hessians = gradients_impl.gradients(gradients,
logits,
name='Hessians')[0]
stats_summary_list = [
array_ops.squeeze(boosted_trees_ops.make_stats_summary(
node_ids=node_ids,
gradients=gradients,
hessians=hessians,
bucketized_features_list=[input_feature_list[f]],
max_splits=max_splits,
num_buckets=num_buckets),
axis=0) for f in range(num_features)
]
def grow_tree_from_stats_summaries(stats_summary_list):
"""Updates ensemble based on the best gains from stats summaries."""
(node_ids_per_feature, gains_list, thresholds_list,
left_node_contribs_list, right_node_contribs_list) = (
boosted_trees_ops.calculate_best_gains_per_feature(
node_id_range=last_layer_nodes_range,
stats_summary_list=stats_summary_list,
l1=tree_hparams.l1,
l2=tree_hparams.l2,
tree_complexity=tree_hparams.tree_complexity,
max_splits=max_splits))
grow_op = boosted_trees_ops.update_ensemble(
# Confirm if local_tree_ensemble or tree_ensemble should be used.
tree_ensemble.resource_handle,
feature_ids=math_ops.range(0,
num_features,
dtype=dtypes.int32),
node_ids=node_ids_per_feature,
gains=gains_list,
thresholds=thresholds_list,
left_node_contribs=left_node_contribs_list,
right_node_contribs=right_node_contribs_list,
learning_rate=tree_hparams.learning_rate,
max_depth=tree_hparams.max_depth,
pruning_mode=boosted_trees_ops.PruningMode.NO_PRUNING)
return grow_op
if train_in_memory and is_single_machine:
train_op.append(distribute_lib.increment_var(global_step))
train_op.append(
grow_tree_from_stats_summaries(stats_summary_list))
else:
summary_accumulator = data_flow_ops.ConditionalAccumulator(
dtype=dtypes.float32,
# The stats consist of gradients and hessians (the last dimension).
shape=[num_features, max_splits, num_buckets, 2],
shared_name='stats_summary_accumulator')
apply_grad = summary_accumulator.apply_grad(
array_ops.stack(stats_summary_list, axis=0), stamp_token)
def grow_tree_from_accumulated_summaries_fn():
"""Updates the tree with the best layer from accumulated summaries."""
# Take out the accumulated summaries from the accumulator and grow.
stats_summary_list = array_ops.unstack(
summary_accumulator.take_grad(1), axis=0)
grow_op = grow_tree_from_stats_summaries(
stats_summary_list)
return grow_op
with ops.control_dependencies([apply_grad]):
train_op.append(distribute_lib.increment_var(global_step))
if config.is_chief:
train_op.append(
control_flow_ops.cond(
math_ops.greater_equal(
summary_accumulator.num_accumulated(),
n_batches_per_layer),
grow_tree_from_accumulated_summaries_fn,
control_flow_ops.no_op,
name='wait_until_n_batches_accumulated'))
return control_flow_ops.group(train_op, name='train_op')
def _train_op_fn(loss):
"""Run one training iteration."""
if training_state_cache:
# Cache logits only after center_bias is complete, if it's in progress.
train_op.append(
control_flow_ops.cond(
center_bias_var, control_flow_ops.no_op,
lambda: training_state_cache.insert(
tree_ids, node_ids, logits)))
if closed_form_grad_and_hess_fn:
gradients, hessians = closed_form_grad_and_hess_fn(
logits, labels)
else:
gradients = gradients_impl.gradients(loss,
logits,
name='Gradients')[0]
hessians = gradients_impl.gradients(gradients,
logits,
name='Hessians')[0]
stats_summaries_list = []
for i, feature_ids in enumerate(feature_ids_list):
num_buckets = bucket_size_list[i]
summaries = [
array_ops.squeeze(boosted_trees_ops.make_stats_summary(
node_ids=node_ids,
gradients=gradients,
hessians=hessians,
bucketized_features_list=[input_feature_list[f]],
max_splits=max_splits,
num_buckets=num_buckets),
axis=0) for f in feature_ids
]
stats_summaries_list.append(summaries)
# ========= Helper methods for both in and not in memory. ==============
def grow_tree_from_stats_summaries(stats_summaries_list,
feature_ids_list):
"""Updates ensemble based on the best gains from stats summaries."""
node_ids_per_feature = []
gains_list = []
thresholds_list = []
left_node_contribs_list = []
right_node_contribs_list = []
all_feature_ids = []
assert len(stats_summaries_list) == len(feature_ids_list)
for i, feature_ids in enumerate(feature_ids_list):
(numeric_node_ids_per_feature, numeric_gains_list,
numeric_thresholds_list, numeric_left_node_contribs_list,
numeric_right_node_contribs_list) = (
boosted_trees_ops.calculate_best_gains_per_feature(
node_id_range=last_layer_nodes_range,
stats_summary_list=stats_summaries_list[i],
l1=tree_hparams.l1,
l2=tree_hparams.l2,
tree_complexity=tree_hparams.tree_complexity,
min_node_weight=tree_hparams.min_node_weight,
max_splits=max_splits))
all_feature_ids += feature_ids
node_ids_per_feature += numeric_node_ids_per_feature
gains_list += numeric_gains_list
thresholds_list += numeric_thresholds_list
left_node_contribs_list += numeric_left_node_contribs_list
right_node_contribs_list += numeric_right_node_contribs_list
grow_op = boosted_trees_ops.update_ensemble(
# Confirm if local_tree_ensemble or tree_ensemble should be used.
tree_ensemble.resource_handle,
feature_ids=all_feature_ids,
node_ids=node_ids_per_feature,
gains=gains_list,
thresholds=thresholds_list,
left_node_contribs=left_node_contribs_list,
right_node_contribs=right_node_contribs_list,
learning_rate=tree_hparams.learning_rate,
max_depth=tree_hparams.max_depth,
pruning_mode=boosted_trees_ops.PruningMode.NO_PRUNING)
return grow_op
def _center_bias_fn(mean_gradients, mean_hessians):
"""Updates the ensembles and cache (if needed) with logits prior."""
continue_centering = boosted_trees_ops.center_bias(
tree_ensemble.resource_handle,
mean_gradients=mean_gradients,
mean_hessians=mean_hessians,
l1=tree_hparams.l1,
l2=tree_hparams.l2)
return center_bias_var.assign(continue_centering)
# ========= End of helper methods. ==============
if train_in_memory and is_single_machine:
train_op.append(distribute_lib.increment_var(global_step))
mean_gradients = array_ops.expand_dims(
math_ops.reduce_mean(gradients, 0), 0)
mean_heassians = array_ops.expand_dims(
math_ops.reduce_mean(hessians, 0), 0)
train_op.append(
control_flow_ops.cond(
center_bias_var, lambda: _center_bias_fn(
mean_gradients, mean_heassians),
functools.partial(grow_tree_from_stats_summaries,
stats_summaries_list,
feature_ids_list)))
else:
def center_bias_not_in_mem():
"""Accumulates the data and updates the logits bias, when ready."""
bias_dependencies = []
bias_accumulator = data_flow_ops.ConditionalAccumulator(
dtype=dtypes.float32,
# The stats consist of grads and hessians means only.
# TODO(nponomareva): this will change for a multiclass
shape=[2, 1],
shared_name='bias_accumulator')
grads_and_hess = array_ops.stack([gradients, hessians],
axis=0)
grads_and_hess = math_ops.reduce_mean(grads_and_hess,
axis=1)
apply_grad = bias_accumulator.apply_grad(
grads_and_hess, stamp_token)
bias_dependencies.append(apply_grad)
def center_bias_from_accumulator():
accumulated = array_ops.unstack(
bias_accumulator.take_grad(1), axis=0)
return _center_bias_fn(
array_ops.expand_dims(accumulated[0], 0),
array_ops.expand_dims(accumulated[1], 0))
with ops.control_dependencies(bias_dependencies):
if config.is_chief:
center_bias_op = control_flow_ops.cond(
math_ops.greater_equal(
bias_accumulator.num_accumulated(),
n_batches_per_layer),
center_bias_from_accumulator,
control_flow_ops.no_op,
name='wait_until_n_batches_for_bias_accumulated'
)
return center_bias_op
else:
return control_flow_ops.no_op()
def grow_not_in_mem():
"""Accumulates the data and grows a layer when ready."""
accumulators = []
dependencies = []
for i, feature_ids in enumerate(feature_ids_list):
stats_summaries = stats_summaries_list[i]
accumulator = data_flow_ops.ConditionalAccumulator(
dtype=dtypes.float32,
# The stats consist of grads and hessians (the last dimension).
shape=[
len(feature_ids), max_splits,
bucket_size_list[i], 2
],
shared_name='numeric_stats_summary_accumulator_' +
str(i))
accumulators.append(accumulator)
apply_grad = accumulator.apply_grad(
array_ops.stack(stats_summaries, axis=0),
stamp_token)
dependencies.append(apply_grad)
def grow_tree_from_accumulated_summaries_fn():
"""Updates tree with the best layer from accumulated summaries."""
# Take out the accumulated summaries from the accumulator and grow.
stats_summaries_list = []
stats_summaries_list = [
array_ops.unstack(accumulator.take_grad(1), axis=0)
for accumulator in accumulators
]
grow_op = grow_tree_from_stats_summaries(
stats_summaries_list, feature_ids_list)
return grow_op
with ops.control_dependencies(dependencies):
if config.is_chief:
min_accumulated = math_ops.reduce_min(
array_ops.stack([
acc.num_accumulated()
for acc in accumulators
]))
grow_model = control_flow_ops.cond(
math_ops.greater_equal(min_accumulated,
n_batches_per_layer),
grow_tree_from_accumulated_summaries_fn,
control_flow_ops.no_op,
name='wait_until_n_batches_accumulated')
return grow_model
else:
return control_flow_ops.no_op()
update_model = control_flow_ops.cond(center_bias_var,
center_bias_not_in_mem,
grow_not_in_mem)
train_op.append(update_model)
with ops.control_dependencies([update_model]):
increment_global = distribute_lib.increment_var(
global_step)
train_op.append(increment_global)
return control_flow_ops.group(train_op, name='train_op')
def _train_op_fn(loss):
"""Run one training iteration."""
train_op = []
if cache:
train_op.append(cache.insert(tree_ids, node_ids, logits))
if closed_form_grad_and_hess_fn:
gradients, hessians = closed_form_grad_and_hess_fn(logits, labels)
else:
gradients = gradients_impl.gradients(loss, logits, name='Gradients')[0]
hessians = gradients_impl.gradients(
gradients, logits, name='Hessians')[0]
stats_summary_list = [
array_ops.squeeze(
boosted_trees_ops.make_stats_summary(
node_ids=node_ids,
gradients=gradients,
hessians=hessians,
bucketized_features_list=[input_feature_list[f]],
max_splits=max_splits,
num_buckets=num_buckets),
axis=0) for f in range(num_features)
]
def grow_tree_from_stats_summaries(stats_summary_list):
"""Updates ensemble based on the best gains from stats summaries."""
(node_ids_per_feature, gains_list, thresholds_list,
left_node_contribs_list, right_node_contribs_list) = (
boosted_trees_ops.calculate_best_gains_per_feature(
node_id_range=last_layer_nodes_range,
stats_summary_list=stats_summary_list,
l1=tree_hparams.l1,
l2=tree_hparams.l2,
tree_complexity=tree_hparams.tree_complexity,
max_splits=max_splits))
grow_op = boosted_trees_ops.update_ensemble(
# Confirm if local_tree_ensemble or tree_ensemble should be used.
tree_ensemble.resource_handle,
feature_ids=math_ops.range(0, num_features, dtype=dtypes.int32),
node_ids=node_ids_per_feature,
gains=gains_list,
thresholds=thresholds_list,
left_node_contribs=left_node_contribs_list,
right_node_contribs=right_node_contribs_list,
learning_rate=tree_hparams.learning_rate,
max_depth=tree_hparams.max_depth,
pruning_mode=boosted_trees_ops.PruningMode.NO_PRUNING)
return grow_op
if train_in_memory and is_single_machine:
train_op.append(distribute_lib.increment_var(global_step))
train_op.append(grow_tree_from_stats_summaries(stats_summary_list))
else:
summary_accumulator = data_flow_ops.ConditionalAccumulator(
dtype=dtypes.float32,
# The stats consist of gradients and hessians (the last dimension).
shape=[num_features, max_splits, num_buckets, 2],
shared_name='stats_summary_accumulator')
apply_grad = summary_accumulator.apply_grad(
array_ops.stack(stats_summary_list, axis=0), stamp_token)
def grow_tree_from_accumulated_summaries_fn():
"""Updates the tree with the best layer from accumulated summaries."""
# Take out the accumulated summaries from the accumulator and grow.
stats_summary_list = array_ops.unstack(
summary_accumulator.take_grad(1), axis=0)
grow_op = grow_tree_from_stats_summaries(stats_summary_list)
return grow_op
with ops.control_dependencies([apply_grad]):
train_op.append(distribute_lib.increment_var(global_step))
if config.is_chief:
train_op.append(
control_flow_ops.cond(
math_ops.greater_equal(
summary_accumulator.num_accumulated(),
n_batches_per_layer),
grow_tree_from_accumulated_summaries_fn,
control_flow_ops.no_op,
name='wait_until_n_batches_accumulated'))
return control_flow_ops.group(train_op, name='train_op')
def _train_op_fn(loss):
"""Run one training iteration."""
if training_state_cache:
# Cache logits only after center_bias is complete, if it's in progress.
train_op.append(
control_flow_ops.cond(
center_bias_var, control_flow_ops.no_op,
lambda: training_state_cache.insert(tree_ids, node_ids, logits))
)
if closed_form_grad_and_hess_fn:
gradients, hessians = closed_form_grad_and_hess_fn(logits, labels)
else:
gradients = gradients_impl.gradients(loss, logits, name='Gradients')[0]
hessians = gradients_impl.gradients(
gradients, logits, name='Hessians')[0]
# TODO(youngheek): perhaps storage could be optimized by storing stats
# with the dimension max_splits_per_layer, instead of max_splits (for the
# entire tree).
max_splits = _get_max_splits(tree_hparams)
stats_summaries_list = []
for i, feature_ids in enumerate(feature_ids_list):
num_buckets = bucket_size_list[i]
summaries = [
array_ops.squeeze(
boosted_trees_ops.make_stats_summary(
node_ids=node_ids,
gradients=gradients,
hessians=hessians,
bucketized_features_list=[input_feature_list[f]],
max_splits=max_splits,
num_buckets=num_buckets),
axis=0) for f in feature_ids
]
stats_summaries_list.append(summaries)
if train_in_memory and is_single_machine:
grower = _InMemoryEnsembleGrower(tree_ensemble, tree_hparams)
else:
grower = _AccumulatorEnsembleGrower(tree_ensemble, tree_hparams,
stamp_token, n_batches_per_layer,
bucket_size_list, config.is_chief)
update_model = control_flow_ops.cond(
center_bias_var,
functools.partial(
grower.center_bias,
center_bias_var,
gradients,
hessians,
),
functools.partial(grower.grow_tree, stats_summaries_list,
feature_ids_list, last_layer_nodes_range))
train_op.append(update_model)
with ops.control_dependencies([update_model]):
increment_global = distribute_lib.increment_var(global_step)
train_op.append(increment_global)
return control_flow_ops.group(train_op, name='train_op')