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 check_predictor(predictor, dataset):
"""Check whether a predictor produces correct predictions for a given dataset"""
dmat = treelite_runtime.DMatrix(dataset_db[dataset].dtest,
dtype=dataset_db[dataset].dtype)
out_margin = predictor.predict(dmat, pred_margin=True)
out_prob = predictor.predict(dmat)
check_predictor_output(dataset, dmat.shape, out_margin, out_prob)
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_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_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_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_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))
dmat = treelite_runtime.DMatrix(X, dtype='float32')
predictor = treelite_runtime.Predictor(libpath=libpath, verbose=True)
assert predictor.num_feature == 127
pytest.raises(treelite_runtime.TreeliteRuntimeError, predictor.predict, dmat)
def compute_annotation(dataset):
model = treelite.Model.load(
dataset_db[dataset].model,
model_format=dataset_db[dataset].format)
if dataset_db[dataset].dtrain is None:
return None
dtrain = treelite_runtime.DMatrix(dataset_db[dataset].dtrain)
annotator = treelite.Annotator()
annotator.annotate_branch(model=model, dmat=dtrain, verbose=True)
annotation_path = os.path.join(tmpdir, f'{dataset}.json')
annotator.save(annotation_path)
with open(annotation_path, 'r') as f:
return f.read()
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))
dmat = treelite_runtime.DMatrix(X, dtype='float32')
predictor = treelite_runtime.Predictor(libpath=libpath, verbose=True)
assert predictor.num_feature == 127
predictor.predict(dmat) # should not crash
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_class == 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_class == 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]])
dmat = treelite_runtime.DMatrix(x, dtype='float32')
pred = predictor.predict(dmat)
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_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_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_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_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_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_model_builder(tmpdir, use_annotation, quantize, toolchain,
test_round_trip):
# pylint: disable=R0914,R0915
"""A simple model"""
num_feature = 127
pred_transform = 'sigmoid'
builder = treelite.ModelBuilder(num_feature=num_feature,
average_tree_output=False,
pred_transform=pred_transform)
# Build mushroom model
tree = treelite.ModelBuilder.Tree()
tree[0].set_numerical_test_node(feature_id=29,
opname='<',
threshold=-9.53674316e-07,
default_left=True,
left_child_key=1,
right_child_key=2)
tree[1].set_numerical_test_node(feature_id=56,
opname='<',
threshold=-9.53674316e-07,
default_left=True,
left_child_key=3,
right_child_key=4)
tree[3].set_numerical_test_node(feature_id=60,
opname='<',
threshold=-9.53674316e-07,
default_left=True,
left_child_key=7,
right_child_key=8)
tree[7].set_leaf_node(leaf_value=1.89899647)
tree[8].set_leaf_node(leaf_value=-1.94736838)
tree[4].set_numerical_test_node(feature_id=21,
opname='<',
threshold=-9.53674316e-07,
default_left=True,
left_child_key=9,
right_child_key=10)
tree[9].set_leaf_node(leaf_value=1.78378379)
tree[10].set_leaf_node(leaf_value=-1.98135197)
tree[2].set_numerical_test_node(feature_id=109,
opname='<',
threshold=-9.53674316e-07,
default_left=True,
left_child_key=5,
right_child_key=6)
tree[5].set_numerical_test_node(feature_id=67,
opname='<',
threshold=-9.53674316e-07,
default_left=True,
left_child_key=11,
right_child_key=12)
tree[11].set_leaf_node(leaf_value=-1.9854598)
tree[12].set_leaf_node(leaf_value=0.938775539)
tree[6].set_leaf_node(leaf_value=1.87096775)
tree[0].set_root()
builder.append(tree)
tree = treelite.ModelBuilder.Tree()
tree[0].set_numerical_test_node(feature_id=29,
opname='<',
threshold=-9.53674316e-07,
default_left=True,
left_child_key=1,
right_child_key=2)
tree[1].set_numerical_test_node(feature_id=21,
opname='<',
threshold=-9.53674316e-07,
default_left=True,
left_child_key=3,
right_child_key=4)
tree[3].set_leaf_node(leaf_value=1.14607906)
tree[4].set_numerical_test_node(feature_id=36,
opname='<',
threshold=-9.53674316e-07,
default_left=True,
left_child_key=7,
right_child_key=8)
tree[7].set_leaf_node(leaf_value=-6.87994671)
tree[8].set_leaf_node(leaf_value=-0.10659159)
tree[2].set_numerical_test_node(feature_id=109,
opname='<',
threshold=-9.53674316e-07,
default_left=True,
left_child_key=5,
right_child_key=6)
tree[5].set_numerical_test_node(feature_id=39,
opname='<',
threshold=-9.53674316e-07,
default_left=True,
left_child_key=9,
right_child_key=10)
tree[9].set_leaf_node(leaf_value=-0.0930657759)
tree[10].set_leaf_node(leaf_value=-1.15261209)
tree[6].set_leaf_node(leaf_value=1.00423074)
tree[0].set_root()
builder.append(tree)
model = builder.commit()
if test_round_trip:
checkpoint_path = os.path.join(tmpdir, 'checkpoint.bin')
model.serialize(checkpoint_path)
model = treelite.Model.deserialize(checkpoint_path)
assert model.num_feature == num_feature
assert model.num_class == 1
assert model.num_tree == 2
annotation_path = os.path.join(tmpdir, 'annotation.json')
if use_annotation:
dtrain = treelite_runtime.DMatrix(load_svmlight_file(
dataset_db['mushroom'].dtrain, zero_based=True)[0],
dtype='float32')
annotator = treelite.Annotator()
annotator.annotate_branch(model=model, dmat=dtrain, verbose=True)
annotator.save(path=annotation_path)
params = {
'annotate_in': (annotation_path if use_annotation else 'NULL'),
'quantize': (1 if quantize else 0),
'parallel_comp': model.num_tree
}
libpath = os.path.join(tmpdir, 'mushroom' + _libext())
model.export_lib(toolchain=toolchain,
libpath=libpath,
params=params,
verbose=True)
predictor = treelite_runtime.Predictor(libpath=libpath, verbose=True)
check_predictor(predictor, 'mushroom')
def test_categorical_data_pandas_df_with_dummies(tmpdir, toolchain):
# pylint: disable=too-many-locals
"""
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
Unlike test_categorical_data(), this test will convert the categorical features into
dummies using OneHotEncoder and store them into a Pandas Dataframe.
"""
import pandas as pd
rng = np.random.default_rng(seed=0)
n_row = 100
x0 = rng.integers(low=0, high=3, size=n_row)
x1 = rng.integers(low=0, high=5, size=n_row)
noise = rng.standard_normal(size=n_row) * 0.1
y = (x0 * 10 - 20) + (x1 - 2) + noise
df = pd.DataFrame({'x0': x0, 'x1': x1}).astype('category')
df_dummies = pd.get_dummies(df, prefix=['x0', 'x1'], prefix_sep='=')
assert df_dummies.columns.tolist() == [
'x0=0', 'x0=1', 'x0=2', 'x1=0', 'x1=1', 'x1=2', 'x1=3', 'x1=4'
]
model_path = os.path.join(tmpdir, 'toy_dummies.txt')
libpath = os.path.join(tmpdir, 'dummies' + _libext())
dtrain = lightgbm.Dataset(df_dummies, label=y)
param = {
'task': 'train',
'boosting_type': 'gbdt',
'objective': 'regression',
'metric': 'rmse',
'num_leaves': 7,
'learning_rate': 0.1
}
bst = lightgbm.train(param,
dtrain,
num_boost_round=15,
valid_sets=[dtrain],
valid_names=['train'])
lgb_out = bst.predict(df_dummies)
bst.save_model(model_path)
model = treelite.Model.load(model_path, model_format='lightgbm')
model.export_lib(toolchain=toolchain,
libpath=libpath,
params={},
verbose=True)
predictor = treelite_runtime.Predictor(libpath=libpath, verbose=True)
tl_out = predictor.predict(
treelite_runtime.DMatrix(df_dummies.values, dtype='float32'))
np.testing.assert_almost_equal(lgb_out, tl_out, decimal=5)
def test_xgb_deserializers(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
# Train xgboost model
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')])
# Serialize xgboost model
model_bin_path = os.path.join(tmpdir, 'serialized.model')
bst.save_model(model_bin_path)
model_json_path = os.path.join(tmpdir, 'serialized.json')
bst.save_model(model_json_path)
# Construct Treelite models from xgboost serializations
model_bin = treelite.Model.load(
model_bin_path, model_format='xgboost'
)
model_json = treelite.Model.load(
model_json_path, model_format='xgboost_json'
)
with open(model_json_path) as file_:
json_str = file_.read()
model_json_str = treelite.Model.from_xgboost_json(json_str)
# Compile models to libraries
model_bin_lib = os.path.join(tmpdir, 'bin{}'.format(_libext()))
model_bin.export_lib(
toolchain=toolchain,
libpath=model_bin_lib,
params={'parallel_comp': model_bin.num_tree}
)
model_json_lib = os.path.join(tmpdir, 'json{}'.format(_libext()))
model_json.export_lib(
toolchain=toolchain,
libpath=model_json_lib,
params={'parallel_comp': model_json.num_tree}
)
model_json_str_lib = os.path.join(tmpdir, 'json_str{}'.format(_libext()))
model_json_str.export_lib(
toolchain=toolchain,
libpath=model_json_str_lib,
params={'parallel_comp': model_json_str.num_tree}
)
# Generate predictors from compiled libraries
predictor_bin = treelite_runtime.Predictor(model_bin_lib)
assert predictor_bin.num_feature == dtrain.num_col()
assert predictor_bin.num_class == 1
assert predictor_bin.pred_transform == 'identity'
assert predictor_bin.global_bias == pytest.approx(0.5)
assert predictor_bin.sigmoid_alpha == pytest.approx(1.0)
predictor_json = treelite_runtime.Predictor(model_json_lib)
assert predictor_json.num_feature == dtrain.num_col()
assert predictor_json.num_class == 1
assert predictor_json.pred_transform == 'identity'
assert predictor_json.global_bias == pytest.approx(0.5)
assert predictor_json.sigmoid_alpha == pytest.approx(1.0)
predictor_json_str = treelite_runtime.Predictor(model_json_str_lib)
assert predictor_json_str.num_feature == dtrain.num_col()
assert predictor_json_str.num_class == 1
assert predictor_json_str.pred_transform == 'identity'
assert predictor_json_str.global_bias == pytest.approx(0.5)
assert predictor_json_str.sigmoid_alpha == pytest.approx(1.0)
# Run inference with each predictor
dmat = treelite_runtime.DMatrix(X_test, dtype='float32')
bin_pred = predictor_bin.predict(dmat)
json_pred = predictor_json.predict(dmat)
json_str_pred = predictor_json_str.predict(dmat)
expected_pred = bst.predict(dtest)
np.testing.assert_almost_equal(bin_pred, expected_pred, decimal=5)
np.testing.assert_almost_equal(json_pred, expected_pred, decimal=5)
np.testing.assert_almost_equal(json_str_pred, expected_pred, decimal=5)
# X_test[0,1:] = np.where(np.isnan(X_test[0,1:]), tofill, X_test[0,1:])
# batch = treelite_runtime.DMatrix(X_test) #shape(1,130)
# y_preds = int((predictor.predict(batch))>0.5)
# res.append(y_preds)
##sum(res)
##output: 6396
# In[ ]:
#调用janestreet API
import janestreet
env = janestreet.make_env() # initialize the environment
iter_test = env.iter_test() # an iterator which loops over the test set
tofill = f_mean.values.reshape((1, -1))
for (test_df, sample_prediction_df) in iter_test:
if test_df['weight'].values[0] == 0:
sample_prediction_df.action = 0
else:
X_test = test_df.loc[:, features].values
if np.isnan(X_test.sum()):
X_test[0, 1:] = np.where(np.isnan(X_test[0, 1:]), tofill,
X_test[0, 1:])
batch = treelite_runtime.DMatrix(X_test)
y_preds = int((predictor.predict(batch)) > 0.5)
sample_prediction_df.action = y_preds
env.predict(sample_prediction_df)
# In[ ]:
def _treelite_format_hook(data):
"""Helper function converting data into treelite format"""
return treelite_runtime.DMatrix(data[0]), data[1]
import treelite
import treelite_runtime
from fast_svmlight_loader import load_svmlight
import sys
import os
import numpy as np
import time
datafile = sys.argv[1]
libfile = sys.argv[2]
modfile = sys.argv[3]
print('Loading data file {}'.format(datafile))
start = time.time()
dmat = load_svmlight(filename=datafile, verbose=False)
end = time.time()
print('Done loading data file {} in {} sec'.format(datafile, end - start))
X = dmat['data']
predictor = treelite_runtime.Predictor(libfile,
verbose=True,
include_master_thread=True)
nrow = X.shape[0]
for batchsize in np.logspace(np.log10(100), 5, 10).astype(np.int):
print('*** batchsize = {}'.format(batchsize))
for i in range(300):
rbegin = np.random.randint(0, nrow - batchsize + 1)
rend = rbegin + batchsize
dmat = treelite_runtime.DMatrix(X[rbegin:rend])
predictor.predict(dmat, pred_margin=True)