Support sklearn.compose.TransformedTargetRegressor

tests/test_trainable.py

@@ -5,12 +5,15 @@
 from tune_sklearn._detect_booster import (
     has_xgboost, has_required_lightgbm_version, has_catboost)
 
+from lightgbm import LGBMRegressor
+from ray.tune.schedulers import AsyncHyperBandScheduler
+from sklearn.compose import TransformedTargetRegressor
 from sklearn.datasets import make_classification
 from sklearn.linear_model import LogisticRegression, SGDClassifier
 from sklearn.model_selection import check_cv
 from sklearn.svm import SVC
 
-from tune_sklearn.utils import _check_multimetric_scoring, get_early_stop_type
+from tune_sklearn.utils import _check_multimetric_scoring, get_early_stop_type, EarlyStopping
 
 
 def create_xgboost():
@@ -203,3 +206,14 @@ def testWarmStart(self):
         trainable.train()
         trainable.train()
         trainable.stop()
+
+
+class TestGetEarlyStopType(unittest.TestCase):
+    def testLGBMRegressor(self):
+        lgbm_regressor = LGBMRegressor()
+        transformed_target_regressor = TransformedTargetRegressor(lgbm_regressor)
+        early_stopping_type = get_early_stop_type(
+            estimator=transformed_target_regressor,
+            early_stopping=AsyncHyperBandScheduler(),
+        )
+        self.assertEqual(early_stopping_type, EarlyStopping.LGBM)

tune_sklearn/utils.py

@@ -1,8 +1,11 @@
 from collections import defaultdict
 from typing import Dict
 
+from sklearn.base import RegressorMixin
+from sklearn.compose import TransformedTargetRegressor
 from sklearn.metrics import check_scoring
 from sklearn.pipeline import Pipeline
+
 from tune_sklearn._detect_booster import (
     is_xgboost_model, is_lightgbm_model_of_required_version, is_catboost_model)
 import numpy as np
@@ -94,8 +97,20 @@ def check_error_warm_start(early_stop_type, estimator_config, estimator):
 
 
 def get_early_stop_type(estimator, early_stopping):
+    # If estimator is TransformedTargetRegressor we should get the wrapped regressor.
+    if isinstance(estimator, TransformedTargetRegressor):
+        estimator = estimator.regressor
+
     if not early_stopping:
         return EarlyStopping.NO_EARLY_STOP
+
+    if check_is_pipeline(estimator):
+        for step_name, step in estimator.steps:
+            is_regressor = isinstance(step, RegressorMixin)
+            if is_regressor:
+                estimator = step
+                break
+
     can_partial_fit = check_partial_fit(estimator)
     can_warm_start_iter = check_warm_start_iter(estimator)
     can_warm_start_ensemble = check_warm_start_ensemble(estimator)