Refactor pl_bolts.datamodules and pl_bolts.datasets

pl_bolts/datamodules/__init__.py

@@ -1,98 +1,36 @@
 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 (
- DiscountedExperienceSource,
- ExperienceSource,
- ExperienceSourceDataset,
- )
-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 SklearnDataModule, SklearnDataset, 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.cityscapes_datamodule import CityscapesDataModule
+from pl_bolts.datamodules.experience_source import DiscountedExperienceSource, ExperienceSource, ExperienceSourceDataset
+from pl_bolts.datamodules.fashion_mnist_datamodule import FashionMNISTDataModule
+from pl_bolts.datamodules.imagenet_datamodule import ImagenetDataModule
+from pl_bolts.datamodules.kitti_datamodule import KittiDataModule
+from pl_bolts.datamodules.mnist_datamodule import MNISTDataModule
+from pl_bolts.datamodules.sklearn_datamodule import SklearnDataModule, SklearnDataset, 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.datasets.kitti_dataset import KittiDataset
+
+__all__ = [
+ 'AsynchronousLoader',
+ 'BinaryMNISTDataModule',
+ 'CIFAR10DataModule',
+ 'TinyCIFAR10DataModule',
+ 'CityscapesDataModule',
+ 'DiscountedExperienceSource',
+ 'ExperienceSource',
+ 'ExperienceSourceDataset',
+ 'FashionMNISTDataModule',
+ 'ImagenetDataModule',
+ 'KittiDataModule',
+ 'MNISTDataModule',
+ 'SklearnDataModule',
+ 'SklearnDataset',
+ 'TensorDataset',
+ 'SSLImagenetDataModule',
+ 'STL10DataModule',
+ 'VOCDetectionDataModule',
+ 'KittiDataset',
+]

pl_bolts/datamodules/binary_mnist_datamodule.py

@@ -2,18 +2,15 @@
 from pytorch_lightning import LightningDataModule
 from torch.utils.data import DataLoader, random_split
 
+from pl_bolts.datasets.mnist_dataset import BinaryMNIST
+from pl_bolts.utils import _TORCHVISION_AVAILABLE
 from pl_bolts.utils.warnings import warn_missing_pkg
 
-try:
+if _TORCHVISION_AVAILABLE:
  from torchvision import transforms as transform_lib
  from torchvision.datasets import MNIST
-
- from pl_bolts.datasets.mnist_dataset import BinaryMNIST
-except ModuleNotFoundError:
- warn_missing_pkg('torchvision') # pragma: no-cover
- _TORCHVISION_AVAILABLE = False
-else:
- _TORCHVISION_AVAILABLE = True
+else: # pragma: no cover
+ warn_missing_pkg('torchvision')
 
 
 class BinaryMNISTDataModule(LightningDataModule):
@@ -67,8 +64,8 @@ def __init__(
  """
  super().__init__(*args, **kwargs)
 
- if not _TORCHVISION_AVAILABLE:
- raise ModuleNotFoundError( # pragma: no-cover
+ if not _TORCHVISION_AVAILABLE: # pragma: no cover
+ raise ModuleNotFoundError(
  'You want to use MNIST dataset loaded from `torchvision` which is not installed yet.'
  )
 

pl_bolts/datamodules/cifar10_datamodule.py

@@ -7,17 +7,14 @@
 
 from pl_bolts.datasets.cifar10_dataset import TrialCIFAR10
 from pl_bolts.transforms.dataset_normalizations import cifar10_normalization
+from pl_bolts.utils import _TORCHVISION_AVAILABLE
 from pl_bolts.utils.warnings import warn_missing_pkg
 
-try:
+if _TORCHVISION_AVAILABLE:
  from torchvision import transforms as transform_lib
  from torchvision.datasets import CIFAR10
-
-except ModuleNotFoundError:
- warn_missing_pkg('torchvision') # pragma: no-cover
- _TORCHVISION_AVAILABLE = False
-else:
- _TORCHVISION_AVAILABLE = True
+else: # pragma: no cover
+ warn_missing_pkg('torchvision')
 
 
 class CIFAR10DataModule(LightningDataModule):
@@ -83,8 +80,8 @@ def __init__(
  """
  super().__init__(*args, **kwargs)
 
- if not _TORCHVISION_AVAILABLE:
- raise ModuleNotFoundError( # pragma: no-cover
+ if not _TORCHVISION_AVAILABLE: # pragma: no cover
+ raise ModuleNotFoundError(
  'You want to use CIFAR10 dataset loaded from `torchvision` which is not installed yet.'
  )
 

pl_bolts/datamodules/cityscapes_datamodule.py

@@ -1,16 +1,14 @@
 from pytorch_lightning import LightningDataModule
 from torch.utils.data import DataLoader
 
+from pl_bolts.utils import _TORCHVISION_AVAILABLE
 from pl_bolts.utils.warnings import warn_missing_pkg
 
-try:
+if _TORCHVISION_AVAILABLE:
  from torchvision import transforms as transform_lib
  from torchvision.datasets import Cityscapes
-except ModuleNotFoundError:
- warn_missing_pkg('torchvision') # pragma: no-cover
- _TORCHVISION_AVAILABLE = False
-else:
- _TORCHVISION_AVAILABLE = True
+else: # pragma: no cover
+ warn_missing_pkg('torchvision')
 
 
 class CityscapesDataModule(LightningDataModule):
@@ -82,8 +80,8 @@ def __init__(
  """
  super().__init__(*args, **kwargs)
 
- if not _TORCHVISION_AVAILABLE:
- raise ModuleNotFoundError( # pragma: no-cover
+ if not _TORCHVISION_AVAILABLE: # pragma: no cover
+ raise ModuleNotFoundError(
  'You want to use CityScapes dataset loaded from `torchvision` which is not installed yet.'
  )
 

pl_bolts/datamodules/experience_source.py

@@ -10,17 +10,16 @@
 import torch
 from torch.utils.data import IterableDataset
 
+from pl_bolts.utils import _GYM_AVAILABLE
 from pl_bolts.utils.warnings import warn_missing_pkg
 
-_GYM_AVAILABLE = importlib.util.find_spec("gym") is not None
 if _GYM_AVAILABLE:
  from gym import Env
-else:
- warn_missing_pkg("gym") # pragma: no-cover
+else: # pragma: no cover
+ warn_missing_pkg("gym")
+ Env = object
 
 
-# Datasets
-
 Experience = namedtuple(
  "Experience", field_names=["state", "action", "reward", "done", "new_state"]
 )
@@ -181,7 +180,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
 
@@ -245,7 +244,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/fashion_mnist_datamodule.py

@@ -2,16 +2,14 @@
 from pytorch_lightning import LightningDataModule
 from torch.utils.data import DataLoader, random_split
 
+from pl_bolts.utils import _TORCHVISION_AVAILABLE
 from pl_bolts.utils.warnings import warn_missing_pkg
 
-try:
+if _TORCHVISION_AVAILABLE:
  from torchvision import transforms as transform_lib
  from torchvision.datasets import FashionMNIST
-except ModuleNotFoundError:
- warn_missing_pkg('torchvision') # pragma: no-cover
- _TORCHVISION_AVAILABLE = False
-else:
- _TORCHVISION_AVAILABLE = True
+else: # pragma: no cover
+ warn_missing_pkg('torchvision')
 
 
 class FashionMNISTDataModule(LightningDataModule):
@@ -64,8 +62,8 @@ def __init__(
  """
  super().__init__(*args, **kwargs)
 
- if not _TORCHVISION_AVAILABLE:
- raise ModuleNotFoundError( # pragma: no-cover
+ if not _TORCHVISION_AVAILABLE: # pragma: no cover
+ raise ModuleNotFoundError(
  'You want to use fashion MNIST dataset loaded from `torchvision` which is not installed yet.'
  )
 

pl_bolts/datamodules/imagenet_datamodule.py

@@ -4,24 +4,15 @@
 from pytorch_lightning import LightningDataModule
 from torch.utils.data import DataLoader
 
+from pl_bolts.datasets.imagenet_dataset import UnlabeledImagenet
 from pl_bolts.transforms.dataset_normalizations import imagenet_normalization
+from pl_bolts.utils import _TORCHVISION_AVAILABLE
 from pl_bolts.utils.warnings import warn_missing_pkg
 
-try:
+if _TORCHVISION_AVAILABLE:
  from torchvision import transforms as transform_lib
-except ModuleNotFoundError:
- warn_missing_pkg('torchvision') # pragma: no-cover
- _TORCHVISION_AVAILABLE = False
-else:
- _TORCHVISION_AVAILABLE = True
-
-try:
- from pl_bolts.datasets.imagenet_dataset import UnlabeledImagenet
-except ModuleNotFoundError:
- warn_missing_pkg('torchvision') # pragma: no-cover
- _TORCHVISION_AVAILABLE = False
-else:
- _TORCHVISION_AVAILABLE = True
+else: # pragma: no cover
+ warn_missing_pkg('torchvision')
 
 
 class ImagenetDataModule(LightningDataModule):
@@ -78,8 +69,8 @@ def __init__(
  """
  super().__init__(*args, **kwargs)
 
- if not _TORCHVISION_AVAILABLE:
- raise ModuleNotFoundError( # pragma: no-cover
+ if not _TORCHVISION_AVAILABLE: # pragma: no cover
+ raise ModuleNotFoundError(
  'You want to use ImageNet dataset loaded from `torchvision` which is not installed yet.'
  )
 

pl_bolts/datamodules/kitti_datamodule.py

@@ -1,4 +1,3 @@
-import importlib
 import os
 
 import torch
@@ -7,13 +6,13 @@
 from torch.utils.data.dataset import random_split
 
 from pl_bolts.datasets.kitti_dataset import KittiDataset
+from pl_bolts.utils import _TORCHVISION_AVAILABLE
 from pl_bolts.utils.warnings import warn_missing_pkg
 
-_TORCHVISION_AVAILABLE = importlib.util.find_spec("torchvision") is not None
 if _TORCHVISION_AVAILABLE:
  import torchvision.transforms as transforms
-else:
- warn_missing_pkg('torchvision') # pragma: no-cover
+else: # pragma: no cover
+ warn_missing_pkg('torchvision')
 
 
 class KittiDataModule(LightningDataModule):
@@ -63,8 +62,8 @@ 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
+ if not _TORCHVISION_AVAILABLE: # pragma: no cover
+ raise ModuleNotFoundError(
  'You want to use `torchvision` which is not installed yet.'
  )