Raise a clear error when ValidSplit gets an IterableDataset (#1151)

Author: aiedwardyi

CHANGES.md

@@ -16,6 +16,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
+- `ValidSplit` now raises a clear error when it receives an `IterableDataset`, instead of an opaque `TypeError` about a missing length (#594)
+
 ## [1.4.0]
 
 ### Added

skorch/dataset.py

@@ -312,7 +312,16 @@ def __call__(self, dataset, y=None, groups=None):
             raise bad_y_error
 
         # pylint: disable=invalid-name
-        len_dataset = get_len(dataset)
+        try:
+            len_dataset = get_len(dataset)
+        except TypeError as exc:
+            if not isinstance(dataset, torch.utils.data.IterableDataset):
+                raise
+            raise ValueError(
+                "Cannot perform a CV split on an IterableDataset because it has "
+                "no length. Set train_split=None to disable the internal "
+                "validation split, or pass a train_split that supports "
+                "IterableDataset.") from exc
         if y is not None:
             len_y = get_len(y)
             if len_dataset != len_y:

skorch/tests/test_net.py

@@ -2215,6 +2215,92 @@ def test_fit_with_dataset_stratified_without_explicit_y_raises(
         msg = "Stratified CV requires explicitly passing a suitable y."
         assert exc.value.args[0] == msg
 
+    @pytest.fixture
+    def iterable_dataset_cls(self):
+        class MyIterableDataset(torch.utils.data.IterableDataset):
+            def __init__(self, X, y):
+                super().__init__()
+                self.X = X
+                self.y = y
+
+            def __iter__(self):
+                return iter(zip(self.X, self.y))
+
+        return MyIterableDataset
+
+    def test_fit_with_iterable_dataset_and_train_split_raises(
+            self, net_cls, module_cls, iterable_dataset_cls, data):
+        from skorch.dataset import ValidSplit
+
+        net = net_cls(
+            module_cls,
+            max_epochs=1,
+            train_split=ValidSplit(stratified=False),
+        )
+        ds = iterable_dataset_cls(*data)
+        with pytest.raises(ValueError) as exc:
+            net.fit(ds, None)
+
+        msg = ("Cannot perform a CV split on an IterableDataset because it has "
+               "no length. Set train_split=None to disable the internal "
+               "validation split, or pass a train_split that supports "
+               "IterableDataset.")
+        assert exc.value.args[0] == msg
+        assert isinstance(exc.value.__cause__, TypeError)
+
+    def test_fit_with_iterable_dataset_no_train_split(
+            self, net_cls, module_cls, iterable_dataset_cls, data):
+        net = net_cls(module_cls, max_epochs=1, train_split=None)
+        ds = iterable_dataset_cls(*data)
+        net.fit(ds, None)  # does not raise
+
+        assert 'train_loss' in net.history[-1]
+
+    @pytest.fixture
+    def sized_iterable_dataset_cls(self):
+        class MySizedIterableDataset(torch.utils.data.IterableDataset):
+            def __init__(self, X, y):
+                super().__init__()
+                self.X = X
+                self.y = y
+
+            def __iter__(self):
+                return iter(zip(self.X, self.y))
+
+            def __len__(self):
+                return len(self.y)
+
+            def __getitem__(self, i):
+                return self.X[i], self.y[i]
+
+        return MySizedIterableDataset
+
+    def test_fit_with_sized_iterable_dataset_and_train_split(
+            self, net_cls, module_cls, sized_iterable_dataset_cls, data):
+        from skorch.dataset import ValidSplit
+
+        net = net_cls(
+            module_cls,
+            max_epochs=1,
+            train_split=ValidSplit(stratified=False),
+        )
+        ds = sized_iterable_dataset_cls(*data)
+        net.fit(ds, None)  # does not raise
+
+        assert 'valid_loss' in net.history[-1]
+
+    def test_fit_with_iterable_dataset_and_custom_train_split(
+            self, net_cls, module_cls, iterable_dataset_cls, data):
+        # a train_split that never indexes the dataset must keep working
+        def train_split(dataset, **kwargs):
+            return dataset, dataset
+
+        net = net_cls(module_cls, max_epochs=1, train_split=train_split)
+        ds = iterable_dataset_cls(*data)
+        net.fit(ds, None)  # does not raise
+
+        assert 'valid_loss' in net.history[-1]
+
     @pytest.fixture
     def dataset_1_item(self):
         class Dataset(torch.utils.data.Dataset):