Training customed LGBM with GPU with Flaml fails

[evansdoe]

I created a customed LGBM learner with the device parameter set to gpu, and when I apply the fit method, FLAML returns with the following ValueError.

ValueError: gpulgbm is not a built-in learner for this task type, only ['xgboost', 'xgb_limitdepth', 'rf', 'lgbm', 'lgbm_spark', 'lrl1', 'lrl2', 'catboost', 'extra_tree', 'kneighbor', 'transformer', 'transformer_ms'] are supported.Please use AutoML.add_learner() to add a customized learner.

Here is the code.

from flaml import AutoML
from flaml.automl.model import LGBMEstimator
from sklearn.model_selection import TimeSeriesSplit
from sklearn.multioutput import MultiOutputRegressor
from sklearn.pipeline import Pipeline
from my_transformer import data_transformer
from train_data import X, y

class GPULGBM(LGBMEstimator):
    def __init__(self, **config):
        super().__init__(device="gpu", **config)


automl = AutoML()
automl.add_learner(learner_name='gpulgbm', learner_class=GPULGBM)


automl_pipeline = Pipeline(
    steps=[
        ("preprocessing", data_transformer()),
        ("automl_estimator", MultiOutputRegressor(estimator=automl, n_jobs=-1)),
    ]
)

tscv = TimeSeriesSplit(gap=0, max_train_size=None, n_splits=5, test_size=None)

automl_settings = {
    "time_budget": 5,
    "metric": "mape", 
    "estimator_list": ["gpulgbm",],  
    "task": "regression",  
    "log_file_name": "flaml_model_training.log",  
    "seed": 7654321,  
    "eval_method": "cv",
    "n_splits": 10,
    "split_type": tscv,
}

pipeline_settings = {f"automl_estimator__{key}": value for key, value in automl_settings.items()}

automl.fit(X, y, **pipeline_settings)

I appreciate any answer.

[sonichi]

It fails because you are using a pipeline and the custom learner is not added within the pipeline.
You can use https://microsoft.github.io/FLAML/docs/Use-Cases/Task-Oriented-AutoML#a-shortcut-to-override-the-search-space
to pass the custom hp "device" instead of using a subclass as a new learner.