111 lines
3.9 KiB
Python
111 lines
3.9 KiB
Python
from pathlib import Path
|
|
|
|
import h5py
|
|
import numpy as np
|
|
import pandas as pd
|
|
|
|
import lightgbm as lgb
|
|
|
|
|
|
class HDFSequence(lgb.Sequence):
|
|
def __init__(self, hdf_dataset, batch_size):
|
|
"""
|
|
Construct a sequence object from HDF5 with required interface.
|
|
|
|
Parameters
|
|
----------
|
|
hdf_dataset : h5py.Dataset
|
|
Dataset in HDF5 file.
|
|
batch_size : int
|
|
Size of a batch. When reading data to construct lightgbm Dataset, each read reads batch_size rows.
|
|
"""
|
|
# We can also open HDF5 file once and get access to
|
|
self.data = hdf_dataset
|
|
self.batch_size = batch_size
|
|
|
|
def __getitem__(self, idx):
|
|
return self.data[idx]
|
|
|
|
def __len__(self):
|
|
return len(self.data)
|
|
|
|
|
|
def create_dataset_from_multiple_hdf(input_flist, batch_size):
|
|
data = []
|
|
ylist = []
|
|
for f in input_flist:
|
|
f = h5py.File(f, "r")
|
|
data.append(HDFSequence(f["X"], batch_size))
|
|
ylist.append(f["Y"][:])
|
|
|
|
params = {
|
|
"bin_construct_sample_cnt": 200000,
|
|
"max_bin": 255,
|
|
}
|
|
y = np.concatenate(ylist)
|
|
dataset = lgb.Dataset(data, label=y, params=params)
|
|
# With binary dataset created, we can use either Python API or cmdline version to train.
|
|
#
|
|
# Note: in order to create exactly the same dataset with the one created in simple_example.py, we need
|
|
# to modify simple_example.py to pass numpy array instead of pandas DataFrame to Dataset constructor.
|
|
# The reason is that DataFrame column names will be used in Dataset. For a DataFrame with Int64Index
|
|
# as columns, Dataset will use column names like ["0", "1", "2", ...]. While for numpy array, column names
|
|
# are using the default one assigned in C++ code (dataset_loader.cpp), like ["Column_0", "Column_1", ...].
|
|
dataset.save_binary("regression.train.from_hdf.bin")
|
|
|
|
|
|
def save2hdf(input_data, fname, batch_size):
|
|
"""Store numpy array to HDF5 file.
|
|
|
|
Please note chunk size settings in the implementation for I/O performance optimization.
|
|
"""
|
|
with h5py.File(fname, "w") as f:
|
|
for name, data in input_data.items():
|
|
nrow, ncol = data.shape
|
|
if ncol == 1:
|
|
# Y has a single column and we read it in single shot. So store it as an 1-d array.
|
|
chunk = (nrow,)
|
|
data = data.values.flatten()
|
|
else:
|
|
# We use random access for data sampling when creating LightGBM Dataset from Sequence.
|
|
# When accessing any element in a HDF5 chunk, it's read entirely.
|
|
# To save I/O for sampling, we should keep number of total chunks much larger than sample count.
|
|
# Here we are just creating a chunk size that matches with batch_size.
|
|
#
|
|
# Also note that the data is stored in row major order to avoid extra copy when passing to
|
|
# lightgbm Dataset.
|
|
chunk = (batch_size, ncol)
|
|
f.create_dataset(name, data=data, chunks=chunk, compression="lzf")
|
|
|
|
|
|
def generate_hdf(input_fname, output_basename, batch_size):
|
|
# Save to 2 HDF5 files for demonstration.
|
|
df = pd.read_csv(input_fname, header=None, sep="\t")
|
|
|
|
mid = len(df) // 2
|
|
df1 = df.iloc[:mid]
|
|
df2 = df.iloc[mid:]
|
|
|
|
# We can store multiple datasets inside a single HDF5 file.
|
|
# Separating X and Y for choosing best chunk size for data loading.
|
|
fname1 = f"{output_basename}1.h5"
|
|
fname2 = f"{output_basename}2.h5"
|
|
save2hdf({"Y": df1.iloc[:, :1], "X": df1.iloc[:, 1:]}, fname1, batch_size)
|
|
save2hdf({"Y": df2.iloc[:, :1], "X": df2.iloc[:, 1:]}, fname2, batch_size)
|
|
|
|
return [fname1, fname2]
|
|
|
|
|
|
def main():
|
|
batch_size = 64
|
|
output_basename = "regression"
|
|
hdf_files = generate_hdf(
|
|
str(Path(__file__).absolute().parents[1] / "regression" / "regression.train"), output_basename, batch_size
|
|
)
|
|
|
|
create_dataset_from_multiple_hdf(hdf_files, batch_size=batch_size)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|