Add SRGAN and datamodules for super resolution

CHANGELOG.md

@@ -10,6 +10,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Added YOLO model ([#552](https://github.com/PyTorchLightning/pytorch-lightning-bolts/pull/552))
 
+
+- Added `SRGAN`, `SRImageLoggerCallback`, `TVTDataModule`, `SRCelebA`, `SRMNIST`, `SRSTL10` ([#466](https://github.com/PyTorchLightning/pytorch-lightning-bolts/pull/466))
+
+
 ### Changed
 
 
@@ -115,8 +119,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
     [#323](https://github.com/PyTorchLightning/lightning-bolts/pull/323))
 - Added data monitor callbacks `ModuleDataMonitor` and `TrainingDataMonitor` ([#285](https://github.com/PyTorchLightning/lightning-bolts/pull/285))
 - Added DCGAN module ([#403](https://github.com/PyTorchLightning/lightning-bolts/pull/403))
-- Added `VisionDataModule` as parent class for `BinaryMNISTDataModule`, `CIFAR10DataModule`, `FashionMNISTDataModule`,
-  and `MNISTDataModule` ([#400](https://github.com/PyTorchLightning/lightning-bolts/pull/400))
+- Added `VisionDataModule` as parent class for `BinaryMNISTDataModule`, `CIFAR10DataModule`, `FashionMNISTDataModule`, and `MNISTDataModule` ([#400](https://github.com/PyTorchLightning/lightning-bolts/pull/400))
 - Added GIoU loss ([#347](https://github.com/PyTorchLightning/lightning-bolts/pull/347))
 - Added IoU loss ([#469](https://github.com/PyTorchLightning/lightning-bolts/pull/469))
 - Added semantic segmentation model `SemSegment` with `UNet` backend ([#259](https://github.com/PyTorchLightning/lightning-bolts/pull/259))

docs/source/deprecated/models/gans.rst

@@ -86,3 +86,86 @@ LSUN Loss curves:
 
 .. autoclass:: pl_bolts.models.gans.DCGAN
    :noindex:
+
+
+SRGAN
+---------
+SRGAN implementation from the paper `Photo-Realistic Single Image Super-Resolution Using a Generative Adversarial
+Network <https://arxiv.org/pdf/1609.04802.pdf>`_. The implementation is based on the version from
+`deeplearning.ai <https://github.com/https-deeplearning-ai/GANs-Public/blob/master/C3W2_SRGAN_(Optional).ipynb>`_.
+
+Implemented by:
+
+    - `Christoph Clement <https://github.com/chris-clem>`_
+
+MNIST results:
+
+    SRGAN MNIST with scale factor of 2 (left: low res, middle: generated high res, right: ground truth high res):
+
+        .. image:: ../../_images/gans/srgan-mnist-scale_factor=2.png
+            :width: 200
+            :alt: SRGAN MNIST with scale factor of 2
+
+    SRGAN MNIST with scale factor of 4:
+
+        .. image:: ../../_images/gans/srgan-mnist-scale_factor=4.png
+            :width: 200
+            :alt: SRGAN MNIST with scale factor of 4
+
+    SRResNet pretraining command used::
+        >>>  python srresnet_module.py --dataset=mnist --data_dir=~/Data --scale_factor=4 --save_model_checkpoint \
+        --batch_size=16 --num_workers=2 --gpus=4 --accelerator=ddp --precision=16 --max_steps=25000
+
+    SRGAN training command used::
+        >>>  python srgan_module.py --dataset=mnist --data_dir=~/Data --scale_factor=4 --batch_size=16 \
+        --num_workers=2 --scheduler_step=29 --gpus=4 --accelerator=ddp --precision=16 --max_steps=50000
+
+STL10 results:
+
+    SRGAN STL10 with scale factor of 2:
+
+        .. image:: ../../_images/gans/srgan-stl10-scale_factor=2.png
+            :width: 200
+            :alt: SRGAN STL10 with scale factor of 2
+
+    SRGAN STL10 with scale factor of 4:
+
+        .. image:: ../../_images/gans/srgan-stl10-scale_factor=4.png
+            :width: 200
+            :alt: SRGAN STL10 with scale factor of 4
+
+    SRResNet pretraining command used::
+        >>>  python srresnet_module.py --dataset=stl10 --data_dir=~/Data --scale_factor=4 --save_model_checkpoint \
+        --batch_size=16 --num_workers=2 --gpus=4 --accelerator=ddp --precision=16 --max_steps=25000
+
+    SRGAN training command used::
+        >>>  python srgan_module.py --dataset=stl10 --data_dir=~/Data --scale_factor=4 --batch_size=16 \
+        --num_workers=2 --scheduler_step=29 --gpus=4 --accelerator=ddp --precision=16 --max_steps=50000
+
+CelebA results:
+
+    SRGAN CelebA with scale factor of 2:
+
+        .. image:: ../../_images/gans/srgan-celeba-scale_factor=2.png
+            :width: 200
+            :alt: SRGAN CelebA with scale factor of 2
+
+    SRGAN CelebA with scale factor of 4:
+
+        .. image:: ../../_images/gans/srgan-celeba-scale_factor=4.png
+            :width: 200
+            :alt: SRGAN CelebA with scale factor of 4
+
+    SRResNet pretraining command used::
+        >>>  python srresnet_module.py --dataset=celeba --data_dir=~/Data --scale_factor=4 --save_model_checkpoint \
+        --batch_size=16 --num_workers=2 --gpus=4 --accelerator=ddp --precision=16 --max_steps=25000
+
+    SRGAN training command used::
+        >>>  python srgan_module.py --dataset=celeba --data_dir=~/Data --scale_factor=4 --batch_size=16 \
+        --num_workers=2 --scheduler_step=29 --gpus=4 --accelerator=ddp --precision=16 --max_steps=50000
+
+.. autoclass:: pl_bolts.models.gans.SRGAN
+   :noindex:
+
+.. autoclass:: pl_bolts.models.gans.SRResNet
+   :noindex:

pl_bolts/callbacks/__init__.py

@@ -9,6 +9,7 @@
 from pl_bolts.callbacks.verification.batch_gradient import BatchGradientVerificationCallback  # type: ignore
 from pl_bolts.callbacks.vision.confused_logit import ConfusedLogitCallback
 from pl_bolts.callbacks.vision.image_generation import TensorboardGenerativeModelImageSampler
+from pl_bolts.callbacks.vision.sr_image_logger import SRImageLoggerCallback
 
 __all__ = [
     "BatchGradientVerificationCallback",
@@ -20,6 +21,7 @@
     "LatentDimInterpolator",
     "ConfusedLogitCallback",
     "TensorboardGenerativeModelImageSampler",
+    "SRImageLoggerCallback",
     "ORTCallback",
     "SparseMLCallback",
 ]

pl_bolts/callbacks/vision/sr_image_logger.py

@@ -0,0 +1,67 @@
+from typing import Tuple
+
+import pytorch_lightning as pl
+import torch
+import torch.nn.functional as F
+from pytorch_lightning import Callback
+
+from pl_bolts.utils import _TORCHVISION_AVAILABLE
+from pl_bolts.utils.warnings import warn_missing_pkg
+
+if _TORCHVISION_AVAILABLE:
+    from torchvision.utils import make_grid
+else:  # pragma: no cover
+    warn_missing_pkg("torchvision")
+
+
+class SRImageLoggerCallback(Callback):
+    """Logs low-res, generated high-res, and ground truth high-res images to TensorBoard Your model must implement
+    the ``forward`` function for generation.
+
+    Requirements::
+
+        # model forward must work generating high-res from low-res image
+        hr_fake = pl_module(lr_image)
+
+    Example::
+
+        from pl_bolts.callbacks import SRImageLoggerCallback
+
+        trainer = Trainer(callbacks=[SRImageLoggerCallback()])
+    """
+
+    def __init__(self, log_interval: int = 1000, scale_factor: int = 4, num_samples: int = 5) -> None:
+        """
+        Args:
+            log_interval: Number of steps between logging. Default: ``1000``.
+            scale_factor: Scale factor used for downsampling the high-res images. Default: ``4``.
+            num_samples: Number of images of displayed in the grid. Default: ``5``.
+        """
+        super().__init__()
+        self.log_interval = log_interval
+        self.scale_factor = scale_factor
+        self.num_samples = num_samples
+
+    def on_train_batch_end(
+        self,
+        trainer: pl.Trainer,
+        pl_module: pl.LightningModule,
+        outputs: torch.Tensor,
+        batch: Tuple[torch.Tensor, torch.Tensor],
+        batch_idx: int,
+        dataloader_idx: int,
+    ) -> None:
+        global_step = trainer.global_step
+        if global_step % self.log_interval == 0:
+            hr_image, lr_image = batch
+            hr_image, lr_image = hr_image.to(pl_module.device), lr_image.to(pl_module.device)
+            hr_fake = pl_module(lr_image)
+            lr_image = F.interpolate(lr_image, scale_factor=self.scale_factor)
+
+            lr_image_grid = make_grid(lr_image[: self.num_samples], nrow=1, normalize=True)
+            hr_fake_grid = make_grid(hr_fake[: self.num_samples], nrow=1, normalize=True)
+            hr_image_grid = make_grid(hr_image[: self.num_samples], nrow=1, normalize=True)
+
+            grid = torch.cat((lr_image_grid, hr_fake_grid, hr_image_grid), -1)
+            title = "sr_images"
+            trainer.logger.experiment.add_image(title, grid, global_step=global_step)

pl_bolts/datamodules/__init__.py

@@ -10,6 +10,7 @@
 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.sr_datamodule import TVTDataModule
 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
@@ -31,6 +32,7 @@
     "SklearnDataModule",
     "SklearnDataset",
     "TensorDataset",
+    "TVTDataModule",
     "SSLImagenetDataModule",
     "STL10DataModule",
     "VOCDetectionDataModule",

pl_bolts/datamodules/sr_datamodule.py

@@ -0,0 +1,73 @@
+from typing import Any
+
+from pytorch_lightning import LightningDataModule
+from torch.utils.data import DataLoader, Dataset
+
+
+class TVTDataModule(LightningDataModule):
+    """Simple DataModule creating train, val, and test dataloaders from given train, val, and test dataset.
+
+    Example::
+        from pl_bolts.datamodules import TVTDataModule
+        from pl_bolts.datasets.sr_mnist_dataset import SRMNIST
+
+        dataset_dev = SRMNIST(scale_factor=4, root=".", train=True)
+        dataset_train, dataset_val = random_split(dataset_dev, lengths=[55_000, 5_000])
+        dataset_test = SRMNIST(scale_factor=4, root=".", train=True)
+        dm = TVTDataModule(dataset_train, dataset_val, dataset_test)
+    """
+
+    def __init__(
+        self,
+        dataset_train: Dataset,
+        dataset_val: Dataset,
+        dataset_test: Dataset,
+        batch_size: int = 16,
+        shuffle: bool = True,
+        num_workers: int = 8,
+        pin_memory: bool = True,
+        drop_last: bool = True,
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
+        """
+        Args:
+            dataset_train: Train dataset
+            dataset_val: Val dataset
+            dataset_test: Test dataset
+            batch_size: How many samples per batch to load
+            num_workers: How many workers to use for loading data
+            shuffle: If true shuffles the train data every epoch
+            pin_memory: If true, the data loader will copy Tensors into CUDA pinned memory before
+                returning them
+            drop_last: If true drops the last incomplete batch
+        """
+        super().__init__()
+
+        self.dataset_train = dataset_train
+        self.dataset_val = dataset_val
+        self.dataset_test = dataset_test
+        self.num_workers = num_workers
+        self.batch_size = batch_size
+        self.shuffle = shuffle
+        self.pin_memory = pin_memory
+        self.drop_last = drop_last
+
+    def train_dataloader(self) -> DataLoader:
+        return self._dataloader(self.dataset_train, shuffle=self.shuffle)
+
+    def val_dataloader(self) -> DataLoader:
+        return self._dataloader(self.dataset_val, shuffle=False)
+
+    def test_dataloader(self) -> DataLoader:
+        return self._dataloader(self.dataset_test, shuffle=False)
+
+    def _dataloader(self, dataset: Dataset, shuffle: bool = True) -> DataLoader:
+        return DataLoader(
+            dataset,
+            batch_size=self.batch_size,
+            shuffle=shuffle,
+            num_workers=self.num_workers,
+            drop_last=self.drop_last,
+            pin_memory=self.pin_memory,
+        )

pl_bolts/datasets/sr_celeba_dataset.py

@@ -0,0 +1,33 @@
+import os
+from typing import Any
+
+from pl_bolts.datasets.sr_dataset_mixin import SRDatasetMixin
+from pl_bolts.utils import _PIL_AVAILABLE, _TORCHVISION_AVAILABLE
+from pl_bolts.utils.warnings import warn_missing_pkg
+
+if _PIL_AVAILABLE:
+    from PIL import Image
+else:  # pragma: no cover
+    warn_missing_pkg("PIL", pypi_name="Pillow")
+
+if _TORCHVISION_AVAILABLE:
+    from torchvision.datasets import CelebA
+else:  # pragma: no cover
+    warn_missing_pkg("torchvision")
+    CelebA = object
+
+
+class SRCelebA(SRDatasetMixin, CelebA):
+    """CelebA dataset that can be used to train Super Resolution models.
+
+    Function __getitem__ (implemented in SRDatasetMixin) returns tuple of high and low resolution image.
+    """
+
+    def __init__(self, scale_factor: int, *args: Any, **kwargs: Any) -> None:
+        hr_image_size = 128
+        lr_image_size = hr_image_size // scale_factor
+        self.image_channels = 3
+        super().__init__(hr_image_size, lr_image_size, self.image_channels, *args, **kwargs)
+
+    def _get_image(self, index: int):
+        return Image.open(os.path.join(self.root, self.base_folder, "img_align_celeba", self.filename[index]))

pl_bolts/datasets/sr_dataset_mixin.py

@@ -0,0 +1,52 @@
+"""Adapted from: https://github.com/https-deeplearning-ai/GANs-Public."""
+from typing import Any, Tuple
+
+import torch
+
+from pl_bolts.utils import _PIL_AVAILABLE, _TORCHVISION_AVAILABLE
+from pl_bolts.utils.warnings import warn_missing_pkg
+
+if _PIL_AVAILABLE:
+    from PIL import Image
+else:  # pragma: no cover
+    warn_missing_pkg("PIL", pypi_name="Pillow")
+
+if _TORCHVISION_AVAILABLE:
+    from torchvision import transforms as transform_lib
+else:  # pragma: no cover
+    warn_missing_pkg("torchvision")
+
+
+class SRDatasetMixin:
+    """Mixin for Super Resolution datasets.
+
+    Scales range of high resolution images to [-1, 1] and range or low resolution images to [0, 1].
+    """
+
+    def __init__(self, hr_image_size: int, lr_image_size: int, image_channels: int, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+
+        self.hr_transforms = transform_lib.Compose(
+            [
+                transform_lib.RandomCrop(hr_image_size),
+                transform_lib.ToTensor(),
+                transform_lib.Normalize(mean=(0.5,) * image_channels, std=(0.5,) * image_channels),
+            ]
+        )
+
+        self.lr_transforms = transform_lib.Compose(
+            [
+                transform_lib.Normalize(mean=(-1.0,) * image_channels, std=(2.0,) * image_channels),
+                transform_lib.ToPILImage(),
+                transform_lib.Resize(lr_image_size, Image.BICUBIC),
+                transform_lib.ToTensor(),
+            ]
+        )
+
+    def __getitem__(self, index: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        image = self._get_image(index)
+
+        hr_image = self.hr_transforms(image)
+        lr_image = self.lr_transforms(hr_image)
+
+        return hr_image, lr_image

pl_bolts/datasets/sr_mnist_dataset.py

@@ -0,0 +1,27 @@
+from typing import Any
+
+from pl_bolts.datasets.mnist_dataset import MNIST
+from pl_bolts.datasets.sr_dataset_mixin import SRDatasetMixin
+from pl_bolts.utils import _PIL_AVAILABLE
+from pl_bolts.utils.warnings import warn_missing_pkg
+
+if _PIL_AVAILABLE:
+    from PIL import Image
+else:  # pragma: no cover
+    warn_missing_pkg("PIL", pypi_name="Pillow")
+
+
+class SRMNIST(SRDatasetMixin, MNIST):
+    """MNIST dataset that can be used to train Super Resolution models.
+
+    Function __getitem__ (implemented in SRDatasetMixin) returns tuple of high and low resolution image.
+    """
+
+    def __init__(self, scale_factor: int, *args: Any, **kwargs: Any) -> None:
+        hr_image_size = 28
+        lr_image_size = hr_image_size // scale_factor
+        self.image_channels = 1
+        super().__init__(hr_image_size, lr_image_size, self.image_channels, *args, **kwargs)
+
+    def _get_image(self, index: int):
+        return Image.fromarray(self.data[index].numpy(), mode="L")