Add Cityscapes DataModule + clean up DM docs

pl_bolts/datamodules/cityscapes_datamodule.py

@@ -0,0 +1,135 @@
+from typing import Optional, Sequence
+
+from pytorch_lightning import LightningDataModule
+from torch.utils.data import DataLoader, random_split
+from torchvision import transforms as transform_lib
+from torchvision.datasets import Cityscapes
+
+
+class CityscapesDataModule(LightningDataModule):
+
+    name = 'Cityscapes'
+    extra_args = {}
+
+    def __init__(
+            self,
+            data_dir,
+            val_split=5000,
+            num_workers=16,
+            batch_size=32,
+            *args,
+            **kwargs,
+    ):
+        """
+        Standard Cityscapes, train, val, test splits and transforms
+
+        Transforms::
+
+            transforms = transform_lib.Compose([
+                transform_lib.ToTensor(),
+            ])
+
+        Example::
+
+            from pl_bolts.datamodules import CityscapesDataModule
+
+            dm = CityscapesDataModule(PATH)
+            model = LitModel(datamodule=dm)
+
+        Or you can set your own transforms
+
+        Example::
+
+            dm.train_transforms = ...
+            dm.test_transforms = ...
+            dm.val_transforms  = ...
+
+        Args:
+            data_dir: where to save/load the data
+            val_split: how many of the training images to use for the validation split
+            num_workers: how many workers to use for loading data
+            batch_size: number of examples per training/eval step
+        """
+        super().__init__(*args, **kwargs)
+        self.dims = (3, 32, 32)
+        self.DATASET = Cityscapes
+        self.data_dir = data_dir
+        self.val_split = val_split
+        self.num_workers = num_workers
+        self.batch_size = batch_size
+
+    @property
+    def num_classes(self):
+        """
+        Return:
+            30
+        """
+        return 30
+
+    def prepare_data(self):
+        """
+        Saves Cityscapes files to data_dir
+        """
+        self.DATASET(self.data_dir, train=True, download=True, transform=transform_lib.ToTensor(), **self.extra_args)
+        self.DATASET(self.data_dir, train=False, download=True, transform=transform_lib.ToTensor(), **self.extra_args)
+
+    def train_dataloader(self):
+        """
+        Cityscapes train set with removed subset to use for validation
+        """
+        transforms = self.default_transforms() if self.train_transforms is None else self.train_transforms
+
+        dataset = self.DATASET(self.data_dir, train=True, download=False, transform=transforms, **self.extra_args)
+        train_length = len(dataset)
+        dataset_train, _ = random_split(dataset, [train_length - self.val_split, self.val_split])
+        loader = DataLoader(
+            dataset_train,
+            batch_size=self.batch_size,
+            shuffle=True,
+            num_workers=self.num_workers,
+            drop_last=True,
+            pin_memory=True
+        )
+        return loader
+
+    def val_dataloader(self):
+        """
+        Cityscapes val set uses a subset of the training set for validation
+        """
+        transforms = self.default_transforms() if self.val_transforms is None else self.val_transforms
+
+        dataset = self.DATASET(self.data_dir, train=True, download=False, transform=transforms, **self.extra_args)
+        train_length = len(dataset)
+        _, dataset_val = random_split(dataset, [train_length - self.val_split, self.val_split])
+        loader = DataLoader(
+            dataset_val,
+            batch_size=self.batch_size,
+            shuffle=False,
+            num_workers=self.num_workers,
+            pin_memory=True,
+            drop_last=True
+        )
+        return loader
+
+    def test_dataloader(self):
+        """
+        Cityscapes test set uses the test split
+        """
+        transforms = self.default_transforms() if self.test_transforms is None else self.test_transforms
+
+        dataset = self.DATASET(self.data_dir, train=False, download=False, transform=transforms, **self.extra_args)
+        loader = DataLoader(
+            dataset,
+            batch_size=self.batch_size,
+            shuffle=False,
+            num_workers=self.num_workers,
+            drop_last=True,
+            pin_memory=True
+        )
+        return loader
+
+    def default_transforms(self):
+        cityscapes_transforms = transform_lib.Compose([
+            transform_lib.ToTensor(),
+        ])
+        return cityscapes_transforms