Refactor pl_bolts.datamodules and pl_bolts.datasets

pl_bolts/datamodules/__init__.py

@@ -1,105 +1,48 @@
 from pl_bolts.datamodules.async_dataloader import AsynchronousLoader
-
-__all__ = []
-
-try:
- from pl_bolts.datamodules.binary_mnist_datamodule import BinaryMNISTDataModule
-except ModuleNotFoundError:
- pass
-else:
- __all__ += ['BinaryMNISTDataModule']
-
-try:
- from pl_bolts.datamodules.cifar10_datamodule import (
- CIFAR10DataModule,
- TinyCIFAR10DataModule,
- )
-except ModuleNotFoundError:
- pass
-else:
- __all__ += ['CIFAR10DataModule', 'TinyCIFAR10DataModule']
-
-try:
- from pl_bolts.datamodules.experience_source import (
- ExperienceSourceDataset,
- ExperienceSource,
- DiscountedExperienceSource,
- )
-except ModuleNotFoundError:
- pass
-else:
- __all__ += ['ExperienceSourceDataset', 'ExperienceSource', 'DiscountedExperienceSource']
-
-try:
- from pl_bolts.datamodules.fashion_mnist_datamodule import FashionMNISTDataModule
-except ModuleNotFoundError:
- pass
-else:
- __all__ += ['FashionMNISTDataModule']
-
-try:
- from pl_bolts.datamodules.imagenet_datamodule import ImagenetDataModule
-except ModuleNotFoundError:
- pass
-else:
- __all__ += ['ImagenetDataModule']
-
-try:
- from pl_bolts.datamodules.mnist_datamodule import MNISTDataModule
-except ModuleNotFoundError:
- pass
-else:
- __all__ += ['MNISTDataModule']
-
-try:
- from pl_bolts.datamodules.sklearn_datamodule import (
- SklearnDataset,
- SklearnDataModule,
- TensorDataset,
- )
-except ModuleNotFoundError:
- pass
-else:
- __all__ += ['SklearnDataset', 'SklearnDataModule', 'TensorDataset']
-
-try:
- from pl_bolts.datamodules.ssl_imagenet_datamodule import SSLImagenetDataModule
-except ModuleNotFoundError:
- pass
-else:
- __all__ += ['SSLImagenetDataModule']
-
-try:
- from pl_bolts.datamodules.stl10_datamodule import STL10DataModule
-except ModuleNotFoundError:
- pass
-else:
- __all__ += ['STL10DataModule']
-
-try:
- from pl_bolts.datamodules.vocdetection_datamodule import VOCDetectionDataModule
-except ModuleNotFoundError:
- pass
-else:
- __all__ += ['VOCDetectionDataModule']
-
-try:
- from pl_bolts.datamodules.cityscapes_datamodule import CityscapesDataModule
-except ModuleNotFoundError: # pragma: no-cover
- pass
-else:
- __all__ += ['CityscapesDataModule']
-
-try:
- from pl_bolts.datasets.kitti_dataset import KittiDataset
-except ModuleNotFoundError:
- pass
-else:
- __all__ += ['KittiDataset']
-
-try:
- from pl_bolts.datamodules.kitti_datamodule import KittiDataModule
-except ModuleNotFoundError:
- pass
-else:
- __all__ += ['KittiDataModule']
+from pl_bolts.datamodules.binary_mnist_datamodule import BinaryMNISTDataModule
+from pl_bolts.datamodules.cifar10_datamodule import (
+ CIFAR10DataModule,
+ TinyCIFAR10DataModule,
+)
+from pl_bolts.datamodules.experience_source import (
+ ExperienceSourceDataset,
+ ExperienceSource,
+ DiscountedExperienceSource,
+)
+from pl_bolts.datamodules.fashion_mnist_datamodule import FashionMNISTDataModule
+from pl_bolts.datamodules.imagenet_datamodule import ImagenetDataModule
+from pl_bolts.datamodules.mnist_datamodule import MNISTDataModule
+from pl_bolts.datamodules.sklearn_datamodule import (
+ SklearnDataset,
+ SklearnDataModule,
+ TensorDataset,
+)
+from pl_bolts.datamodules.ssl_imagenet_datamodule import SSLImagenetDataModule
+from pl_bolts.datamodules.stl10_datamodule import STL10DataModule
+from pl_bolts.datamodules.vocdetection_datamodule import VOCDetectionDataModule
+from pl_bolts.datamodules.cityscapes_datamodule import CityscapesDataModule
+from pl_bolts.datasets.kitti_dataset import KittiDataset
+from pl_bolts.datamodules.kitti_datamodule import KittiDataModule
+
+
+__all__ = [
+ 'BinaryMNISTDataModule',
+ 'CIFAR10DataModule',
+ 'TinyCIFAR10DataModule',
+ 'ExperienceSourceDataset',
+ 'ExperienceSource',
+ 'DiscountedExperienceSource',
+ 'FashionMNISTDataModule',
+ 'ImagenetDataModule',
+ 'MNISTDataModule',
+ 'SklearnDataset',
+ 'SklearnDataModule',
+ 'TensorDataset',
+ 'SSLImagenetDataModule',
+ 'STL10DataModule',
+ 'VOCDetectionDataModule',
+ 'CityscapesDataModule',
+ 'KittiDataset',
+ 'KittiDataModule',
+ 'AsynchronousLoader',
+]

pl_bolts/datamodules/experience_source.py

@@ -4,12 +4,16 @@
 """
 from abc import ABC
 from collections import deque, namedtuple
+import importlib
 from typing import Iterable, Callable, Tuple, List
 
 import torch
-from gym import Env
 from torch.utils.data import IterableDataset
 
+_GYM_AVAILABLE = importlib.util.find_spec("gym") is not None
+if _GYM_AVAILABLE:
+ from gym import Env
+
 # Datasets
 
 Experience = namedtuple(
@@ -172,7 +176,7 @@ def env_actions(self, device) -> List[List[int]]:
 
  return actions
 
- def env_step(self, env_idx: int, env: Env, action: List[int]) -> Experience:
+ def env_step(self, env_idx: int, env: "Env", action: List[int]) -> Experience:
  """
  Carries out a step through the given environment using the given action
 
@@ -236,7 +240,7 @@ def pop_rewards_steps(self):
 class DiscountedExperienceSource(ExperienceSource):
  """Outputs experiences with a discounted reward over N steps"""
 
- def __init__(self, env: Env, agent, n_steps: int = 1, gamma: float = 0.99):
+ def __init__(self, env: "Env", agent, n_steps: int = 1, gamma: float = 0.99):
  super().__init__(env, agent, (n_steps + 1))
  self.gamma = gamma
  self.steps = n_steps

pl_bolts/datamodules/imagenet_datamodule.py

@@ -1,3 +1,4 @@
+import importlib
 import os
 from typing import Optional
 from warnings import warn
@@ -7,23 +8,13 @@
 
 from pl_bolts.transforms.dataset_normalizations import imagenet_normalization
 
-try:
+_TORCHVISION_AVAILABLE = importlib.util.find_spec("torchvision") is not None
+if _TORCHVISION_AVAILABLE:
  from torchvision import transforms as transform_lib
-except ModuleNotFoundError:
- warn('You want to use `torchvision` which is not installed yet,' # pragma: no-cover
- ' install it with `pip install torchvision`.')
- _TORCHVISION_AVAILABLE = False
-else:
- _TORCHVISION_AVAILABLE = True
-
-try:
  from pl_bolts.datasets.imagenet_dataset import UnlabeledImagenet
-except ModuleNotFoundError:
+else:
  warn('You want to use `torchvision` which is not installed yet,' # pragma: no-cover
  ' install it with `pip install torchvision`.')
- _TORCHVISION_AVAILABLE = False
-else:
- _TORCHVISION_AVAILABLE = True
 
 
 class ImagenetDataModule(LightningDataModule):

pl_bolts/datamodules/kitti_datamodule.py

@@ -1,13 +1,17 @@
+import importlib
 import os
 
 import torch
-import torchvision.transforms as transforms
 from pytorch_lightning import LightningDataModule
 from torch.utils.data import DataLoader
 from torch.utils.data.dataset import random_split
 
 from pl_bolts.datasets.kitti_dataset import KittiDataset
 
+_TORCHVISION_AVAILABLE = importlib.util.find_spec("torchvision") is not None
+if _TORCHVISION_AVAILABLE:
+ import torchvision.transforms as transforms
+
 
 class KittiDataModule(LightningDataModule):
 
@@ -56,6 +60,11 @@ def __init__(
  batch_size: the batch size
  seed: random seed to be used for train/val/test splits
  """
+ if not _TORCHVISION_AVAILABLE:
+ raise ModuleNotFoundError( # pragma: no-cover
+ 'You want to use `transforms` from `torchvision` which is not installed yet.'
+ )
+
  super().__init__(*args, **kwargs)
  self.data_dir = data_dir if data_dir is not None else os.getcwd()
  self.batch_size = batch_size

pl_bolts/datamodules/sklearn_datamodule.py

@@ -1,3 +1,4 @@
+from importlib.util import find_spec
 import math
 from typing import Any
 
@@ -6,14 +7,10 @@
 from pytorch_lightning import LightningDataModule
 from torch.utils.data import Dataset, DataLoader
 
-try:
+
+_SKLEARN_AVAILABLE = find_spec("sklearn") is not None
+if _SKLEARN_AVAILABLE:
  from sklearn.utils import shuffle as sk_shuffle
-except ModuleNotFoundError:
- raise ModuleNotFoundError('You want to use `sklearn` which is not installed yet,' # pragma: no-cover
- ' install it with `pip install sklearn`.')
- _SKLEARN_AVAILABLE = False
-else:
- _SKLEARN_AVAILABLE = True
 
 
 class SklearnDataset(Dataset):

pl_bolts/datamodules/vocdetection_datamodule.py

@@ -1,19 +1,17 @@
+import importlib
 from warnings import warn
 
 import torch
-import torchvision.transforms as T
 from pytorch_lightning import LightningDataModule
 from torch.utils.data import DataLoader
 
-try:
+_TORCHVISION_AVAILABLE = importlib.util.find_spec("torchvision") is not None
+if _TORCHVISION_AVAILABLE:
+ import torchvision.transforms as T
  from torchvision.datasets import VOCDetection
-
-except ModuleNotFoundError:
+else:
  warn('You want to use `torchvision` which is not installed yet,' # pragma: no-cover
  ' install it with `pip install torchvision`.')
- _TORCHVISION_AVAILABLE = False
-else:
- _TORCHVISION_AVAILABLE = True
 
 
 class Compose(object):
@@ -118,13 +116,13 @@ def __init__(
  *args,
  **kwargs,
  ):
- super().__init__(*args, **kwargs)
-
  if not _TORCHVISION_AVAILABLE:
  raise ModuleNotFoundError( # pragma: no-cover
  'You want to use VOC dataset loaded from `torchvision` which is not installed yet.'
  )
 
+ super().__init__(*args, **kwargs)
+
  self.year = year
  self.data_dir = data_dir
  self.num_workers = num_workers

pl_bolts/datasets/imagenet_dataset.py

@@ -1,5 +1,6 @@
 import gzip
 import hashlib
+import importlib
 import os
 import shutil
 import tarfile
@@ -11,13 +12,12 @@
 import torch
 from torch._six import PY3
 
-try:
+_TORCHVISION_AVAILABLE = importlib.util.find_spec("torchvision") is not None
+if _TORCHVISION_AVAILABLE:
  from torchvision.datasets import ImageNet
  from torchvision.datasets.imagenet import load_meta_file
-except ModuleNotFoundError as err:
- raise ModuleNotFoundError( # pragma: no-cover
- 'You want to use `torchvision` which is not installed yet, install it with `pip install torchvision`.'
- ) from err
+else:
+ ImageNet = object # pragma: no-cover
 
 
 class UnlabeledImagenet(ImageNet):
@@ -48,6 +48,11 @@ def __init__(
  download:
  kwargs:
  """
+ if not _TORCHVISION_AVAILABLE:
+ raise ModuleNotFoundError( # pragma: no-cover
+ 'You want to use `torchvision` which is not installed yet, install it with `pip install torchvision`.'
+ )
+
  root = self.root = os.path.expanduser(root)
 
  # [train], [val] --> [train, val], [test]

pl_bolts/datasets/kitti_dataset.py

@@ -1,9 +1,14 @@
+import importlib
 import os
 
 import numpy as np
-from PIL import Image
 from torch.utils.data import Dataset
 
+_PIL_AVAILABLE = importlib.util.find_spec("PIL") is not None
+if _PIL_AVAILABLE:
+ from PIL import Image
+
+
 DEFAULT_VOID_LABELS = (0, 1, 2, 3, 4, 5, 6, 9, 10, 14, 15, 16, 18, 29, 30, -1)
 DEFAULT_VALID_LABELS = (7, 8, 11, 12, 13, 17, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 31, 32, 33)
 
@@ -41,6 +46,11 @@ def __init__(
  void_labels: useless classes to be excluded from training
  valid_labels: useful classes to include
  """
+ if not _PIL_AVAILABLE:
+ raise ModuleNotFoundError( # pragma: no-cover
+ 'You want to use `Pillow` which is not installed yet.'
+ )
+
  self.img_size = img_size
  self.void_labels = void_labels
  self.valid_labels = valid_labels