gbnet.xgbmodule

Classes

XGBModule	XGBoost Module that wraps XGBoost boosting into a PyTorch Module.
XGBObj	Helper class for use with XGBoost as a backend for XGBModule

Module Contents

class gbnet.xgbmodule.XGBModule(batch_size, input_dim, output_dim, params={}, min_hess=0)

Bases: gbnet.base.BaseGBModule

XGBoost Module that wraps XGBoost boosting into a PyTorch Module.

This module allows integration of XGBoost gradient boosting with PyTorch neural networks. It maintains the boosting model state and handles both training and inference.

Parameters:

• batch_size (int) – Size of training data

• input_dim (int) – Dimension of input features

• output_dim (int) – Dimension of output predictions

• params (dict, optional) – Parameters passed to LightGBM. Defaults to {}.

• min_hess (float, optional) – Minimum hessian value submitted to LightGBM. Defaults to 0.

Variables:

• batch_size (int) – Size of mini-batches

• input_dim (int) – Input feature dimension

• output_dim (int) – Output prediction dimension

• params (dict) – LightGBM parameters

• bst (lightgbm.Booster) – The underlying LightGBM booster

• FX (torch.nn.Parameter) – Current predictions tensor

• train_dat (lightgbm.Dataset) – Training dataset used for caching

• min_hess (float) – Minimum hessian threshold

batch_size

input_dim

output_dim

params

n_completed_boost_rounds = 0

min_hess = 0

bst

dtrain = None

training_n = None

FX

_check_training_data()

_input_checking_setting(input_data)

Parameters:

input_data (Union[xgboost.DMatrix, pandas.DataFrame, numpy.ndarray])

forward(input_data, return_tensor=True)

Forward pass through the XGBoost module.

Parameters:

• input_dataset (Union[xgb.DMatrix, np.ndarray, pd.DataFrame]) – Input data for prediction. Can be a XGBoost DMatrix, numpy array, or pandas DataFrame.

• return_tensor (bool, optional) – Whether to return predictions as a PyTorch tensor. Defaults to True.

• input_data (Union[xgboost.DMatrix, numpy.ndarray, pandas.DataFrame])

Returns:

Model predictions. Returns a PyTorch tensor if return_tensor=True, otherwise returns a numpy array.

Return type:

Union[torch.Tensor, np.ndarray]

The forward pass handles both train and eval - In train mode, maintains state between iterations and updates internal FX tensor - In eval mode, generates predictions on new data using the trained model

gb_calc()

gb_step()

Performs a gradient boosting step to update the model.

This method: 1. Computes gradients and hessians from the current predictions 3. Updates the internal boosting model

The gradients are scaled by batch size and hessians are clipped to a minimum value to ensure numerical stability.

Returns:

None

_gb_step_grad_hess(grad, hess)

get_extra_state()

set_extra_state(state)

class gbnet.xgbmodule.XGBObj(grad, hess)

Helper class for use with XGBoost as a backend for XGBModule

grad

hess

__call__(preds, dtrain)