NeuralNetClassifier doesn't fit to a skorch.dataset.Dataset with the default train_split #1122

Problem

The Skorch docs claim that it is possible to pass a Skorch or PyTorch Dataset to the .fit() method of a NeuralNet, with the caveat that y=None must be passed explicitly with the subclasses NeuralNetClassifier and NeuralNetRegressor.

However, this currently fails with a non-obvious error unless you override the default train_split argument.

Example

Suppose you have this setup:

import numpy as np
from sklearn.datasets import load_breast_cancer
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import minmax_scale
from skorch.dataset import Dataset

cancer = load_breast_cancer()
X_train, X_test, y_train, y_test = train_test_split(
    cancer.data.astype(np.float32),
    cancer.target,
    stratify=cancer.target
)
train_ds = Dataset(X_train, y_train)

Now suppose you have the following simple model:

import skorch
from skorch.toy import MLPModule
import torch.nn as nn

model = skorch.NeuralNetClassifier(
    MLPModule,
    module__input_units=cancer.data.shape[1],
    module__output_units=len(np.unique(cancer.target)),
    module__hidden_units=1000,
    criterion=nn.CrossEntropyLoss,
)

This works:

model.fit(X_train, y_train)

But this fails:

model.fit(train_ds, y=None)

Error:

ValueError: Stratified CV requires explicitly passing a suitable y.

It is not clear from the error message how a user should proceed. See this StackOverflow post for example.

Why?

The problem arises because the default value for train_split is ValidSplit(5, stratified=True), but the code in ValidSplit.__call__ doesn't handle stratification for a Dataset object.

Possible solution

It would improve usability if Skorch's NeuralNetClassifier were just to work with a Dataset passed to .fit(). Here are some ways to make this happen:

1. If passed a skorch.dataset.Dataset, make use of its .y attribute for determining class proportions for stratification. This would be a simple partial fix that addresses the specific case above.
2. If passed a torch.data.TensorDataset called as TensorDataset(X, y), we could likewise access y from its attribute .tensors[1].
3. If passed some other PyTorch Dataset, Skorch could disable stratification and emit an appropriate warning (not an error) to inform the user of an appropriate workaround.
4. Change the default train_split value for NeuralNetClassifier to ValidSplit(5, stratified=False), as for a NeuralNet.

User workaround

So what is an appropriate workaround to point the user to? The error message above implies the user must dig out all y values from the Dataset, which could be very expensive. But the most appropriate workaround is probably to pass a different argument to train_split= than the default -- either (a) using predefined_split(), (b) disabling stratification by passing train_split=ValidSplit(5, stratified=False), or perhaps even (c) to disable validation entirely by passing train_split=None.

Class proportions

All stratification requires is knowing the proportions of each class label in the dataset, which is much less demanding than knowing all the actual y values. Ideally Skorch would allow for passing class proportions for its stratification -- perhaps as a dictionary like {0: 0.4, 1: 0.4, 2: 0.2} or a dense array from calling np.bincount(y). So a further proposal is:

5. Enhance ValidSplit() to accept class proportions for its classification. This could be via the existing stratified= kwarg or perhaps via a new kwarg like stratify_proportions=.

My recommendation would be to adopt (1), (2), (3), and (5). I'm happy to accept feedback and create draft a PR for this.