fit_on_batch for SklearnModels where BaseEstimator supports partial_fit

[Jack-42]

🚀 Feature

Some scikit-learn models (BaseEstimators) support the method partial_fit (see here). For cases where these models are being wrapped by SklearnModel, it may make sense to allow the user to call fit_on_batch for the SklearnModel.

Motivation

Assuming the given BaseEstimator supports partial_fit, allowing calls to fit_on_batch would allow for SklearnModel to train on data that does not fit completely into memory (e.g., data from a DiskDataset).

Additional context

If the BaseEstimator does not support partial_fit, we could raise an AttributeError or other appropriate error if fit_on_batch is called.

Another option is to inform users in the docs that they can subclass SklearnModel to implement fit_on_batch, e.g. :

# extension of SklearnModel to implement fit_on_batch
class SklearnModelFoB(SklearnModel):
    def __init__(self, model: BaseEstimator, model_dir: Optional[str] = None, **kwargs):
        super().__init__(model, model_dir, **kwargs)
        # check that the model implements partial_fit
        partial_fit_op = getattr(model, "partial_fit", None)
        if callable(partial_fit_op):
            self.implements_partial_fit = True
        else:
            self.implements_partial_fit = False

    def fit_on_batch(
        self,
        X: np.ndarray,
        y: np.ndarray,
        w: np.ndarray,
        classes: Optional[np.ndarray] = None,
    ):
        if not self.implements_partial_fit:
            raise NotImplementedError(
                "Given sklearn model does not implement partial_fit!"
            )
        if self.use_weights:
            self.model.partial_fit(X, y, classes=classes, sample_weight=w)
        else:
            self.model.partial_fit(X, y, classes=classes)

[rbharath]

I'd be open to this as a new feature. Sounds potentially useful for the community. Would you be able to come by OH (9am PST, MWF ) some day to discuss with us?

[Jack-42]

I'd be open to this as a new feature. Sounds potentially useful for the community. Would you be able to come by OH (9am PST, MWF ) some day to discuss with us?

Sure! I can drop by the next OH (May 15th).

[rbharath]

Great! Please join the discord (https://discord.gg/ArRuv9Eu) if you haven't already. We announce timing adjustments for OH there.