def test_node_insert_delete(tmpdir, toolchain):
"""Test ability to add and remove nodes"""
num_feature = 3
builder = treelite.ModelBuilder(num_feature=num_feature)
builder.append(treelite.ModelBuilder.Tree())
builder[0][1].set_root()
builder[0][1].set_numerical_test_node(feature_id=2,
opname='<',
threshold=-0.5,
default_left=True,
left_child_key=5,
right_child_key=10)
builder[0][5].set_leaf_node(-1)
builder[0][10].set_numerical_test_node(feature_id=0,
opname='<=',
threshold=0.5,
default_left=False,
left_child_key=7,
right_child_key=8)
builder[0][7].set_leaf_node(0.0)
builder[0][8].set_leaf_node(1.0)
del builder[0][1]
del builder[0][5]
builder[0][5].set_categorical_test_node(feature_id=1,
left_categories=[1, 2, 4],
default_left=True,
left_child_key=20,
right_child_key=10)
builder[0][20].set_leaf_node(2.0)
builder[0][5].set_root()
model = builder.commit()
assert model.num_feature == num_feature
assert model.num_output_group == 1
assert model.num_tree == 1
libpath = os.path.join(tmpdir, 'libtest' + _libext())
model.export_lib(toolchain=toolchain, libpath=libpath, verbose=True)
predictor = treelite_runtime.Predictor(libpath=libpath)
assert predictor.num_feature == num_feature
assert predictor.num_output_group == 1
assert predictor.pred_transform == 'identity'
assert predictor.global_bias == 0.0
assert predictor.sigmoid_alpha == 1.0
for f0 in [-0.5, 0.5, 1.5, np.nan]:
for f1 in [0, 1, 2, 3, 4, np.nan]:
for f2 in [-1.0, -0.5, 1.0, np.nan]:
x = np.array([f0, f1, f2])
pred = predictor.predict_instance(x)
if f1 in [1, 2, 4] or np.isnan(f1):
expected_pred = 2.0
elif f0 <= 0.5 and not np.isnan(f0):
expected_pred = 0.0
else:
expected_pred = 1.0
if pred != expected_pred:
raise ValueError(
f'Prediction wrong for f0={f0}, f1={f1}, f2={f2}: ' +
f'expected_pred = {expected_pred} vs actual_pred = {pred}'
)
def test_model_query(tmpdir, dataset):
"""Test all query functions for every example model"""
if dataset == 'sparse_categorical':
if os_platform() == 'windows':
pytest.xfail('MSVC cannot handle long if conditional')
elif os_platform() == 'osx':
pytest.xfail('Apple Clang cannot handle long if conditional')
if dataset == 'letor' and os_platform() == 'windows':
pytest.xfail('export_lib() is too slow for letor on MSVC')
libpath = os.path.join(tmpdir, dataset_db[dataset].libname + _libext())
model = treelite.Model.load(dataset_db[dataset].model, model_format=dataset_db[dataset].format)
assert model.num_feature == _model_facts[dataset].num_feature
assert model.num_class == _model_facts[dataset].num_class
assert model.num_tree == _model_facts[dataset].num_tree
toolchain = os_compatible_toolchains()[0]
model.export_lib(toolchain=toolchain, libpath=libpath,
params={'quantize': 1, 'parallel_comp': model.num_tree}, verbose=True)
predictor = treelite_runtime.Predictor(libpath=libpath, verbose=True)
assert predictor.num_feature == _model_facts[dataset].num_feature
assert predictor.num_class == _model_facts[dataset].num_class
assert predictor.pred_transform == _model_facts[dataset].pred_transform
assert predictor.global_bias == _model_facts[dataset].global_bias
assert predictor.sigmoid_alpha == _model_facts[dataset].sigmoid_alpha
assert predictor.threshold_type == _model_facts[dataset].threshold_type
assert predictor.leaf_output_type == _model_facts[dataset].leaf_output_type
def test_code_folding(tmpdir, annotation, dataset, toolchain,
code_folding_factor):
"""Test suite for testing code folding feature"""
if dataset == 'letor' and os_platform() == 'windows':
pytest.xfail('export_lib() is too slow for letor on MSVC')
libpath = os.path.join(tmpdir, dataset_db[dataset].libname + _libext())
model = treelite.Model.load(dataset_db[dataset].model,
model_format=dataset_db[dataset].format)
annotation_path = os.path.join(tmpdir, 'annotation.json')
if annotation[dataset] is None:
annotation_path = None
else:
with open(annotation_path, 'w') as f:
f.write(annotation[dataset])
params = {
'annotate_in': (annotation_path if annotation_path else 'NULL'),
'quantize': 1,
'parallel_comp': model.num_tree,
'code_folding_req': code_folding_factor
}
model.export_lib(toolchain=toolchain,
libpath=libpath,
params=params,
verbose=True)
predictor = treelite_runtime.Predictor(libpath=libpath, verbose=True)
check_predictor(predictor, dataset)
def test_skl_converter_regressor(tmpdir, clazz, toolchain): # pylint: disable=too-many-locals
"""Convert scikit-learn regressor"""
X, y = load_boston(return_X_y=True)
kwargs = {}
if clazz == GradientBoostingRegressor:
kwargs['init'] = 'zero'
clf = clazz(max_depth=3, random_state=0, n_estimators=10, **kwargs)
clf.fit(X, y)
expected_pred = clf.predict(X)
model = treelite.sklearn.import_model(clf)
assert model.num_feature == clf.n_features_
assert model.num_output_group == 1
assert model.num_tree == clf.n_estimators
dtrain = treelite_runtime.DMatrix(X, dtype='float32')
annotation_path = os.path.join(tmpdir, 'annotation.json')
annotator = treelite.Annotator()
annotator.annotate_branch(model=model, dmat=dtrain, verbose=True)
annotator.save(path=annotation_path)
libpath = os.path.join(tmpdir, 'skl' + _libext())
model.export_lib(toolchain=toolchain,
libpath=libpath,
params={'annotate_in': annotation_path},
verbose=True)
predictor = treelite_runtime.Predictor(libpath=libpath)
assert predictor.num_feature == clf.n_features_
assert model.num_output_group == 1
assert predictor.pred_transform == 'identity'
assert predictor.global_bias == 0.0
assert predictor.sigmoid_alpha == 1.0
out_pred = predictor.predict(dtrain)
np.testing.assert_almost_equal(out_pred, expected_pred, decimal=5)
def test_skl_converter_iforest(tmpdir, toolchain): # pylint: disable=W0212
# pylint: disable=too-many-locals
"""Convert scikit-learn regressor"""
X = load_boston(return_X_y=True)[0]
clf = IsolationForest(max_samples=64, random_state=0, n_estimators=10)
clf.fit(X)
expected_pred = clf._compute_chunked_score_samples(X) # pylint: disable=W0212
model = treelite.sklearn.import_model(clf)
assert model.num_feature == clf.n_features_
assert model.num_class == 1
assert model.num_tree == clf.n_estimators
dtrain = treelite_runtime.DMatrix(X, dtype='float32')
annotation_path = os.path.join(tmpdir, 'annotation.json')
annotator = treelite.Annotator()
annotator.annotate_branch(model=model, dmat=dtrain, verbose=True)
annotator.save(path=annotation_path)
libpath = os.path.join(tmpdir, 'skl' + _libext())
model.export_lib(toolchain=toolchain,
libpath=libpath,
params={'annotate_in': annotation_path},
verbose=True)
predictor = treelite_runtime.Predictor(libpath=libpath)
assert predictor.num_feature == clf.n_features_
assert model.num_class == 1
assert predictor.pred_transform == 'exponential_standard_ratio'
assert predictor.global_bias == 0.0
assert predictor.sigmoid_alpha == 1.0
assert predictor.ratio_c > 0
assert predictor.ratio_c < 64
out_pred = predictor.predict(dtrain)
np.testing.assert_almost_equal(out_pred, expected_pred, decimal=2)
def test_basic(tmpdir, annotation, dataset, use_annotation, quantize, parallel_comp, toolchain):
# pylint: disable=too-many-arguments
"""Test 'ast_native' compiler"""
if dataset == 'letor' and os_platform() == 'windows':
pytest.xfail('export_lib() is too slow for letor on MSVC')
libpath = os.path.join(tmpdir, dataset_db[dataset].libname + _libext())
model = treelite.Model.load(dataset_db[dataset].model, model_format=dataset_db[dataset].format)
annotation_path = os.path.join(tmpdir, 'annotation.json')
if use_annotation:
if annotation[dataset] is None:
pytest.skip('No training data available. Skipping annotation')
with open(annotation_path, 'w') as f:
f.write(annotation[dataset])
params = {
'annotate_in': (annotation_path if use_annotation else 'NULL'),
'quantize': (1 if quantize else 0),
'parallel_comp': (parallel_comp if parallel_comp else 0)
}
model.export_lib(toolchain=toolchain, libpath=libpath, params=params, verbose=True)
predictor = treelite_runtime.Predictor(libpath=libpath, verbose=True)
check_predictor(predictor, dataset)
def test_sparse_categorical_model(tmpdir, quantize, toolchain):
"""
LightGBM is able to produce categorical splits directly, so that
categorical data don't have to be one-hot encoded. Test if Treelite is
able to handle categorical splits.
This example produces a model with high-cardinality categorical variables.
The training data has many missing values, so we need to match LightGBM
when it comes to handling missing values
"""
if toolchain == 'clang':
pytest.xfail(reason='Clang cannot handle long if conditional')
if os_platform() == 'windows':
pytest.xfail(reason='MSVC cannot handle long if conditional')
dataset = 'sparse_categorical'
libpath = os.path.join(tmpdir, dataset_db[dataset].libname + _libext())
model = treelite.Model.load(dataset_db[dataset].model,
model_format=dataset_db[dataset].format)
params = {'quantize': (1 if quantize else 0)}
model.export_lib(toolchain=toolchain,
libpath=libpath,
params=params,
verbose=True,
options=['-O0'])
predictor = treelite_runtime.Predictor(libpath=libpath, verbose=True)
check_predictor(predictor, dataset)
def test_nonlinear_objective(tmpdir, objective, max_label,
expected_global_bias, toolchain, model_format):
# pylint: disable=too-many-locals,too-many-arguments
"""Test non-linear objectives with dummy data"""
np.random.seed(0)
nrow = 16
ncol = 8
X = np.random.randn(nrow, ncol)
y = np.random.randint(0, max_label, size=nrow)
assert np.min(y) == 0
assert np.max(y) == max_label - 1
num_round = 4
dtrain = xgboost.DMatrix(X, y)
bst = xgboost.train({
'objective': objective,
'base_score': 0.5,
'seed': 0
},
dtrain=dtrain,
num_boost_round=num_round)
objective_tag = objective.replace(':', '_')
if model_format == 'json':
model_name = f'nonlinear_{objective_tag}.json'
else:
model_name = f'nonlinear_{objective_tag}.bin'
model_path = os.path.join(tmpdir, model_name)
bst.save_model(model_path)
model = treelite.Model.load(
filename=model_path,
model_format=('xgboost_json' if model_format == 'json' else 'xgboost'))
assert model.num_feature == dtrain.num_col()
assert model.num_output_group == 1
assert model.num_tree == num_round
libpath = os.path.join(tmpdir, objective_tag + _libext())
model.export_lib(toolchain=toolchain,
libpath=libpath,
params={},
verbose=True)
expected_pred_transform = {
'binary:logistic': 'sigmoid',
'count:poisson': 'exponential'
}
predictor = treelite_runtime.Predictor(libpath=libpath, verbose=True)
assert predictor.num_feature == dtrain.num_col()
assert predictor.num_output_group == 1
assert predictor.pred_transform == expected_pred_transform[objective]
np.testing.assert_almost_equal(predictor.global_bias,
expected_global_bias,
decimal=5)
assert predictor.sigmoid_alpha == 1.0
dmat = treelite_runtime.DMatrix(X, dtype='float32')
out_pred = predictor.predict(dmat)
expected_pred = bst.predict(dtrain)
np.testing.assert_almost_equal(out_pred, expected_pred, decimal=5)
def test_sparse_ranking_model(tmpdir, quantize, toolchain):
# pylint: disable=too-many-locals
"""Generate a LightGBM ranking model with highly sparse data.
This example is inspired by https://github.com/dmlc/treelite/issues/222. It verifies that
Treelite is able to accommodate the unique behavior of LightGBM when it comes to handling
missing values.
LightGBM offers two modes of handling missing values:
1. Assign default direction per each test node: This is similar to how XGBoost handles missing
values.
2. Replace missing values with zeroes (0.0): This behavior is unique to LightGBM.
The mode is controlled by the missing_value_to_zero_ field of each test node.
This example is crafted so as to invoke the second mode of missing value handling.
"""
rng = np.random.default_rng(seed=2020)
X = scipy.sparse.random(m=10,
n=206947,
format='csr',
dtype=np.float64,
random_state=0,
density=0.0001)
X.data = rng.standard_normal(size=X.data.shape[0], dtype=np.float64)
y = rng.integers(low=0, high=5, size=X.shape[0])
params = {
'objective': 'lambdarank',
'num_leaves': 32,
'lambda_l1': 0.0,
'lambda_l2': 0.0,
'min_gain_to_split': 0.0,
'learning_rate': 1.0,
'min_data_in_leaf': 1
}
model_path = os.path.join(tmpdir, 'sparse_ranking_lightgbm.txt')
libpath = os.path.join(tmpdir, 'sparse_ranking_lgb' + _libext())
dtrain = lightgbm.Dataset(X, label=y, group=[X.shape[0]])
bst = lightgbm.train(params, dtrain, num_boost_round=1)
lgb_out = bst.predict(X)
bst.save_model(model_path)
model = treelite.Model.load(model_path, model_format='lightgbm')
params = {'quantize': (1 if quantize else 0)}
model.export_lib(toolchain=toolchain,
libpath=libpath,
params=params,
verbose=True)
predictor = treelite_runtime.Predictor(libpath, verbose=True)
dmat = treelite_runtime.DMatrix(X)
out = predictor.predict(dmat)
np.testing.assert_almost_equal(out, lgb_out)
def test_xgb_iris(tmpdir, toolchain, objective, model_format,
expected_pred_transform, num_parallel_tree):
# pylint: disable=too-many-locals, too-many-arguments
"""Test Iris data (multi-class classification)"""
X, y = load_iris(return_X_y=True)
X_train, X_test, y_train, y_test = train_test_split(X,
y,
test_size=0.2,
shuffle=False)
dtrain = xgboost.DMatrix(X_train, label=y_train)
dtest = xgboost.DMatrix(X_test, label=y_test)
num_class = 3
num_round = 10
param = {
'max_depth': 6,
'eta': 0.05,
'num_class': num_class,
'verbosity': 0,
'objective': objective,
'metric': 'mlogloss',
'num_parallel_tree': num_parallel_tree
}
bst = xgboost.train(param,
dtrain,
num_boost_round=num_round,
evals=[(dtrain, 'train'), (dtest, 'test')])
if model_format == 'json':
model_name = 'iris.json'
model_path = os.path.join(tmpdir, model_name)
bst.save_model(model_path)
model = treelite.Model.load(filename=model_path,
model_format='xgboost_json')
else:
model_name = 'iris.model'
model_path = os.path.join(tmpdir, model_name)
bst.save_model(model_path)
model = treelite.Model.load(filename=model_path,
model_format='xgboost')
assert model.num_feature == dtrain.num_col()
assert model.num_class == num_class
assert model.num_tree == num_round * num_class * num_parallel_tree
libpath = os.path.join(tmpdir, 'iris' + _libext())
model.export_lib(toolchain=toolchain,
libpath=libpath,
params={},
verbose=True)
predictor = treelite_runtime.Predictor(libpath=libpath, verbose=True)
assert predictor.num_feature == dtrain.num_col()
assert predictor.num_class == num_class
assert predictor.pred_transform == expected_pred_transform
assert predictor.global_bias == 0.5
assert predictor.sigmoid_alpha == 1.0
dmat = treelite_runtime.DMatrix(X_test, dtype='float32')
out_pred = predictor.predict(dmat)
expected_pred = bst.predict(dtest)
np.testing.assert_almost_equal(out_pred, expected_pred, decimal=5)
def test_treelite_compiled(tmpdir, toy_model, test_data, ref_pred):
libpath = os.path.join(tmpdir, 'mylib' + _libext())
toolchain = os_compatible_toolchains()[0]
toy_model.export_lib(toolchain=toolchain, libpath=libpath)
predictor = treelite_runtime.Predictor(libpath=libpath)
dmat = treelite_runtime.DMatrix(test_data)
pred = predictor.predict(dmat)
np.testing.assert_equal(pred, ref_pred)
def test_single_inst(tmpdir, annotation, dataset, toolchain):
# pylint: disable=too-many-locals
"""Run end-to-end test"""
libpath = os.path.join(tmpdir, dataset_db[dataset].libname + _libext())
model = treelite.Model.load(dataset_db[dataset].model,
model_format=dataset_db[dataset].format)
annotation_path = os.path.join(tmpdir, 'annotation.json')
if annotation[dataset] is None:
annotation_path = None
else:
with open(annotation_path, 'w') as f:
f.write(annotation[dataset])
params = {
'annotate_in': (annotation_path if annotation_path else 'NULL'),
'quantize': 1,
'parallel_comp': model.num_tree
}
model.export_lib(toolchain=toolchain,
libpath=libpath,
params=params,
verbose=True)
predictor = treelite_runtime.Predictor(libpath=libpath, verbose=True)
from sklearn.datasets import load_svmlight_file
X_test, _ = load_svmlight_file(dataset_db[dataset].dtest, zero_based=True)
out_prob = [[] for _ in range(4)]
out_margin = [[] for _ in range(4)]
for i in range(X_test.shape[0]):
x = X_test[i, :]
# Scipy CSR matrix
out_prob[0].append(predictor.predict_instance(x))
out_margin[0].append(predictor.predict_instance(x, pred_margin=True))
# NumPy 1D array with 0 as missing value
x = x.toarray().flatten()
out_prob[1].append(predictor.predict_instance(x, missing=0.0))
out_margin[1].append(
predictor.predict_instance(x, missing=0.0, pred_margin=True))
# NumPy 1D array with np.nan as missing value
np.place(x, x == 0.0, [np.nan])
out_prob[2].append(predictor.predict_instance(x, missing=np.nan))
out_margin[2].append(
predictor.predict_instance(x, missing=np.nan, pred_margin=True))
# NumPy 1D array with np.nan as missing value
# (default when `missing` parameter is unspecified)
out_prob[3].append(predictor.predict_instance(x))
out_margin[3].append(predictor.predict_instance(x, pred_margin=True))
for i in range(4):
check_predictor_output(dataset,
X_test.shape,
out_margin=np.squeeze(np.array(out_margin[i])),
out_prob=np.squeeze(np.array(out_prob[i])))
def test_xgb_boston(tmpdir, toolchain, objective, model_format,
num_parallel_tree):
# pylint: disable=too-many-locals
"""Test Boston data (regression)"""
X, y = load_boston(return_X_y=True)
X_train, X_test, y_train, y_test = train_test_split(X,
y,
test_size=0.2,
shuffle=False)
dtrain = xgboost.DMatrix(X_train, label=y_train)
dtest = xgboost.DMatrix(X_test, label=y_test)
param = {
'max_depth': 8,
'eta': 1,
'silent': 1,
'objective': objective,
'num_parallel_tree': num_parallel_tree
}
num_round = 10
bst = xgboost.train(param,
dtrain,
num_boost_round=num_round,
evals=[(dtrain, 'train'), (dtest, 'test')])
if model_format == 'json':
model_name = 'boston.json'
model_path = os.path.join(tmpdir, model_name)
bst.save_model(model_path)
model = treelite.Model.load(filename=model_path,
model_format='xgboost_json')
else:
model_name = 'boston.model'
model_path = os.path.join(tmpdir, model_name)
bst.save_model(model_path)
model = treelite.Model.load(filename=model_path,
model_format='xgboost')
assert model.num_feature == dtrain.num_col()
assert model.num_class == 1
assert model.num_tree == num_round * num_parallel_tree
libpath = os.path.join(tmpdir, 'boston' + _libext())
model.export_lib(toolchain=toolchain,
libpath=libpath,
params={'parallel_comp': model.num_tree},
verbose=True)
predictor = treelite_runtime.Predictor(libpath=libpath, verbose=True)
assert predictor.num_feature == dtrain.num_col()
assert predictor.num_class == 1
assert predictor.pred_transform == 'identity'
assert predictor.global_bias == 0.5
assert predictor.sigmoid_alpha == 1.0
dmat = treelite_runtime.DMatrix(X_test, dtype='float32')
out_pred = predictor.predict(dmat)
expected_pred = bst.predict(dtest)
np.testing.assert_almost_equal(out_pred, expected_pred, decimal=5)
def test_xgb_dart(tmpdir, toolchain, model_format):
# pylint: disable=too-many-locals,too-many-arguments
"""Test dart booster with dummy data"""
np.random.seed(0)
nrow = 16
ncol = 8
X = np.random.randn(nrow, ncol)
y = np.random.randint(0, 2, size=nrow)
assert np.min(y) == 0
assert np.max(y) == 1
num_round = 50
dtrain = xgboost.DMatrix(X, label=y)
param = {
'booster': 'dart',
'max_depth': 5,
'learning_rate': 0.1,
'objective': 'binary:logistic',
'sample_type': 'uniform',
'normalize_type': 'tree',
'rate_drop': 0.1,
'skip_drop': 0.5
}
bst = xgboost.train(param, dtrain=dtrain, num_boost_round=num_round)
if model_format == 'json':
model_json_path = os.path.join(tmpdir, 'serialized.json')
bst.save_model(model_json_path)
model = treelite.Model.load(model_json_path,
model_format='xgboost_json')
else:
model_bin_path = os.path.join(tmpdir, 'serialized.model')
bst.save_model(model_bin_path)
model = treelite.Model.load(model_bin_path, model_format='xgboost')
assert model.num_feature == dtrain.num_col()
assert model.num_class == 1
assert model.num_tree == num_round
libpath = os.path.join(tmpdir, "dart" + _libext())
model.export_lib(toolchain=toolchain,
libpath=libpath,
params={},
verbose=True)
predictor = treelite_runtime.Predictor(libpath=libpath, verbose=True)
assert predictor.num_feature == dtrain.num_col()
assert predictor.num_class == 1
assert predictor.pred_transform == 'sigmoid'
np.testing.assert_almost_equal(predictor.global_bias, 0, decimal=5)
assert predictor.sigmoid_alpha == 1.0
dmat = treelite_runtime.DMatrix(X, dtype='float32')
out_pred = predictor.predict(dmat)
expected_pred = bst.predict(dtrain)
np.testing.assert_almost_equal(out_pred, expected_pred, decimal=5)
def test_deficient_matrix(tmpdir):
"""Test if Treelite correctly handles sparse matrix with fewer columns than the training data
used for the model. In this case, the matrix should be padded with zeros."""
libpath = os.path.join(tmpdir, dataset_db['mushroom'].libname + _libext())
model = treelite.Model.load(dataset_db['mushroom'].model, model_format='xgboost')
toolchain = os_compatible_toolchains()[0]
model.export_lib(toolchain=toolchain, libpath=libpath, params={'quantize': 1}, verbose=True)
X = csr_matrix(([], ([], [])), shape=(3, 3))
batch = treelite_runtime.Batch.from_csr(X)
predictor = treelite_runtime.Predictor(libpath=libpath, verbose=True)
assert predictor.num_feature == 127
predictor.predict(batch) # should not crash
def test_too_wide_matrix(tmpdir):
"""Test if Treelite correctly handles sparse matrix with more columns than the training data
used for the model. In this case, an exception should be thrown"""
libpath = os.path.join(tmpdir, dataset_db['mushroom'].libname + _libext())
model = treelite.Model.load(dataset_db['mushroom'].model, model_format='xgboost')
toolchain = os_compatible_toolchains()[0]
model.export_lib(toolchain=toolchain, libpath=libpath, params={'quantize': 1}, verbose=True)
X = csr_matrix(([], ([], [])), shape=(3, 1000))
batch = treelite_runtime.Batch.from_csr(X)
predictor = treelite_runtime.Predictor(libpath=libpath, verbose=True)
assert predictor.num_feature == 127
pytest.raises(treelite_runtime.TreeliteRuntimeError, predictor.predict, batch)
def test_xgb_categorical_split(tmpdir, toolchain, quantize, parallel_comp):
"""Test toy XGBoost model with categorical splits"""
dataset = 'xgb_toy_categorical'
model = treelite.Model.load(dataset_db[dataset].model, model_format='xgboost_json')
libpath = os.path.join(tmpdir, dataset_db[dataset].libname + _libext())
params = {
'quantize': (1 if quantize else 0),
'parallel_comp': (parallel_comp if parallel_comp else 0)
}
model.export_lib(toolchain=toolchain, libpath=libpath, params=params, verbose=True)
predictor = treelite_runtime.Predictor(libpath)
check_predictor(predictor, dataset)
def test_lightgbm_multiclass_classification(tmpdir, objective, boosting_type,
toolchain):
# pylint: disable=too-many-locals
"""Test a multi-class classifier"""
model_path = os.path.join(tmpdir, 'iris_lightgbm.txt')
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
X, y = load_iris(return_X_y=True)
X_train, X_test, y_train, y_test = train_test_split(X,
y,
test_size=0.2,
shuffle=False)
dtrain = lightgbm.Dataset(X_train, y_train, free_raw_data=False)
dtest = lightgbm.Dataset(X_test,
y_test,
reference=dtrain,
free_raw_data=False)
param = {
'task': 'train',
'boosting': boosting_type,
'objective': objective,
'metric': 'multi_logloss',
'num_class': 3,
'num_leaves': 31,
'learning_rate': 0.05
}
if boosting_type == 'rf':
param.update({'bagging_fraction': 0.8, 'bagging_freq': 1})
bst = lightgbm.train(param,
dtrain,
num_boost_round=10,
valid_sets=[dtrain, dtest],
valid_names=['train', 'test'])
bst.save_model(model_path)
model = treelite.Model.load(model_path, model_format='lightgbm')
libpath = os.path.join(tmpdir, f'iris_{objective}' + _libext())
model.export_lib(toolchain=toolchain,
libpath=libpath,
params={'quantize': 1},
verbose=True)
predictor = treelite_runtime.Predictor(libpath=libpath, verbose=True)
dmat = treelite_runtime.DMatrix(X_test, dtype='float64')
out_pred = predictor.predict(dmat)
expected_pred = bst.predict(X_test)
np.testing.assert_almost_equal(out_pred, expected_pred, decimal=5)
def test_xgb_iris(tmpdir, toolchain): # pylint: disable=too-many-locals
"""Test Iris data (multi-class classification)"""
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
X, y = load_iris(return_X_y=True)
X_train, X_test, y_train, y_test = train_test_split(X,
y,
test_size=0.2,
shuffle=False)
dtrain = xgboost.DMatrix(X_train, label=y_train)
dtest = xgboost.DMatrix(X_test, label=y_test)
num_class = 3
num_round = 10
param = {
'max_depth': 6,
'eta': 0.05,
'num_class': num_class,
'verbosity': 0,
'objective': 'multi:softmax',
'metric': 'mlogloss'
}
bst = xgboost.train(param,
dtrain,
num_boost_round=num_round,
evals=[(dtrain, 'train'), (dtest, 'test')])
model = treelite.Model.from_xgboost(bst)
assert model.num_feature == dtrain.num_col()
assert model.num_output_group == num_class
assert model.num_tree == num_round * num_class
libpath = os.path.join(tmpdir, 'iris' + _libext())
model.export_lib(toolchain=toolchain,
libpath=libpath,
params={},
verbose=True)
predictor = treelite_runtime.Predictor(libpath=libpath, verbose=True)
assert predictor.num_feature == dtrain.num_col()
assert predictor.num_output_group == num_class
assert predictor.pred_transform == 'max_index'
assert predictor.global_bias == 0.5
assert predictor.sigmoid_alpha == 1.0
batch = treelite_runtime.Batch.from_npy2d(X_test)
out_pred = predictor.predict(batch)
expected_pred = bst.predict(dtest)
np.testing.assert_almost_equal(out_pred, expected_pred, decimal=5)
def test_lightgbm_regression(tmpdir, objective, reg_sqrt, toolchain):
# pylint: disable=too-many-locals
"""Test a regressor"""
model_path = os.path.join(tmpdir, 'boston_lightgbm.txt')
from sklearn.datasets import load_boston
from sklearn.model_selection import train_test_split
X, y = load_boston(return_X_y=True)
X_train, X_test, y_train, y_test = train_test_split(X,
y,
test_size=0.2,
shuffle=False)
dtrain = lightgbm.Dataset(X_train, y_train, free_raw_data=False)
dtest = lightgbm.Dataset(X_test,
y_test,
reference=dtrain,
free_raw_data=False)
param = {
'task': 'train',
'boosting_type': 'gbdt',
'objective': objective,
'reg_sqrt': reg_sqrt,
'metric': 'rmse',
'num_leaves': 31,
'learning_rate': 0.05
}
bst = lightgbm.train(param,
dtrain,
num_boost_round=10,
valid_sets=[dtrain, dtest],
valid_names=['train', 'test'])
bst.save_model(model_path)
model = treelite.Model.load(model_path, model_format='lightgbm')
libpath = os.path.join(tmpdir, f'boston_{objective}' + _libext())
model.export_lib(toolchain=toolchain,
libpath=libpath,
params={'quantize': 1},
verbose=True)
predictor = treelite_runtime.Predictor(libpath=libpath, verbose=True)
dmat = treelite_runtime.DMatrix(X_test, dtype='float64')
out_pred = predictor.predict(dmat)
expected_pred = bst.predict(X_test)
np.testing.assert_almost_equal(out_pred, expected_pred, decimal=5)
def test_nonlinear_objective(tmpdir, objective, max_label, global_bias,
toolchain):
# pylint: disable=too-many-locals
"""Test non-linear objectives with dummy data"""
np.random.seed(0)
nrow = 16
ncol = 8
X = np.random.randn(nrow, ncol)
y = np.random.randint(0, max_label, size=nrow)
assert np.min(y) == 0
assert np.max(y) == max_label - 1
num_round = 4
dtrain = xgboost.DMatrix(X, y)
bst = xgboost.train({'objective': objective},
dtrain=dtrain,
num_boost_round=num_round)
model = treelite.Model.from_xgboost(bst)
assert model.num_feature == dtrain.num_col()
assert model.num_output_group == 1
assert model.num_tree == num_round
libpath = os.path.join(tmpdir,
f'{objective.replace(":", "_")}' + _libext())
model.export_lib(toolchain=toolchain,
libpath=libpath,
params={},
verbose=True)
expected_pred_transform = {
'binary:logistic': 'sigmoid',
'count:poisson': 'exponential'
}
predictor = treelite_runtime.Predictor(libpath=libpath, verbose=True)
assert predictor.num_feature == dtrain.num_col()
assert predictor.num_output_group == 1
assert predictor.pred_transform == expected_pred_transform[objective]
np.testing.assert_almost_equal(predictor.global_bias,
global_bias,
decimal=5)
assert predictor.sigmoid_alpha == 1.0
batch = treelite_runtime.Batch.from_npy2d(X)
out_pred = predictor.predict(batch)
expected_pred = bst.predict(dtrain)
np.testing.assert_almost_equal(out_pred, expected_pred, decimal=5)
def test_skl_converter_multiclass_classifier(tmpdir, clazz, toolchain):
# pylint: disable=too-many-locals
"""Convert scikit-learn multi-class classifier"""
if clazz == RandomForestClassifier:
pytest.xfail(
reason='Need to use float64 thresholds to obtain correct result')
X, y = load_iris(return_X_y=True)
kwargs = {}
if clazz == GradientBoostingClassifier:
kwargs['init'] = 'zero'
clf = clazz(max_depth=3, random_state=0, n_estimators=10, **kwargs)
clf.fit(X, y)
expected_prob = clf.predict_proba(X)
model = treelite.sklearn.import_model(clf)
assert model.num_feature == clf.n_features_
assert model.num_output_group == clf.n_classes_
assert (model.num_tree == clf.n_estimators *
(clf.n_classes_ if clazz == GradientBoostingClassifier else 1))
dtrain = treelite.DMatrix(X)
annotation_path = os.path.join(tmpdir, 'annotation.json')
annotator = treelite.Annotator()
annotator.annotate_branch(model=model, dmat=dtrain, verbose=True)
annotator.save(path=annotation_path)
libpath = os.path.join(tmpdir, 'skl' + _libext())
model.export_lib(toolchain=toolchain,
libpath=libpath,
params={'annotate_in': annotation_path},
verbose=True)
predictor = treelite_runtime.Predictor(libpath=libpath)
assert predictor.num_feature == clf.n_features_
assert predictor.num_output_group == clf.n_classes_
assert (predictor.pred_transform == ('softmax'
if clazz == GradientBoostingClassifier
else 'identity_multiclass'))
assert predictor.global_bias == 0.0
assert predictor.sigmoid_alpha == 1.0
batch = treelite_runtime.Batch.from_npy2d(X)
out_prob = predictor.predict(batch)
np.testing.assert_almost_equal(out_prob, expected_prob)
def test_failsafe_compiler(tmpdir, dataset, use_elf, toolchain):
"""Test 'failsafe' compiler"""
libpath = os.path.join(tmpdir, dataset_db[dataset].libname + _libext())
model = treelite.Model.load(dataset_db[dataset].model, model_format=dataset_db[dataset].format)
params = {'dump_array_as_elf': (1 if use_elf else 0)}
is_linux = sys.platform.startswith('linux')
# Expect Treelite to throw error if we try to use dump_array_as_elf on non-Linux OS
# Also, failsafe compiler is only available for XGBoost models
if ((not is_linux) and use_elf) or dataset_db[dataset].format != 'xgboost':
expect_raises = pytest.raises(treelite.TreeliteError)
else:
expect_raises = does_not_raise()
with expect_raises:
model.export_lib(compiler='failsafe', toolchain=toolchain, libpath=libpath, params=params,
verbose=True)
predictor = treelite_runtime.Predictor(libpath=libpath, verbose=True)
check_predictor(predictor, dataset)
def test_skl_converter_multiclass_classifier(tmpdir, import_method, clazz,
toolchain):
# pylint: disable=too-many-locals
"""Convert scikit-learn multi-class classifier"""
X, y = load_iris(return_X_y=True)
kwargs = {}
if clazz == GradientBoostingClassifier:
kwargs['init'] = 'zero'
clf = clazz(max_depth=3, random_state=0, n_estimators=10, **kwargs)
clf.fit(X, y)
expected_prob = clf.predict_proba(X)
if import_method == 'import_new':
model = treelite.sklearn.import_model(clf)
else:
model = treelite.sklearn.import_model_with_model_builder(clf)
assert model.num_feature == clf.n_features_
assert model.num_class == clf.n_classes_
assert (model.num_tree == clf.n_estimators *
(clf.n_classes_ if clazz == GradientBoostingClassifier else 1))
dtrain = treelite_runtime.DMatrix(X, dtype='float64')
annotation_path = os.path.join(tmpdir, 'annotation.json')
annotator = treelite.Annotator()
annotator.annotate_branch(model=model, dmat=dtrain, verbose=True)
annotator.save(path=annotation_path)
libpath = os.path.join(tmpdir, 'skl' + _libext())
model.export_lib(toolchain=toolchain,
libpath=libpath,
params={'annotate_in': annotation_path},
verbose=True)
predictor = treelite_runtime.Predictor(libpath=libpath)
assert predictor.num_feature == clf.n_features_
assert predictor.num_class == clf.n_classes_
assert (predictor.pred_transform == ('softmax'
if clazz == GradientBoostingClassifier
else 'identity_multiclass'))
assert predictor.global_bias == 0.0
assert predictor.sigmoid_alpha == 1.0
out_prob = predictor.predict(dtrain)
np.testing.assert_almost_equal(out_prob, expected_prob)
def test_round_trip(tmpdir, dataset, quantize, toolchain):
"""Perform round-trip tests"""
libpath = os.path.join(tmpdir, dataset_db[dataset].libname + _libext())
pb_path = os.path.join(tmpdir, 'my.buffer')
model = treelite.Model.load(dataset_db[dataset].model,
model_format=dataset_db[dataset].format)
model.export_protobuf(pb_path)
model2 = treelite.Model.load(pb_path, model_format='protobuf')
params = {
'quantize': (1 if quantize else 0),
'parallel_comp': model.num_tree
}
model2.export_lib(toolchain=toolchain,
libpath=libpath,
params=params,
verbose=True)
predictor = treelite_runtime.Predictor(libpath=libpath, verbose=True)
check_predictor(predictor, dataset)
def test_lightgbm_binary_classification(tmpdir, objective, toolchain):
# pylint: disable=too-many-locals
"""Test a binary classifier"""
dataset = 'mushroom'
model_path = os.path.join(tmpdir, 'mushroom_lightgbm.txt')
dtest_path = dataset_db[dataset].dtest
dtrain = lightgbm.Dataset(dataset_db[dataset].dtrain)
dtest = lightgbm.Dataset(dtest_path, reference=dtrain)
param = {
'task': 'train',
'boosting_type': 'gbdt',
'objective': objective,
'metric': 'auc',
'num_leaves': 7,
'learning_rate': 0.1
}
bst = lightgbm.train(param,
dtrain,
num_boost_round=10,
valid_sets=[dtrain, dtest],
valid_names=['train', 'test'])
bst.save_model(model_path)
expected_prob = bst.predict(dtest_path)
expected_margin = bst.predict(dtest_path, raw_score=True)
model = treelite.Model.load(model_path, model_format='lightgbm')
libpath = os.path.join(tmpdir, f'agaricus_{objective}' + _libext())
dmat = treelite_runtime.DMatrix(load_svmlight_file(dtest_path,
zero_based=True)[0],
dtype='float64')
model.export_lib(toolchain=toolchain,
libpath=libpath,
params={},
verbose=True)
predictor = treelite_runtime.Predictor(libpath, verbose=True)
out_prob = predictor.predict(dmat)
np.testing.assert_almost_equal(out_prob, expected_prob, decimal=5)
out_margin = predictor.predict(dmat, pred_margin=True)
np.testing.assert_almost_equal(out_margin, expected_margin, decimal=5)
def test_code_folding(tmpdir, annotation, dataset, toolchain,
code_folding_factor):
"""Test suite for testing code folding feature"""
if dataset == 'letor' and os_platform() == 'windows':
pytest.xfail('export_lib() is too slow for letor on MSVC')
libpath = os.path.join(tmpdir, dataset_db[dataset].libname + _libext())
model = treelite.Model.load(dataset_db[dataset].model,
model_format=dataset_db[dataset].format)
annotation_path = os.path.join(tmpdir, 'annotation.json')
if annotation[dataset] is None:
annotation_path = None
else:
with open(annotation_path, 'w') as f:
f.write(annotation[dataset])
params = {
'annotate_in': (annotation_path if annotation_path else 'NULL'),
'quantize': 1,
'parallel_comp': model.num_tree,
'code_folding_req': code_folding_factor
}
# For the LightGBM model, we expect this error:
# !t3->convert_missing_to_zero: Code folding not supported, because a categorical split is
# supposed to convert missing values into zeros, and this is not possible with current code
# folding implementation.
if dataset == 'toy_categorical' and code_folding_factor < 2.0:
expect_raises = pytest.raises(treelite.TreeliteError)
else:
expect_raises = does_not_raise()
with expect_raises:
model.export_lib(toolchain=toolchain,
libpath=libpath,
params=params,
verbose=True)
predictor = treelite_runtime.Predictor(libpath=libpath, verbose=True)
check_predictor(predictor, dataset)
def test_xgb_boston(tmpdir, toolchain): # pylint: disable=too-many-locals
"""Test Boston data (regression)"""
from sklearn.datasets import load_boston
from sklearn.model_selection import train_test_split
X, y = load_boston(return_X_y=True)
X_train, X_test, y_train, y_test = train_test_split(X,
y,
test_size=0.2,
shuffle=False)
dtrain = xgboost.DMatrix(X_train, label=y_train)
dtest = xgboost.DMatrix(X_test, label=y_test)
param = {'max_depth': 8, 'eta': 1, 'silent': 1, 'objective': 'reg:linear'}
num_round = 10
bst = xgboost.train(param,
dtrain,
num_boost_round=num_round,
evals=[(dtrain, 'train'), (dtest, 'test')])
model = treelite.Model.from_xgboost(bst)
assert model.num_feature == dtrain.num_col()
assert model.num_output_group == 1
assert model.num_tree == num_round
libpath = os.path.join(tmpdir, 'boston' + _libext())
model.export_lib(toolchain=toolchain,
libpath=libpath,
params={'parallel_comp': model.num_tree},
verbose=True)
predictor = treelite_runtime.Predictor(libpath=libpath, verbose=True)
assert predictor.num_feature == dtrain.num_col()
assert predictor.num_output_group == 1
assert predictor.pred_transform == 'identity'
assert predictor.global_bias == 0.5
assert predictor.sigmoid_alpha == 1.0
batch = treelite_runtime.Batch.from_npy2d(X_test)
out_pred = predictor.predict(batch)
expected_pred = bst.predict(dtest)
np.testing.assert_almost_equal(out_pred, expected_pred, decimal=5)
def test_nan_handling_with_categorical_splits(tmpdir, toolchain):
"""Test that NaN inputs are handled correctly in categorical splits"""
# Test case taken from https://github.com/dmlc/treelite/issues/277
X = np.array(30 * [[1]] + 30 * [[2]] + 30 * [[0]])
y = np.array(60 * [5] + 30 * [10])
train_data = lightgbm.Dataset(X, label=y, categorical_feature=[0])
bst = lightgbm.train({}, train_data, 1)
model_path = os.path.join(tmpdir, 'dummy_categorical.txt')
libpath = os.path.join(tmpdir, 'dummy_categorical_lgb' + _libext())
input_with_nan = np.array([[np.NaN], [0.0]])
lgb_pred = bst.predict(input_with_nan)
bst.save_model(model_path)
model = treelite.Model.load(model_path, model_format='lightgbm')
model.export_lib(toolchain=toolchain, libpath=libpath)
predictor = treelite_runtime.Predictor(libpath)
dmat = treelite_runtime.DMatrix(input_with_nan)
tl_pred = predictor.predict(dmat)
np.testing.assert_almost_equal(tl_pred, lgb_pred)
def test_categorical_data(tmpdir, quantize, parallel_comp, toolchain):
"""
LightGBM is able to produce categorical splits directly, so that
categorical data don't have to be one-hot encoded. Test if Treelite is
able to handle categorical splits.
This toy example contains two features, both of which are categorical.
The first has cardinality 3 and the second 5. The label was generated using
the formula
y = f(x0) + g(x1) + [noise with std=0.1]
where f and g are given by the tables
x0 f(x0) x1 g(x1)
0 -20 0 -2
1 -10 1 -1
2 0 2 0
3 1
4 2
"""
dataset = 'toy_categorical'
libpath = os.path.join(tmpdir, dataset_db[dataset].libname + _libext())
model = treelite.Model.load(dataset_db[dataset].model,
model_format=dataset_db[dataset].format)
params = {
'quantize': (1 if quantize else 0),
'parallel_comp': (parallel_comp if parallel_comp else 0)
}
model.export_lib(toolchain=toolchain,
libpath=libpath,
params=params,
verbose=True)
predictor = treelite_runtime.Predictor(libpath=libpath, verbose=True)
check_predictor(predictor, dataset)
