Backward compatible api change:BoostedTreesUpdateEnsembleV2 works on list of feature_ids.

PiperOrigin-RevId: 289687663
Change-Id: I5d12d044ae42fc34f03a3eaa357bf71b7cb06eec
This commit is contained in:
A. Unique TensorFlower 2020-01-14 11:02:55 -08:00 committed by TensorFlower Gardener
parent 857b55bc2d
commit 8a2a86318b
4 changed files with 35 additions and 20 deletions

View File

@ -91,6 +91,14 @@ END
name: "logits_dimension"
description: <<END
scalar, dimension of the logits
END
}
attr {
name: "num_groups"
description: <<END
Number of groups of split information to process, where a group contains feature
ids that are processed together in BoostedTreesCalculateBestFeatureSplitOpV2.
INFERRED.
END
}
summary: "Updates the tree ensemble by adding a layer to the last tree being grown"

View File

@ -269,9 +269,11 @@ class BoostedTreesUpdateEnsembleV2Op : public OpKernel {
OP_REQUIRES_OK(context,
context->input_list("split_types", &split_types_list));
const Tensor* feature_ids_t;
OP_REQUIRES_OK(context, context->input("feature_ids", &feature_ids_t));
const auto feature_ids = feature_ids_t->vec<int32>();
OpInputList feature_ids_list;
OP_REQUIRES_OK(context,
context->input_list("feature_ids", &feature_ids_list));
// TODO(crawles): Read groups of feature ids and find best splits among all.
const auto feature_ids = feature_ids_list[0].vec<int32>();
const Tensor* max_depth_t;
OP_REQUIRES_OK(context, context->input("max_depth", &max_depth_t));

View File

@ -618,7 +618,7 @@ REGISTER_OP("BoostedTreesUpdateEnsemble")
REGISTER_OP("BoostedTreesUpdateEnsembleV2")
.Input("tree_ensemble_handle: resource")
.Input("feature_ids: int32")
.Input("feature_ids: num_groups * int32")
.Input("dimension_ids: num_features * int32")
.Input("node_ids: num_features * int32")
.Input("gains: num_features * float")
@ -631,13 +631,18 @@ REGISTER_OP("BoostedTreesUpdateEnsembleV2")
.Input("pruning_mode: int32")
.Attr("num_features: int >= 0") // Inferred.
.Attr("logits_dimension: int = 1")
.Attr("num_groups: int = 1") // Number of groups to process.
.SetShapeFn([](shape_inference::InferenceContext* c) {
shape_inference::ShapeHandle shape_handle;
int num_features;
TF_RETURN_IF_ERROR(c->GetAttr("num_features", &num_features));
int num_groups;
TF_RETURN_IF_ERROR(c->GetAttr("num_groups", &num_groups));
// Feature_ids, should be one for each feature.
shape_inference::ShapeHandle feature_ids_shape;
// TODO(crawles): remove 1 hardcode once kernel operates on multiple
// groups.
TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 1, &feature_ids_shape));
TF_RETURN_IF_ERROR(
c->Merge(c->input(1), c->Vector(num_features), &shape_handle));

View File

@ -180,7 +180,7 @@ class UpdateTreeEnsembleOpTest(test_util.TensorFlowTestCase):
pruning_mode=boosted_trees_ops.PruningMode.NO_PRUNING,
# Tree will be finalized now, since we will reach depth 1.
max_depth=1,
feature_ids=feature_ids,
feature_ids=[feature_ids],
dimension_ids=[feature1_dimensions, feature2_dimensions],
node_ids=[feature1_nodes, feature2_nodes],
gains=[feature1_gains, feature2_gains],
@ -289,7 +289,7 @@ class UpdateTreeEnsembleOpTest(test_util.TensorFlowTestCase):
pruning_mode=boosted_trees_ops.PruningMode.NO_PRUNING,
# Tree will be finalized now, since we will reach depth 1.
max_depth=1,
feature_ids=feature_ids,
feature_ids=[feature_ids],
dimension_ids=[feature1_dimensions, feature2_dimensions],
node_ids=[feature1_nodes, feature2_nodes],
gains=[feature1_gains, feature2_gains],
@ -401,7 +401,7 @@ class UpdateTreeEnsembleOpTest(test_util.TensorFlowTestCase):
pruning_mode=boosted_trees_ops.PruningMode.NO_PRUNING,
# Tree will be finalized now, since we will reach depth 1.
max_depth=1,
feature_ids=feature_ids,
feature_ids=[feature_ids],
dimension_ids=[feature1_dimensions, feature2_dimensions],
node_ids=[feature1_nodes, feature2_nodes],
gains=[feature1_gains, feature2_gains],
@ -809,7 +809,7 @@ class UpdateTreeEnsembleOpTest(test_util.TensorFlowTestCase):
pruning_mode=boosted_trees_ops.PruningMode.NO_PRUNING,
# tree is going to be finalized now, since we reach depth 2.
max_depth=2,
feature_ids=feature_ids,
feature_ids=[feature_ids],
dimension_ids=[
feature1_dimensions, feature2_dimensions, feature3_dimensions
],
@ -1014,7 +1014,7 @@ class UpdateTreeEnsembleOpTest(test_util.TensorFlowTestCase):
pruning_mode=boosted_trees_ops.PruningMode.NO_PRUNING,
# tree is going to be finalized now, since we reach depth 2.
max_depth=2,
feature_ids=feature_ids,
feature_ids=[feature_ids],
dimension_ids=[
feature1_dimensions, feature2_dimensions, feature3_dimensions
],
@ -1230,7 +1230,7 @@ class UpdateTreeEnsembleOpTest(test_util.TensorFlowTestCase):
pruning_mode=boosted_trees_ops.PruningMode.NO_PRUNING,
# tree is going to be finalized now, since we reach depth 2.
max_depth=2,
feature_ids=feature_ids,
feature_ids=[feature_ids],
dimension_ids=[
feature1_dimensions, feature2_dimensions, feature3_dimensions
],
@ -1610,7 +1610,7 @@ class UpdateTreeEnsembleOpTest(test_util.TensorFlowTestCase):
pruning_mode=boosted_trees_ops.PruningMode.NO_PRUNING,
learning_rate=0.1,
max_depth=2,
feature_ids=feature_ids,
feature_ids=[feature_ids],
dimension_ids=[feature1_dimensions],
node_ids=[feature1_nodes],
gains=[feature1_gains],
@ -1769,7 +1769,7 @@ class UpdateTreeEnsembleOpTest(test_util.TensorFlowTestCase):
pruning_mode=boosted_trees_ops.PruningMode.NO_PRUNING,
learning_rate=0.1,
max_depth=2,
feature_ids=feature_ids,
feature_ids=[feature_ids],
dimension_ids=[feature1_dimensions],
node_ids=[feature1_nodes],
gains=[feature1_gains],
@ -1942,7 +1942,7 @@ class UpdateTreeEnsembleOpTest(test_util.TensorFlowTestCase):
pruning_mode=boosted_trees_ops.PruningMode.NO_PRUNING,
learning_rate=0.1,
max_depth=2,
feature_ids=feature_ids,
feature_ids=[feature_ids],
dimension_ids=[feature1_dimensions],
node_ids=[feature1_nodes],
gains=[feature1_gains],
@ -2309,7 +2309,7 @@ class UpdateTreeEnsembleOpTest(test_util.TensorFlowTestCase):
pruning_mode=boosted_trees_ops.PruningMode.PRE_PRUNING,
# tree is going to be finalized now, since we reach depth 2.
max_depth=3,
feature_ids=feature_ids,
feature_ids=[feature_ids],
dimension_ids=[
feature1_dimensions, feature2_dimensions, feature3_dimensions
],
@ -3041,7 +3041,7 @@ class UpdateTreeEnsembleOpTest(test_util.TensorFlowTestCase):
learning_rate=1.0,
pruning_mode=boosted_trees_ops.PruningMode.POST_PRUNING,
max_depth=3,
feature_ids=feature_ids,
feature_ids=[feature_ids],
dimension_ids=[feature1_dimensions, feature2_dimensions],
node_ids=[feature1_nodes, feature2_nodes],
gains=[feature1_gains, feature2_gains],
@ -3140,7 +3140,7 @@ class UpdateTreeEnsembleOpTest(test_util.TensorFlowTestCase):
learning_rate=1.0,
pruning_mode=boosted_trees_ops.PruningMode.POST_PRUNING,
max_depth=3,
feature_ids=feature_ids,
feature_ids=[feature_ids],
dimension_ids=[feature1_dimensions],
node_ids=[feature1_nodes],
gains=[feature1_gains],
@ -3293,7 +3293,7 @@ class UpdateTreeEnsembleOpTest(test_util.TensorFlowTestCase):
learning_rate=1.0,
pruning_mode=boosted_trees_ops.PruningMode.POST_PRUNING,
max_depth=3,
feature_ids=feature_ids,
feature_ids=[feature_ids],
dimension_ids=[feature1_dimensions],
node_ids=[feature1_nodes],
gains=[feature1_gains],
@ -3679,7 +3679,7 @@ class UpdateTreeEnsembleOpTest(test_util.TensorFlowTestCase):
learning_rate=1.0,
pruning_mode=boosted_trees_ops.PruningMode.POST_PRUNING,
max_depth=2,
feature_ids=feature_ids,
feature_ids=[feature_ids],
dimension_ids=[feature1_dimensions, feature2_dimensions],
node_ids=[feature1_nodes, feature2_nodes],
gains=[feature1_gains, feature2_gains],
@ -3778,7 +3778,7 @@ class UpdateTreeEnsembleOpTest(test_util.TensorFlowTestCase):
learning_rate=1.0,
pruning_mode=boosted_trees_ops.PruningMode.POST_PRUNING,
max_depth=2,
feature_ids=feature_ids,
feature_ids=[feature_ids],
dimension_ids=[feature1_dimensions],
node_ids=[feature1_nodes],
gains=[feature1_gains],
@ -4014,7 +4014,7 @@ class UpdateTreeEnsembleOpTest(test_util.TensorFlowTestCase):
learning_rate=1.0,
pruning_mode=boosted_trees_ops.PruningMode.POST_PRUNING,
max_depth=1,
feature_ids=feature_ids,
feature_ids=[feature_ids],
dimension_ids=[feature1_dimensions, feature2_dimensions],
node_ids=[feature1_nodes, feature2_nodes],
gains=[feature1_gains, feature2_gains],