1 // Copyright 2017 The TensorFlow Authors. All Rights Reserved.
2 //
3 // Licensed under the Apache License, Version 2.0 (the "License");
4 // you may not use this file except in compliance with the License.
5 // You may obtain a copy of the License at
6 //
7 //     http://www.apache.org/licenses/LICENSE-2.0
8 //
9 // Unless required by applicable law or agreed to in writing, software
10 // distributed under the License is distributed on an "AS IS" BASIS,
11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 // See the License for the specific language governing permissions and
13 // limitations under the License.
14 // =============================================================================
15 #ifndef TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_TESTUTIL_RANDOM_TREE_GEN_H_
16 #define TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_TESTUTIL_RANDOM_TREE_GEN_H_
17 
18 #include <memory>
19 
20 #include "tensorflow/contrib/boosted_trees/proto/tree_config.pb.h"  // NOLINT
21 #include "tensorflow/core/lib/random/simple_philox.h"
22 #include "tensorflow/core/platform/macros.h"
23 
24 namespace tensorflow {
25 namespace boosted_trees {
26 namespace testutil {
27 
28 // Randomly generate a balanced tree, for performance benchmarking purposes,
29 // that assume all features are sparse float features, for now.
30 class RandomTreeGen {
31  public:
32   RandomTreeGen(tensorflow::random::SimplePhilox* rng, int dense_feature_size,
33                 int sparse_feature_size);
34 
35   // Required: depth must be >= 1.
36   // If one wants to generate multiple trees with the same depth, see also the
37   // overload below.
38   boosted_trees::trees::DecisionTreeConfig Generate(int depth);
39 
40   // Randomly generate a new tree with the same depth (and tree structure)
41   // as the given tree. This is faster.
42   boosted_trees::trees::DecisionTreeConfig Generate(
43       const boosted_trees::trees::DecisionTreeConfig& tree);
44 
45   // Required: depth >= 1; tree_count >= 1.
46   boosted_trees::trees::DecisionTreeEnsembleConfig GenerateEnsemble(
47       int dept, int tree_count);
48 
49  private:
50   tensorflow::random::SimplePhilox* rng_;
51   const int dense_feature_size_;
52   const int sparse_feature_size_;
53 
54   // Put together a deeper tree by combining two trees.
55   void Combine(boosted_trees::trees::DecisionTreeConfig* root,
56                boosted_trees::trees::DecisionTreeConfig* left_branch,
57                boosted_trees::trees::DecisionTreeConfig* right_branch);
58 
59   // For each node in the provided tree, shift its referenced left/right index
60   // by shift.
61   void ShiftNodeIndex(boosted_trees::trees::DecisionTreeConfig* tree,
62                       int shift);
63 
64   // Generate a sparse split in the node.
65   void GenerateSplit(boosted_trees::trees::TreeNode* node, int left_id,
66                      int right_id);
67 
68   TF_DISALLOW_COPY_AND_ASSIGN(RandomTreeGen);
69 };
70 
71 }  // namespace testutil
72 }  // namespace boosted_trees
73 }  // namespace tensorflow
74 
75 #endif  // TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_TESTUTIL_RANDOM_TREE_GEN_H_
76