hummingbird.ml.operator_converters.xgb

Converters for XGBoost models.

hummingbird.ml.operator_converters.xgb._get_tree_parameters(tree_info, extra_config)[source]

Parse the tree and returns an in-memory friendly representation of its structure.

hummingbird.ml.operator_converters.xgb._tree_traversal(tree_info, lefts, rights, features, thresholds, values)[source]

Recursive function for parsing a tree and filling the input data structures.

hummingbird.ml.operator_converters.xgb.convert_sklearn_xgb_classifier(operator, device, extra_config)[source]

Converter for xgboost.XGBClassifier (trained using the Sklearn API).

Args:

operator: An operator wrapping a xgboost.XGBClassifier model device: String defining the type of device the converted operator should be run on extra_config: Extra configuration used to select the best conversion strategy

Returns:

A PyTorch model

hummingbird.ml.operator_converters.xgb.convert_sklearn_xgb_regressor(operator, device, extra_config)[source]

Converter for xgboost.XGBRegressor (trained using the Sklearn API).

Args:

operator: An operator wrapping a xgboost.XGBRegressor model device: String defining the type of device the converted operator should be run on extra_config: Extra configuration used to select the best conversion strategy

Returns:

A PyTorch model