diff --git a/src/mlpack/methods/decision_tree/best_binary_numeric_split_impl.hpp b/src/mlpack/methods/decision_tree/best_binary_numeric_split_impl.hpp index 45f4551600..06abf935d6 100644 --- a/src/mlpack/methods/decision_tree/best_binary_numeric_split_impl.hpp +++ b/src/mlpack/methods/decision_tree/best_binary_numeric_split_impl.hpp @@ -227,9 +227,34 @@ double BestBinaryNumericSplit::SplitIfBetter( // Force a minimum leaf size of 1 (empty children don't make sense). const size_t minimum = std::max(minimumLeafSize, (size_t) 1); + double totalWeight = 0.0; + double totalLeftWeight = 0.0; + double totalRightWeight = 0.0; + + if (UseWeights) + { + totalWeight = arma::accu(sortedWeights); + bestFoundGain *= totalWeight; + + for (size_t i = 0; i < minimum - 1; ++i) + totalLeftWeight += sortedWeights[i]; + + for (size_t i = minimum - 1; i < data.n_elem; ++i) + totalRightWeight += sortedWeights[i]; + } + else + { + bestFoundGain *= data.n_elem; + } + // Loop through all possible split points, choosing the best one. for (size_t index = minimum; index < data.n_elem - minimum + 1; ++index) { + if (UseWeights) + { + totalLeftWeight += sortedWeights[index - 1]; + totalRightWeight -= sortedWeights[index - 1]; + } // Make sure that the value has changed. if (data[sortedIndices[index]] == data[sortedIndices[index - 1]]) continue; @@ -245,7 +270,17 @@ double BestBinaryNumericSplit::SplitIfBetter( const double rightGain = FitnessFunction::template Evaluate(sortedLabels, sortedWeights, index, labels.n_elem); - double gain = leftGain + rightGain; + double gain; + if (UseWeights) + { + gain = totalLeftWeight * leftGain + totalRightWeight * rightGain; + } + else + { + // Calculate the gain at this split point. + gain = double(index) * leftGain + + double(sortedLabels.n_elem - index) * rightGain; + } // Corner case: is this the best possible split? if (gain >= 0.0) @@ -273,6 +308,11 @@ double BestBinaryNumericSplit::SplitIfBetter( if (!improved) return DBL_MAX; + if (UseWeights) + bestFoundGain /= totalWeight; + else + bestFoundGain /= data.n_elem; + return bestFoundGain; } diff --git a/src/mlpack/methods/decision_tree/mse_gain.hpp b/src/mlpack/methods/decision_tree/mse_gain.hpp index 340aee99a9..e12ae34ae4 100644 --- a/src/mlpack/methods/decision_tree/mse_gain.hpp +++ b/src/mlpack/methods/decision_tree/mse_gain.hpp @@ -73,6 +73,8 @@ class MSEGain for (size_t i = begin; i < end; ++i) mse += std::pow(labels[i] - mean, 2); + + mse /= (double) (end - begin); } return -mse; diff --git a/src/mlpack/tests/decision_tree_regressor_test.cpp b/src/mlpack/tests/decision_tree_regressor_test.cpp index ad95a3b891..496c6ab2a3 100644 --- a/src/mlpack/tests/decision_tree_regressor_test.cpp +++ b/src/mlpack/tests/decision_tree_regressor_test.cpp @@ -471,14 +471,14 @@ TEST_CASE("DecisionTreeRegressorEnergyTest", "[DecisionTreeRegressorTest]") arma::rowvec r = m.row(0); m.shed_row(0); - DecisionTreeRegressor<> d(m, r, 1, 0.0, 4); + DecisionTreeRegressor<> d(m, r, 1, 0.0, 0); arma::rowvec p; d.Predict(m, p); arma::rowvec weights = arma::ones(r.n_elem); const double mse = arma::accu(arma::square(p - r)) / p.n_elem; - REQUIRE(mse < 0.5); + REQUIRE(mse == Approx(0.0).epsilon(1e-4)); } /**