Implement sklearn API Metadata routing

[DCoupry]

The main problem here is the ability to use weird input schemes (mixed lengths, Datasets etc) with complex pipelines, which often await tabular data. This is for example the case with torch_geometric Data objects, and complex, non-GNN architectures receiving multiple types of "X".
see #1084 for a relevant issue.

proposed solution: metadata routing. If we can successfully route the data (e.g: an array or pandas DataFrame of torch_geometric.data.Data), it would be trivial to pass them as additional fit params, and use a collate_fn to handle stuff from there.
Metadata routing is not implemented for now in skorch NeuralNet, but it looks like it's been partailly done for LLMs?

[DCoupry]

relevant documentation sklearn-side: https://scikit-learn.org/stable/auto_examples/miscellaneous/plot_metadata_routing.html#third-party-development-and-scikit-learn-dependency

[BenjaminBossan]

Note that in #1078, we changed skorch's net classes to inherit from BaseEstimator (there is no release yet, so installing from source is required). From my understanding, this should be largely enough to support metadata routing, although I haven't tested it. Did you give it a try?

[DCoupry]

Checked the PR, and it does mention metadata routing wit the set_fit_request flag, but it was bumped in another PR

an example of a very simple thing which is not really doable in a pipeline:

from skorch.regressor import NeuralNetRegressor
import sklearn
from sklearn.datasets import make_regression
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler
import torch.nn as nn
sklearn.set_config(enable_metadata_routing=True)

class RegressionModule(nn.Module):
    def __init__(self):
        super(RegressionModule, self).__init__()
        self.fc = nn.Linear(10,1)

    def forward(self, X, Z):
        return self.fc(X + Z)

X, y = make_regression(n_samples=1000, n_features=20)
X, Z = X[:,:10], X[:,10:]
net = NeuralNetRegressor(RegressionModule)
# Fails
net.fit(X, y, Z=Z)
# fails
net.set_fit_request(Z=True)
# fails, and is ultimately what I need to do: scale X, not Z, pass Z to neural net.
# cannot use skorch helper classes here (slicedicts and datasets)
pipe = make_pipeline(StandardScaler(), net)
pipe.fit(X, y, Z=Z)

[BenjaminBossan]

net.set_fit_request(Z=True)
# fails, and is ultimately what I need to do: scale X, not Z, pass Z to neural net.

I could reproduce, set_fit_request does not exist. I thought that inheriting from BaseEstimator would be enough to get the method. For reference, here is the MRO:

>>> type(net).mro()
[skorch.regressor.NeuralNetRegressor,
 sklearn.base.RegressorMixin,
 skorch.net.NeuralNet,
 sklearn.base.BaseEstimator,
 sklearn.utils._estimator_html_repr._HTMLDocumentationLinkMixin,
 sklearn.utils._metadata_requests._MetadataRequester,
 object]

So _MetadataRequester is there, I'm not sure what we're still missing. Maybe @adrinjalali can help us.

[DCoupry]

I don't think _MetadataRequester is needed:

class _MetadataRequester:
"""Mixin class for adding metadata request functionality.

---> BaseEstimator inherits from this Mixin.

.. versionadded:: 1.3
"""

A skorch specific part would be to make sure the additional fit.predict/score params get correctly passed to the torch module through the collate_fn?

[adrinjalali]

set_fit_request (and other set_{method}_request methods are only generated by default for methods which explicitly have those "metadata" in their signature. For cases where the method signature doesn't explicitly include them, and instead accepts a **kwargs kinda thing, you'd need to implement the set_{method}_request yourself.

An example of that in scikit-learn is here:

https://github.com/scikit-learn/scikit-learn/blob/bac46762be1d2926663f99de306d9794e63e1fb2/sklearn/metrics/_scorer.py#L303

[BenjaminBossan]

Thanks for the pointer. So basically, we would just have to copy more or less the same method in skorch and we're good? Probably this is the relevant part:

        for param, alias in kwargs.items():
            self._metadata_request.score.add_request(param=param, alias=alias)

[adrinjalali]

Yes, you'd change that to be for the fit method. We probably can add a Mixin in scikit-learn so that you could avoid using the private type of API call.

[BenjaminBossan]

Okay, so I tested this:

# net.py: add to NeuralNet

    def set_fit_request(self, **kwargs):
        from sklearn.utils.metadata_routing import MetadataRequest
        self._metadata_request = MetadataRequest(owner=self.__class__.__name__)
        for param, alias in kwargs.items():
            self._metadata_request.fit.add_request(param=param, alias=alias)
        return self

Then run this script:

from skorch.regressor import NeuralNetRegressor
import sklearn
from sklearn.datasets import make_regression
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler
import torch.nn as nn
sklearn.set_config(enable_metadata_routing=True)

class RegressionModule(nn.Module):
    def __init__(self):
        super(RegressionModule, self).__init__()
        self.fc = nn.Linear(10,1)

    def forward(self, X, Z):
        return self.fc(X + Z)

X, y = make_regression(n_samples=1000, n_features=20)
X, Z = X[:,:10], X[:,10:]
net = NeuralNetRegressor(RegressionModule)
net.set_fit_request(Z=True)
net.fit(X, y, Z=Z)

Now we get a bit further, but then get the error:

File ~/work/skorch/skorch/regressor.py:85, in NeuralNetRegressor.fit(self, X, y, **fit_params)
     74 """See ``NeuralNet.fit``.
     75 
     76 In contrast to ``NeuralNet.fit``, ``y`` is non-optional to
   (...)
     80 
     81 """
     82 # pylint: disable=useless-super-delegation
     83 # this is actually a pylint bug:
     84 # https://github.com/PyCQA/pylint/issues/1085
---> 85 return super(NeuralNetRegressor, self).fit(X, y, **fit_params)

File ~/work/skorch/skorch/net.py:1351, in NeuralNet.fit(self, X, y, **fit_params)
   1348 if not self.warm_start or not self.initialized_:
   1349     self.initialize()
-> 1351 self.partial_fit(X, y, **fit_params)
   1352 return self

File ~/work/skorch/skorch/net.py:1310, in NeuralNet.partial_fit(self, X, y, classes, **fit_params)
   1308 self.notify('on_train_begin', X=X, y=y)
   1309 try:
-> 1310     self.fit_loop(X, y, **fit_params)
   1311 except KeyboardInterrupt:
   1312     pass

File ~/work/skorch/skorch/net.py:1208, in NeuralNet.fit_loop(self, X, y, epochs, **fit_params)
   1205 self.check_training_readiness()
   1206 epochs = epochs if epochs is not None else self.max_epochs
-> 1208 dataset_train, dataset_valid = self.get_split_datasets(
   1209     X, y, **fit_params)
   1210 on_epoch_kwargs = {
   1211     'dataset_train': dataset_train,
   1212     'dataset_valid': dataset_valid,
   1213 }
   1214 iterator_train = self.get_iterator(dataset_train, training=True)

File ~/work/skorch/skorch/net.py:1803, in NeuralNet.get_split_datasets(self, X, y, **fit_params)
   1800 if y is None:
   1801     return self.train_split(dataset, **fit_params)
-> 1803 return self.train_split(dataset, y, **fit_params)

Probably not unexpected. We may have to avoid passing each fit_params to the train_split or add a **kwargs there, but I'd have to think about it.

@DCoupry Would this help you progress?

We probably can add a Mixin in scikit-learn so that you could avoid using the private type of API call.

IMO not high priority, unless other packages have the same issue.