aboutsummaryrefslogtreecommitdiffhomepage
path: root/tensorflow/contrib/tensor_forest/core/ops/tree_utils.cc
blob: f3fc41605543a000fc9269bb358d082be965971b (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
// Copyright 2016 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// =============================================================================
#include "tensorflow/contrib/tensor_forest/core/ops/tree_utils.h"

namespace tensorflow {
namespace tensorforest {

using tensorflow::Tensor;

int32 BestFeatureClassification(
    const Tensor& total_counts, const Tensor& split_counts,
    int32 accumulator) {
  int32 best_feature_index = -1;
  // We choose the split with the lowest score.
  float best_score = kint64max;
  const int32 num_splits = static_cast<int32>(split_counts.shape().dim_size(1));
  const int32 num_classes = static_cast<int32>(
      split_counts.shape().dim_size(2));
  // Ideally, Eigen::Tensor::chip would be best to use here but it results
  // in seg faults, so we have to go with flat views of these tensors.  However,
  // it is still pretty efficient because we put off evaluation until the
  // score is actually returned.
  const auto tc = total_counts.Slice(
      accumulator, accumulator + 1).unaligned_flat<float>();
  const auto splits = split_counts.Slice(
      accumulator, accumulator + 1).unaligned_flat<float>();
  Eigen::array<int, 1> bcast;
  bcast[0] = num_splits;
  const auto rights = tc.broadcast(bcast) - splits;

  for (int i = 0; i < num_splits; i++) {
    Eigen::array<int, 1> offsets;
    offsets[0] = i * num_classes;
    Eigen::array<int, 1> extents;
    extents[0] = num_classes;
    float score = WeightedGiniImpurity(splits.slice(offsets, extents)) +
        WeightedGiniImpurity(rights.slice(offsets, extents));

    if (score < best_score) {
      best_score = score;
      best_feature_index = i;
    }
  }
  return best_feature_index;
}

int32 BestFeatureRegression(
    const Tensor& total_sums, const Tensor& total_squares,
    const Tensor& split_sums, const Tensor& split_squares,
    int32 accumulator) {
  int32 best_feature_index = -1;
  // We choose the split with the lowest score.
  float best_score = kint64max;
  const int32 num_splits = static_cast<int32>(split_sums.shape().dim_size(1));
  const int32 num_regression_dims = static_cast<int32>(
      split_sums.shape().dim_size(2));
  // Ideally, Eigen::Tensor::chip would be best to use here but it results
  // in seg faults, so we have to go with flat views of these tensors.  However,
  // it is still pretty efficient because we put off evaluation until the
  // score is actually returned.
  const auto tc_sum = total_sums.Slice(
      accumulator, accumulator + 1).unaligned_flat<float>();
  const auto tc_square = total_squares.Slice(
      accumulator, accumulator + 1).unaligned_flat<float>();
  const auto splits_sum = split_sums.Slice(
      accumulator, accumulator + 1).unaligned_flat<float>();
  const auto splits_square = split_squares.Slice(
      accumulator, accumulator + 1).unaligned_flat<float>();
  // Eigen is infuriating to work with, usually resulting in all kinds of
  // unhelpful compiler errors when trying something that seems sane.  This
  // helps us do a simple thing like access the first element (the counts)
  // of these tensors so we can calculate expected value in Variance().
  const auto splits_count_accessor = split_sums.tensor<float, 3>();
  const auto totals_count_accessor = total_sums.tensor<float, 2>();

  Eigen::array<int, 1> bcast;
  bcast[0] = num_splits;
  const auto right_sums = tc_sum.broadcast(bcast) - splits_sum;
  const auto right_squares = tc_square.broadcast(bcast) - splits_square;

  for (int i = 0; i < num_splits; i++) {
    Eigen::array<int, 1> offsets;
    offsets[0] = i * num_regression_dims;
    Eigen::array<int, 1> extents;
    extents[0] = num_regression_dims;
    float left_count = splits_count_accessor(accumulator, i, 0);
    float right_count = totals_count_accessor(accumulator, 0) - left_count;

    float score = 0;

    // Guard against divide-by-zero.
    if (left_count > 0) {
      score += WeightedVariance(
        splits_sum.slice(offsets, extents),
        splits_square.slice(offsets, extents), left_count);
    }

    if (right_count > 0) {
        score += WeightedVariance(right_sums.slice(offsets, extents),
                                  right_squares.slice(offsets, extents),
                                  right_count);
    }

    if (score < best_score) {
      best_score = score;
      best_feature_index = i;
    }
  }
  return best_feature_index;
}

bool DecideNode(const Tensor& point, int32 feature, float bias) {
  const auto p = point.unaligned_flat<float>();
  CHECK_LT(feature, p.size());
  return p(feature) > bias;
}

bool IsAllInitialized(const Tensor& features) {
  const auto feature_vec = features.unaligned_flat<int32>();
  return feature_vec(feature_vec.size() - 1) >= 0;
}


}  // namespace tensorforest
}  // namespace tensorflow