Return probabilities instead of raw values for Classification Trees (#37

)
BayesWitnesses · Jan 29, 2019 · 5aabac1 · 5aabac1
1 parent 6a9833a
commit 5aabac1
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/m2cgen/assemblers/tree.py b/m2cgen/assemblers/tree.py
@@ -37,8 +37,9 @@ def _assemble_leaf(self, node_id):
         scores = self._tree.value[node_id][0]
         if self._is_vector_output:
             outputs = []
+            score_sum = scores.sum() or 1.0
             for s in scores:
-                outputs.append(ast.NumVal(s))
+                outputs.append(ast.NumVal(s / score_sum))
             return ast.VectorExpr(outputs)
         else:
             assert len(scores) == 1, "Unexpected number of outputs"

diff --git a/tests/utils.py b/tests/utils.py
@@ -73,7 +73,7 @@ def _train_model(estimator, dataset, test_fraction):
     if isinstance(estimator, LinearClassifierMixin):
         y_pred = estimator.decision_function(X_test)
     elif isinstance(estimator, DecisionTreeClassifier):
-        y_pred = estimator.tree_.predict(X_test.astype(np.float32))
+        y_pred = estimator.predict_proba(X_test.astype(np.float32))
     else:
         y_pred = estimator.predict(X_test)