Fix inplace predict missing value. (#6787)

This commit is contained in:
Jiaming Yuan
2021-03-27 05:36:10 +08:00
committed by GitHub
parent 5c87c2bba8
commit a59c7323b4
8 changed files with 97 additions and 33 deletions

View File

@@ -103,31 +103,37 @@ class TestInplacePredict:
'''Tests for running inplace prediction'''
@classmethod
def setup_class(cls):
cls.rows = 100
cls.rows = 1000
cls.cols = 10
cls.missing = 11 # set to integer for testing
cls.rng = np.random.RandomState(1994)
cls.X = cls.rng.randn(cls.rows, cls.cols)
missing_idx = [i for i in range(0, cls.cols, 4)]
cls.X[:, missing_idx] = cls.missing # set to be missing
cls.y = cls.rng.randn(cls.rows)
dtrain = xgb.DMatrix(cls.X, cls.y)
cls.test = xgb.DMatrix(cls.X[:10, ...], missing=cls.missing)
cls.booster = xgb.train({'tree_method': 'hist'}, dtrain, num_boost_round=10)
cls.test = xgb.DMatrix(cls.X[:10, ...])
def test_predict(self):
booster = self.booster
X = self.X
test = self.test
predt_from_array = booster.inplace_predict(X[:10, ...])
predt_from_array = booster.inplace_predict(X[:10, ...], missing=self.missing)
predt_from_dmatrix = booster.predict(test)
np.testing.assert_allclose(predt_from_dmatrix, predt_from_array)
predt_from_array = booster.inplace_predict(X[:10, ...], iteration_range=(0, 4))
predt_from_array = booster.inplace_predict(
X[:10, ...], iteration_range=(0, 4), missing=self.missing
)
predt_from_dmatrix = booster.predict(test, ntree_limit=4)
np.testing.assert_allclose(predt_from_dmatrix, predt_from_array)