Estimator base class wrapping model and preprocessing pipeline

Viewed 228

I'm trying to work out how to easily wrap existing estimators along with a preprocessing pipeline and a target encoder, essentially generalizing the idea behind scikit's TransformedTargetRegressor. I have a possible solution but I'm wondering if I'm missing any repercussions of the design that are not immediately obvious. The basic idea is this:

class BaseModel:

    Model = None

    def __init__(self, feature_encoder=None, target_encoder=None, **params):
        steps = [("features", feature_encoder)] if feature_encoder else []
        steps.append(("model", self.Model(**params)))
        self.pipe_ = Pipeline(steps=steps)
        self.target_encoder_ = target_encoder

    def get_params(self, deep=True):
        """Argument `deep` essentially differentiates the use of the resulting params dict.

            - `deep=True` is used in GridSearch etc. to know which parameters can be set with `set_params`.
            - `deep=False` is used by clone(), where the resulting keys must correspond to __init__ args.
        """
        if deep:
            return self.pipe_.get_params(deep=deep)

        params = {
            "feature_encoder": self.pipe_.named_steps["features"],
            "target_encoder": self.target_encoder_,
        }
        params.update(self.pipe_.named_steps["model"].get_params())
        return params

    def set_params(self, **kwargs):
        self.pipe_.set_params(**kwargs)
        return self

    def prepare_fit(self, X, y=None):
        """Encode target, determine model parameters dynamically depending on data etc."""
        ...

    def fit(self, X, y=None):
        y = self.prepare_fit(X, y)
        self.pipe_.fit(X, y)
        ...
        return self

    def predict(self, X):
        """Predict and decode (inverse transform) target"""
        yp = self.pipe_.predict(X)
        ...

And so wrapping CatBoost e.g. would simply be

class CatboostClassifier(BaseModel):
    Model = catboost.CatBoostClassifier

The crucial part is getting get_params and set_params right such that the wrapped models play nice with scikit's grid search etc., even though the design doesn't follow the official guidelines of having model attributes match its __init__ args.

It seems to work, in the sense that getting and setting any of the model or preprocessing pipeline's parameters works, as does cloning:

model = CatboostClassifier(verbose=50, iterations=100, feature_encoder=..., target_encoder=...)

params = {
    "model__iterations": 95,
    "features__datetime__encode__components": ["year", "month"],
}

print("Params before:", {k: v for k, v in model.get_params().items() if k in params})
cloned = clone(clone(model).set_params(**params))
print("Params after:", {k: v for k, v in cloned.get_params().items() if k in params})
>> Params before: {'features__datetime__encode__components': ['year'], 'model__iterations': 100}
>> Params after: {'features__datetime__encode__components': ['year', 'month'], 'model__iterations': 95}

And GridSearchCV also seems to work:

params = {
    "model__iterations": [100, 200],
    "model__learning_rate": [0.2, 0.3],
}

search = GridSearchCV(model, param_grid=params, cv=5)
search.fit(X, y)
print(search.best_params_)
print(search.best_estimator_.get_params()["model__iterations"])

>> {'model__iterations': 100, 'model__learning_rate': 0.2}
>> 100

So... the question I have is if there is any other way this design may go wrong, e.g. any pending deprecations or planned changes that may break this use of get_params and set_params, e.g.? From what I can see, the only trick necessary to have an arbitrary constructor like this (breaking the official developer guidelines), is that get_params(deep=True) should return all settable parameters, while get_params(deep=False) is used by base.clone() only and needs to return any and all parameters necessary to call the constructor and make a (correct) copy.

I know this is a rather long "question", but I'd be grateful to know about any caveats I should be aware of regarding the proposed pattern.

0 Answers
Related