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,