Test a small random forest with leaf vector output
(predict_kind)
| 87 | |
| 88 | @pytest.mark.parametrize("predict_kind", ["default", "raw", "leaf_id"]) |
| 89 | def test_leaf_vector_rf(predict_kind): |
| 90 | """Test a small random forest with leaf vector output""" |
| 91 | builder = ModelBuilder( |
| 92 | threshold_type="float32", |
| 93 | leaf_output_type="float32", |
| 94 | metadata=Metadata( |
| 95 | num_feature=1, |
| 96 | task_type="kMultiClf", |
| 97 | average_tree_output=True, |
| 98 | num_target=1, |
| 99 | num_class=[3], |
| 100 | leaf_vector_shape=(1, 3), |
| 101 | ), |
| 102 | tree_annotation=TreeAnnotation(num_tree=2, target_id=[0, 0], class_id=[-1, -1]), |
| 103 | postprocessor=PostProcessorFunc(name="identity_multiclass"), |
| 104 | base_scores=[100.0, 200.0, 300.0], |
| 105 | ) |
| 106 | |
| 107 | def make_tree_stump(left_child_val, right_child_val): |
| 108 | builder.start_tree() |
| 109 | builder.start_node(0) |
| 110 | builder.numerical_test( |
| 111 | feature_id=0, |
| 112 | threshold=0.0, |
| 113 | default_left=False, |
| 114 | opname="<", |
| 115 | left_child_key=1, |
| 116 | right_child_key=2, |
| 117 | ) |
| 118 | builder.end_node() |
| 119 | builder.start_node(1) |
| 120 | builder.leaf(left_child_val) |
| 121 | builder.end_node() |
| 122 | builder.start_node(2) |
| 123 | builder.leaf(right_child_val) |
| 124 | builder.end_node() |
| 125 | builder.end_tree() |
| 126 | |
| 127 | make_tree_stump([1.0, 0.0, 0.0], [0.0, 0.5, 0.5]) |
| 128 | make_tree_stump([1.0, 0.0, 0.0], [0.0, 0.5, 0.5]) |
| 129 | |
| 130 | model = builder.commit() |
| 131 | dmat = np.array([[1.0], [-1.0]], dtype=np.float32) |
| 132 | if predict_kind in "default": |
| 133 | expected_pred = np.array([[[100.0, 200.5, 300.5]], [[101.0, 200.0, 300.0]]]) |
| 134 | pred = treelite.gtil.predict(model, dmat, pred_margin=False) |
| 135 | elif predict_kind == "raw": |
| 136 | expected_pred = np.array([[[100.0, 200.5, 300.5]], [[101.0, 200.0, 300.0]]]) |
| 137 | pred = treelite.gtil.predict(model, dmat, pred_margin=True) |
| 138 | else: |
| 139 | expected_pred = np.array([[2, 2], [1, 1]]) |
| 140 | pred = treelite.gtil.predict_leaf(model, dmat) |
| 141 | np.testing.assert_almost_equal(pred, expected_pred, decimal=5) |
nothing calls this directly
no test coverage detected