Merge pull request #127 from Continvvm/continuumDataset

Add kwargs to pytorch datasets

Author: arthurdouillard

continuum/datasets/base.py

@@ -66,21 +66,19 @@ class _SemanticSegmentationDataset(_ContinuumDataset):
     def data_type(self) -> str:
         return "segmentation"
 
-
 class PyTorchDataset(_ContinuumDataset):
     """Continuum version of torchvision datasets.
-
     :param dataset_type: A Torchvision dataset, like MNIST or CIFAR100.
+    :param train: train flag
+    :param download: download
     """
 
     # TODO: some datasets have a different structure, like SVHN for ex. Handle it.
     def __init__(
-        self, data_path: str = "", dataset_type=None, train: bool = True, download: bool = True
-    ):
+            self, data_path: str = "", dataset_type=None, train: bool = True, download: bool = True, **kwargs):
         super().__init__(data_path=data_path, train=train, download=download)
-
         self.dataset_type = dataset_type
-        self.dataset = self.dataset_type(self.data_path, download=self.download, train=self.train)
+        self.dataset = self.dataset_type(self.data_path, download=self.download, train=self.train, **kwargs)
 
     def get_data(self) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
         x, y = np.array(self.dataset.data), np.array(self.dataset.targets)
@@ -97,13 +95,13 @@ class InMemoryDataset(_ContinuumDataset):
     """
 
     def __init__(
-        self,
-        x: np.ndarray,
-        y: np.ndarray,
-        t: Union[None, np.ndarray] = None,
-        data_type: str = "image_array",
-        train: bool = True,
-        download: bool = True,
+            self,
+            x: np.ndarray,
+            y: np.ndarray,
+            t: Union[None, np.ndarray] = None,
+            data_type: str = "image_array",
+            train: bool = True,
+            download: bool = True,
     ):
         super().__init__(train=train, download=download)
 
@@ -141,7 +139,6 @@ def __init__(self, data_path: str, train: bool = True, download: bool = True, da
         self.data_path = data_path
         super().__init__(data_path=data_path, train=train, download=download)
 
-
         allowed_data_types = ("image_path", "segmentation")
         if data_type not in allowed_data_types:
             raise ValueError(f"Invalid data_type={data_type}, allowed={allowed_data_types}.")

tests/test_dataset_attributes.py

@@ -1,6 +1,8 @@
 import pytest
 
 from continuum import datasets as cont_datasets
+from torchvision.datasets import EMNIST, KMNIST
+from continuum.datasets import PyTorchDataset
 
 ATTRS = ["get_data", "_download"]
 
@@ -11,3 +13,13 @@ def test_has_attr(dataset_name):
 
     for attr in ATTRS:
         assert hasattr(d, attr), (dataset_name, attr)
+
+
+@pytest.mark.slow
+def test_PytorchDataset_EMNIST(tmpdir):
+    dataset_train = PyTorchDataset(tmpdir, dataset_type=EMNIST, train=True, download=True, split='letters')
+
+
+@pytest.mark.slow
+def test_PytorchDataset_KMNIST(tmpdir):
+    dataset_train = PyTorchDataset(tmpdir, dataset_type=KMNIST, train=True, download=True)