Hendrik/xgboost
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Fix typo in dask interface. (#6240)

Browse Source
This commit is contained in:
Jiaming Yuan
2020-10-15 15:26:29 +08:00
committed by GitHub
parent 06e453ddf4
commit 3da5a69dc9
2 changed files with 21 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
tests/python/test_with_dask.py
View File

@@ -40,10 +40,13 @@ kCols = 10
kWorkers = 5
def generate_array():
def generate_array(with_weights=False):
partition_size = 20
X = da.random.random((kRows, kCols), partition_size)
y = da.random.random(kRows, partition_size)
if with_weights:
w = da.random.random(kRows, partition_size)
return X, y, w
return X, y
@@ -252,11 +255,11 @@ def test_dask_missing_value_cls():
def test_dask_regressor():
with LocalCluster(n_workers=kWorkers) as cluster:
with Client(cluster) as client:
X, y = generate_array()
X, y, w = generate_array(with_weights=True)
regressor = xgb.dask.DaskXGBRegressor(verbosity=1, n_estimators=2)
regressor.set_params(tree_method='hist')
regressor.client = client
regressor.fit(X, y, eval_set=[(X, y)])
regressor.fit(X, y, sample_weight=w, eval_set=[(X, y)])
prediction = regressor.predict(X)
assert prediction.ndim == 1
@@ -274,12 +277,12 @@ def test_dask_regressor():
def test_dask_classifier():
with LocalCluster(n_workers=kWorkers) as cluster:
with Client(cluster) as client:
X, y = generate_array()
X, y, w = generate_array(with_weights=True)
y = (y * 10).astype(np.int32)
classifier = xgb.dask.DaskXGBClassifier(
verbosity=1, n_estimators=2, eval_metric='merror')
classifier.client = client
classifier.fit(X, y, eval_set=[(X, y)])
classifier.fit(X, y, sample_weight=w, eval_set=[(X, y)])
prediction = classifier.predict(X)
assert prediction.ndim == 1