Skip to content

GitLab

Unverified Commit ab604263 authored by Nikita Titov's avatar Nikita Titov Committed by GitHub
Browse files

[ci][python-package] add testing for the shape of raw score predictions (#6765) 


Co-authored-by: default avatarJames Lamb <jaylamb20@gmail.com>
parent 799199a8
Hide whitespace changes
Inline Side-by-side
Showing with 6 additions and 0 deletions
tests/python_package_test/test_engine.py
View file @ ab604263
...@@ -3909,12 +3909,14 @@ def test_predict_regression_output_shape(): ...@@ -3909,12 +3909,14 @@ def test_predict_regression_output_shape():
# 1-round model # 1-round model
bst = lgb.train(params, dtrain, num_boost_round=1) bst = lgb.train(params, dtrain, num_boost_round=1)
assert bst.predict(X).shape == (n_samples,) assert bst.predict(X).shape == (n_samples,)
assert bst.predict(X, raw_score=True).shape == (n_samples,)
assert bst.predict(X, pred_contrib=True).shape == (n_samples, n_features + 1) assert bst.predict(X, pred_contrib=True).shape == (n_samples, n_features + 1)
assert bst.predict(X, pred_leaf=True).shape == (n_samples, 1) assert bst.predict(X, pred_leaf=True).shape == (n_samples, 1)
# 2-round model # 2-round model
bst = lgb.train(params, dtrain, num_boost_round=2) bst = lgb.train(params, dtrain, num_boost_round=2)
assert bst.predict(X).shape == (n_samples,) assert bst.predict(X).shape == (n_samples,)
assert bst.predict(X, raw_score=True).shape == (n_samples,)
assert bst.predict(X, pred_contrib=True).shape == (n_samples, n_features + 1) assert bst.predict(X, pred_contrib=True).shape == (n_samples, n_features + 1)
assert bst.predict(X, pred_leaf=True).shape == (n_samples, 2) assert bst.predict(X, pred_leaf=True).shape == (n_samples, 2)
...@@ -3929,12 +3931,14 @@ def test_predict_binary_classification_output_shape(): ...@@ -3929,12 +3931,14 @@ def test_predict_binary_classification_output_shape():
# 1-round model # 1-round model
bst = lgb.train(params, dtrain, num_boost_round=1) bst = lgb.train(params, dtrain, num_boost_round=1)
assert bst.predict(X).shape == (n_samples,) assert bst.predict(X).shape == (n_samples,)
assert bst.predict(X, raw_score=True).shape == (n_samples,)
assert bst.predict(X, pred_contrib=True).shape == (n_samples, n_features + 1) assert bst.predict(X, pred_contrib=True).shape == (n_samples, n_features + 1)
assert bst.predict(X, pred_leaf=True).shape == (n_samples, 1) assert bst.predict(X, pred_leaf=True).shape == (n_samples, 1)
# 2-round model # 2-round model
bst = lgb.train(params, dtrain, num_boost_round=2) bst = lgb.train(params, dtrain, num_boost_round=2)
assert bst.predict(X).shape == (n_samples,) assert bst.predict(X).shape == (n_samples,)
assert bst.predict(X, raw_score=True).shape == (n_samples,)
assert bst.predict(X, pred_contrib=True).shape == (n_samples, n_features + 1) assert bst.predict(X, pred_contrib=True).shape == (n_samples, n_features + 1)
assert bst.predict(X, pred_leaf=True).shape == (n_samples, 2) assert bst.predict(X, pred_leaf=True).shape == (n_samples, 2)
...@@ -3950,12 +3954,14 @@ def test_predict_multiclass_classification_output_shape(): ...@@ -3950,12 +3954,14 @@ def test_predict_multiclass_classification_output_shape():
# 1-round model # 1-round model
bst = lgb.train(params, dtrain, num_boost_round=1) bst = lgb.train(params, dtrain, num_boost_round=1)
assert bst.predict(X).shape == (n_samples, n_classes) assert bst.predict(X).shape == (n_samples, n_classes)
assert bst.predict(X, raw_score=True).shape == (n_samples, n_classes)
assert bst.predict(X, pred_contrib=True).shape == (n_samples, n_classes * (n_features + 1)) assert bst.predict(X, pred_contrib=True).shape == (n_samples, n_classes * (n_features + 1))
assert bst.predict(X, pred_leaf=True).shape == (n_samples, n_classes) assert bst.predict(X, pred_leaf=True).shape == (n_samples, n_classes)
# 2-round model # 2-round model
bst = lgb.train(params, dtrain, num_boost_round=2) bst = lgb.train(params, dtrain, num_boost_round=2)
assert bst.predict(X).shape == (n_samples, n_classes) assert bst.predict(X).shape == (n_samples, n_classes)
assert bst.predict(X, raw_score=True).shape == (n_samples, n_classes)
assert bst.predict(X, pred_contrib=True).shape == (n_samples, n_classes * (n_features + 1)) assert bst.predict(X, pred_contrib=True).shape == (n_samples, n_classes * (n_features + 1))
assert bst.predict(X, pred_leaf=True).shape == (n_samples, n_classes * 2) assert bst.predict(X, pred_leaf=True).shape == (n_samples, n_classes * 2)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment