Add DCGAN module

CHANGELOG.md

@@ -13,6 +13,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
     [#348](https://github.com/PyTorchLightning/pytorch-lightning-bolts/pull/348),
     [#323](https://github.com/PyTorchLightning/pytorch-lightning-bolts/pull/323))
 - Added data monitor callbacks `ModuleDataMonitor` and `TrainingDataMonitor` ([#285](https://github.com/PyTorchLightning/pytorch-lightning-bolts/pull/285))
+- Added DCGAN module ([#403](https://github.com/PyTorchLightning/pytorch-lightning-bolts/pull/403)) 
 - Added `VisionDataModule` as parent class for `BinaryMNISTDataModule`, `CIFAR10DataModule`, `FashionMNISTDataModule`, 
   and `MNISTDataModule` ([#400](https://github.com/PyTorchLightning/pytorch-lightning-bolts/pull/400))
 - Added GIoU loss ([#347](https://github.com/PyTorchLightning/pytorch-lightning-bolts/pull/347))

docs/source/gans.rst

@@ -40,4 +40,49 @@ Loss curves:
 
 
 .. autoclass:: pl_bolts.models.gans.GAN
-   :noindex:
+   :noindex:
+
+DCGAN
+---------
+DCGAN implementation from the paper `Unsupervised Representation Learning with Deep Convolutional Generative
+Adversarial Networks <https://arxiv.org/pdf/1511.06434.pdf>`_. The implementation is based on the version from
+PyTorch's `examples <https://github.com/pytorch/examples/blob/master/dcgan/main.py>`_.
+
+Implemented by:
+
+    - `Christoph Clement <https://github.com/chris-clem>`_
+
+Example MNIST outputs:
+
+    .. image:: _images/gans/dcgan_mnist_outputs.png
+        :width: 400
+        :alt: DCGAN generated MNIST samples
+
+Example LSUN bedroom outputs:
+
+    .. image:: _images/gans/dcgan_lsun_outputs.png
+        :width: 400
+        :alt: DCGAN generated LSUN bedroom samples
+
+MNIST Loss curves:
+
+    .. image:: _images/gans/dcgan_mnist_dloss.png
+        :width: 200
+        :alt: DCGAN MNIST disc loss
+
+    .. image:: _images/gans/dcgan_mnist_gloss.png
+        :width: 200
+        :alt: DCGAN MNIST gen loss
+
+LSUN Loss curves:
+
+    .. image:: _images/gans/dcgan_lsun_dloss.png
+        :width: 200
+        :alt: DCGAN LSUN disc loss
+
+    .. image:: _images/gans/dcgan_lsun_gloss.png
+        :width: 200
+        :alt: DCGAN LSUN gen loss
+
+.. autoclass:: pl_bolts.models.gans.DCGAN
+   :noindex:

pl_bolts/callbacks/vision/image_generation.py

@@ -6,7 +6,7 @@
 try:
     import torchvision
 except ModuleNotFoundError:
-    warn_missing_pkg('torchvision')  # pragma: no-cover
+    warn_missing_pkg("torchvision")  # pragma: no-cover
 
 
 class TensorboardGenerativeModelImageSampler(Callback):

pl_bolts/datamodules/kitti_datamodule.py

@@ -81,11 +81,6 @@ def __init__(
         self.pin_memory = pin_memory
         self.drop_last = drop_last
 
-        self.default_transforms = transforms.Compose([
-            transforms.ToTensor(),
-            transforms.Normalize(mean=[0.35675976, 0.37380189, 0.3764753], std=[0.32064945, 0.32098866, 0.32325324])
-        ])
-
         # split into train, val, test
         kitti_dataset = KittiDataset(self.data_dir, transform=self._default_transforms())
 

pl_bolts/models/gans/__init__.py

@@ -1 +1,2 @@
 from pl_bolts.models.gans.basic.basic_gan_module import GAN  # noqa: F401
+from pl_bolts.models.gans.dcgan.dcgan_module import DCGAN  # noqa: F401

pl_bolts/models/gans/dcgan/components.py

@@ -0,0 +1,95 @@
+# Based on https://github.com/pytorch/examples/blob/master/dcgan/main.py
+import torch
+from torch import nn
+
+
+class DCGANGenerator(nn.Module):
+
+    def __init__(self, latent_dim: int, feature_maps: int, image_channels: int) -> None:
+        """
+        Args:
+            latent_dim: Dimension of the latent space
+            feature_maps: Number of feature maps to use
+            image_channels: Number of channels of the images from the dataset
+        """
+        super().__init__()
+        self.gen = nn.Sequential(
+            self._make_gen_block(latent_dim, feature_maps * 8, kernel_size=4, stride=1, padding=0),
+            self._make_gen_block(feature_maps * 8, feature_maps * 4),
+            self._make_gen_block(feature_maps * 4, feature_maps * 2),
+            self._make_gen_block(feature_maps * 2, feature_maps),
+            self._make_gen_block(feature_maps, image_channels, last_block=True),
+        )
+
+    @staticmethod
+    def _make_gen_block(
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int = 4,
+        stride: int = 2,
+        padding: int = 1,
+        bias: bool = False,
+        last_block: bool = False,
+    ) -> nn.Sequential:
+        if not last_block:
+            gen_block = nn.Sequential(
+                nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride, padding, bias=bias),
+                nn.BatchNorm2d(out_channels),
+                nn.ReLU(True),
+            )
+        else:
+            gen_block = nn.Sequential(
+                nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride, padding, bias=bias),
+                nn.Tanh(),
+            )
+
+        return gen_block
+
+    def forward(self, noise: torch.Tensor) -> torch.Tensor:
+        return self.gen(noise)
+
+
+class DCGANDiscriminator(nn.Module):
+
+    def __init__(self, feature_maps: int, image_channels: int) -> None:
+        """
+        Args:
+            feature_maps: Number of feature maps to use
+            image_channels: Number of channels of the images from the dataset
+        """
+        super().__init__()
+        self.disc = nn.Sequential(
+            self._make_disc_block(image_channels, feature_maps, batch_norm=False),
+            self._make_disc_block(feature_maps, feature_maps * 2),
+            self._make_disc_block(feature_maps * 2, feature_maps * 4),
+            self._make_disc_block(feature_maps * 4, feature_maps * 8),
+            self._make_disc_block(feature_maps * 8, 1, kernel_size=4, stride=1, padding=0, last_block=True),
+        )
+
+    @staticmethod
+    def _make_disc_block(
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int = 4,
+        stride: int = 2,
+        padding: int = 1,
+        bias: bool = False,
+        batch_norm: bool = True,
+        last_block: bool = False,
+    ) -> nn.Sequential:
+        if not last_block:
+            disc_block = nn.Sequential(
+                nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, bias=bias),
+                nn.BatchNorm2d(out_channels) if batch_norm else nn.Identity(),
+                nn.LeakyReLU(0.2, inplace=True),
+            )
+        else:
+            disc_block = nn.Sequential(
+                nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, bias=bias),
+                nn.Sigmoid(),
+            )
+
+        return disc_block
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.disc(x).view(-1, 1).squeeze(1)