Adding types to some of datamodules

pl_bolts/datamodules/async_dataloader.py

@@ -1,10 +1,11 @@
 import re
 from queue import Queue
 from threading import Thread
+from typing import Any, Optional, Union
 
 import torch
 from torch._six import container_abcs, string_classes
-from torch.utils.data import DataLoader
+from torch.utils.data import DataLoader, Dataset
 
 
 class AsynchronousLoader(object):
@@ -26,7 +27,14 @@ class AsynchronousLoader(object):
             constructing one here
     """
 
-    def __init__(self, data, device=torch.device('cuda', 0), q_size=10, num_batches=None, **kwargs):
+    def __init__(
+            self,
+            data: Union[DataLoader, Dataset],
+            device: torch.device = torch.device('cuda', 0),
+            q_size: int = 10,
+            num_batches: Optional[int] = None,
+            **kwargs: Any
+    ):
         if isinstance(data, torch.utils.data.DataLoader):
             self.dataloader = data
         else:
@@ -105,5 +113,5 @@ def __next__(self):
         self.idx += 1
         return out
 
-    def __len__(self):
+    def __len__(self) -> int:
         return self.num_batches

pl_bolts/datamodules/binary_mnist_datamodule.py

@@ -56,7 +56,7 @@ def __init__(
         drop_last: bool = False,
         *args: Any,
         **kwargs: Any,
-    ) -> None:
+    ):
         """
         Args:
             data_dir: Where to save/load the data

pl_bolts/datamodules/cifar10_datamodule.py

@@ -71,7 +71,7 @@ def __init__(
         drop_last: bool = False,
         *args: Any,
         **kwargs: Any,
-    ) -> None:
+    ):
         """
         Args:
             data_dir: Where to save/load the data
@@ -146,14 +146,14 @@ class TinyCIFAR10DataModule(CIFAR10DataModule):
 
     def __init__(
         self,
-        data_dir: str,
+        data_dir: Optional[str] = None,
         val_split: int = 50,
         num_workers: int = 16,
         num_samples: int = 100,
         labels: Optional[Sequence] = (1, 5, 8),
         *args: Any,
         **kwargs: Any,
-    ) -> None:
+    ):
         """
         Args:
             data_dir: where to save/load the data

pl_bolts/datamodules/cityscapes_datamodule.py

@@ -1,3 +1,5 @@
+from typing import Any
+
 from pytorch_lightning import LightningDataModule
 from torch.utils.data import DataLoader
 
@@ -69,8 +71,8 @@ def __init__(
             shuffle: bool = False,
             pin_memory: bool = False,
             drop_last: bool = False,
-            *args,
-            **kwargs,
+            *args: Any,
+            **kwargs: Any,
     ):
         """
         Args:
@@ -109,14 +111,14 @@ def __init__(
         self.target_transforms = None
 
     @property
-    def num_classes(self):
+    def num_classes(self) -> int:
         """
         Return:
             30
         """
         return 30
 
-    def train_dataloader(self):
+    def train_dataloader(self) -> DataLoader:
         """
         Cityscapes train set
         """
@@ -141,7 +143,7 @@ def train_dataloader(self):
         )
         return loader
 
-    def val_dataloader(self):
+    def val_dataloader(self) -> DataLoader:
         """
         Cityscapes val set
         """
@@ -166,7 +168,7 @@ def val_dataloader(self):
         )
         return loader
 
-    def test_dataloader(self):
+    def test_dataloader(self) -> DataLoader:
         """
         Cityscapes test set
         """

pl_bolts/datamodules/fashion_mnist_datamodule.py

@@ -57,7 +57,7 @@ def __init__(
         drop_last: bool = False,
         *args: Any,
         **kwargs: Any,
-    ) -> None:
+    ):
         """
         Args:
             data_dir: Where to save/load the data

pl_bolts/datamodules/imagenet_datamodule.py

@@ -1,5 +1,5 @@
 import os
-from typing import Optional
+from typing import Any, Optional
 
 from pytorch_lightning import LightningDataModule
 from torch.utils.data import DataLoader
@@ -58,8 +58,8 @@ def __init__(
             shuffle: bool = False,
             pin_memory: bool = False,
             drop_last: bool = False,
-            *args,
-            **kwargs,
+            *args: Any,
+            **kwargs: Any,
     ):
         """
         Args:
@@ -94,7 +94,7 @@ def __init__(
         self.num_samples = 1281167 - self.num_imgs_per_val_class * self.num_classes
 
     @property
-    def num_classes(self):
+    def num_classes(self) -> int:
         """
         Return:
 
@@ -103,7 +103,7 @@ def num_classes(self):
         """
         return 1000
 
-    def _verify_splits(self, data_dir, split):
+    def _verify_splits(self, data_dir: str, split: str):
         dirs = os.listdir(data_dir)
 
         if split not in dirs:
@@ -138,7 +138,7 @@ def prepare_data(self):
                 UnlabeledImagenet.generate_meta_bins(path)
                 """)
 
-    def train_dataloader(self):
+    def train_dataloader(self) -> DataLoader:
         """
         Uses the train split of imagenet2012 and puts away a portion of it for the validation split
         """
@@ -160,7 +160,7 @@ def train_dataloader(self):
         )
         return loader
 
-    def val_dataloader(self):
+    def val_dataloader(self) -> DataLoader:
         """
         Uses the part of the train split of imagenet2012  that was not used for training via `num_imgs_per_val_class`
 
@@ -185,7 +185,7 @@ def val_dataloader(self):
         )
         return loader
 
-    def test_dataloader(self):
+    def test_dataloader(self) -> DataLoader:
         """
         Uses the validation split of imagenet2012 for testing
         """

pl_bolts/datamodules/kitti_datamodule.py

@@ -1,4 +1,5 @@
 import os
+from typing import Any, Optional
 
 import torch
 from pytorch_lightning import LightningDataModule
@@ -21,7 +22,7 @@ class KittiDataModule(LightningDataModule):
 
     def __init__(
             self,
-            data_dir: str,
+            data_dir: Optional[str] = None,
             val_split: float = 0.2,
             test_split: float = 0.1,
             num_workers: int = 16,
@@ -30,8 +31,8 @@ def __init__(
             shuffle: bool = False,
             pin_memory: bool = False,
             drop_last: bool = False,
-            *args,
-            **kwargs,
+            *args: Any,
+            **kwargs: Any,
     ):
         """
         Kitti train, validation and test dataloaders.
@@ -100,7 +101,7 @@ def __init__(
                                                                 lengths=[train_len, val_len, test_len],
                                                                 generator=torch.Generator().manual_seed(self.seed))
 
-    def train_dataloader(self):
+    def train_dataloader(self) -> DataLoader:
         loader = DataLoader(
             self.trainset,
             batch_size=self.batch_size,
@@ -111,7 +112,7 @@ def train_dataloader(self):
         )
         return loader
 
-    def val_dataloader(self):
+    def val_dataloader(self) -> DataLoader:
         loader = DataLoader(
             self.valset,
             batch_size=self.batch_size,
@@ -122,7 +123,7 @@ def val_dataloader(self):
         )
         return loader
 
-    def test_dataloader(self):
+    def test_dataloader(self) -> DataLoader:
         loader = DataLoader(
             self.testset,
             batch_size=self.batch_size,

pl_bolts/datamodules/mnist_datamodule.py

@@ -56,7 +56,7 @@ def __init__(
         drop_last: bool = False,
         *args: Any,
         **kwargs: Any,
-    ) -> None:
+    ):
         """
         Args:
             data_dir: Where to save/load the data

pl_bolts/datamodules/sklearn_datamodule.py

@@ -1,5 +1,5 @@
 import math
-from typing import Any
+from typing import Any, Tuple
 
 import numpy as np
 import torch
@@ -42,10 +42,10 @@ def __init__(self, X: np.ndarray, y: np.ndarray, X_transform: Any = None, y_tran
         self.X_transform = X_transform
         self.y_transform = y_transform
 
-    def __len__(self):
+    def __len__(self) -> int:
         return len(self.X)
 
-    def __getitem__(self, idx):
+    def __getitem__(self, idx) -> Tuple[np.ndarray, np.ndarray]:
         x = self.X[idx].astype(np.float32)
         y = self.Y[idx]
 
@@ -89,10 +89,10 @@ def __init__(self, X: torch.Tensor, y: torch.Tensor, X_transform: Any = None, y_
         self.X_transform = X_transform
         self.y_transform = y_transform
 
-    def __len__(self):
+    def __len__(self) -> int:
         return len(self.X)
 
-    def __getitem__(self, idx):
+    def __getitem__(self, idx) -> Tuple[torch.Tensor, torch.Tensor]:
         x = self.X[idx].float()
         y = self.Y[idx]
 
@@ -145,14 +145,14 @@ def __init__(
             x_val=None, y_val=None,
             x_test=None, y_test=None,
             val_split=0.2, test_split=0.1,
-            num_workers=2,
-            random_state=1234,
-            shuffle=True,
+            num_workers: int = 2,
+            random_state: int = 1234,
+            shuffle: bool = True,
             batch_size: int = 16,
-            pin_memory=False,
-            drop_last=False,
-            *args,
-            **kwargs,
+            pin_memory: bool = False,
+            drop_last: bool = False,
+            *args: Any,
+            **kwargs: Any,
     ):
 
         super().__init__(*args, **kwargs)
@@ -193,12 +193,20 @@ def __init__(
 
         self._init_datasets(X, y, x_val, y_val, x_test, y_test)
 
-    def _init_datasets(self, X, y, x_val, y_val, x_test, y_test):
+    def _init_datasets(
+        self,
+        X: np.ndarray,
+        y: np.ndarray,
+        x_val: np.ndarray,
+        y_val: np.ndarray,
+        x_test: np.ndarray,
+        y_test: np.ndarray
+    ):
         self.train_dataset = SklearnDataset(X, y)
         self.val_dataset = SklearnDataset(x_val, y_val)
         self.test_dataset = SklearnDataset(x_test, y_test)
 
-    def train_dataloader(self):
+    def train_dataloader(self) -> DataLoader:
         loader = DataLoader(
             self.train_dataset,
             batch_size=self.batch_size,
@@ -209,7 +217,7 @@ def train_dataloader(self):
         )
         return loader
 
-    def val_dataloader(self):
+    def val_dataloader(self) -> DataLoader:
         loader = DataLoader(
             self.val_dataset,
             batch_size=self.batch_size,
@@ -220,7 +228,7 @@ def val_dataloader(self):
         )
         return loader
 
-    def test_dataloader(self):
+    def test_dataloader(self) -> DataLoader:
         loader = DataLoader(
             self.test_dataset,
             batch_size=self.batch_size,