wehub-resource-sync
51 lines
1.4 KiB
Python
51 lines
1.4 KiB
Python
from pathlib import Path
|
|
|
|
import dask.array as da
|
|
import numpy as np
|
|
from distributed import Client, LocalCluster
|
|
from sklearn.datasets import load_svmlight_file
|
|
|
|
import lightgbm as lgb
|
|
|
|
if __name__ == "__main__":
|
|
print("loading data")
|
|
|
|
rank_example_dir = Path(__file__).absolute().parents[2] / "lambdarank"
|
|
X, y = load_svmlight_file(str(rank_example_dir / "rank.train"))
|
|
group = np.loadtxt(str(rank_example_dir / "rank.train.query"))
|
|
|
|
print("initializing a Dask cluster")
|
|
|
|
cluster = LocalCluster(n_workers=2)
|
|
client = Client(cluster)
|
|
|
|
print("created a Dask LocalCluster")
|
|
|
|
print("distributing training data on the Dask cluster")
|
|
|
|
# split training data into two partitions
|
|
rows_in_part1 = int(np.sum(group[:100]))
|
|
rows_in_part2 = X.shape[0] - rows_in_part1
|
|
num_features = X.shape[1]
|
|
|
|
# make this array dense because we're splitting across
|
|
# a sparse boundary to partition the data
|
|
X = X.toarray()
|
|
|
|
dX = da.from_array(x=X, chunks=[(rows_in_part1, rows_in_part2), (num_features,)])
|
|
dy = da.from_array(
|
|
x=y,
|
|
chunks=[
|
|
(rows_in_part1, rows_in_part2),
|
|
],
|
|
)
|
|
dg = da.from_array(x=group, chunks=[(100, group.size - 100)])
|
|
|
|
print("beginning training")
|
|
|
|
dask_model = lgb.DaskLGBMRanker(n_estimators=10)
|
|
dask_model.fit(dX, dy, group=dg)
|
|
assert dask_model.fitted_
|
|
|
|
print("done training")
|