Reviewing GAN basics, VisionDataModule, MNISTDataModule, CIFAR10DataModule

pl_bolts/models/gans/basic/basic_gan_module.py

@@ -2,13 +2,12 @@
 
 import torch
 from pytorch_lightning import LightningModule, Trainer, seed_everything
+from pytorch_lightning.callbacks.progress import TQDMProgressBar
 from torch.nn import functional as F
 
 from pl_bolts.models.gans.basic.components import Discriminator, Generator
-from pl_bolts.utils.stability import under_review
 
 
-@under_review()
 class GAN(LightningModule):
     """Vanilla GAN implementation.
 
@@ -22,7 +21,7 @@ class GAN(LightningModule):
     Example CLI::
 
         # mnist
-        python  basic_gan_module.py --gpus 1
+        python basic_gan_module.py --gpus 1
 
         # imagenet
         python  basic_gan_module.py --gpus 1 --dataset 'imagenet2012'
@@ -166,7 +165,6 @@ def add_model_specific_args(parent_parser):
         return parser
 
 
-@under_review()
 def cli_main(args=None):
     from pl_bolts.callbacks import LatentDimInterpolator, TensorboardGenerativeModelImageSampler
     from pl_bolts.datamodules import CIFAR10DataModule, ImagenetDataModule, MNISTDataModule, STL10DataModule
@@ -193,8 +191,12 @@ def cli_main(args=None):
 
     dm = dm_cls.from_argparse_args(args)
     model = GAN(*dm.size(), **vars(args))
-    callbacks = [TensorboardGenerativeModelImageSampler(), LatentDimInterpolator(interpolate_epoch_interval=5)]
-    trainer = Trainer.from_argparse_args(args, callbacks=callbacks, progress_bar_refresh_rate=20)
+    callbacks = [
+        TensorboardGenerativeModelImageSampler(),
+        LatentDimInterpolator(interpolate_epoch_interval=5),
+        TQDMProgressBar(refresh_rate=20),
+    ]
+    trainer = Trainer.from_argparse_args(args, callbacks=callbacks)
     trainer.fit(model, datamodule=dm)
     return dm, model, trainer
 

pl_bolts/models/gans/basic/components.py

@@ -3,10 +3,7 @@
 from torch import nn
 from torch.nn import functional as F
 
-from pl_bolts.utils.stability import under_review
 
-
-@under_review()
 class Generator(nn.Module):
     def __init__(self, latent_dim, img_shape, hidden_dim=256):
         super().__init__()
@@ -27,7 +24,6 @@ def forward(self, z):
         return img
 
 
-@under_review()
 class Discriminator(nn.Module):
     def __init__(self, img_shape, hidden_dim=1024):
         super().__init__()

tests/models/test_gans.py → tests/models/gans/integration/test_gans.py

@@ -3,7 +3,7 @@
 from torch.utils.data.dataloader import DataLoader
 from torchvision import transforms as transform_lib
 
-from pl_bolts.datamodules import CIFAR10DataModule, MNISTDataModule
+from pl_bolts.datamodules import CIFAR10DataModule, MNISTDataModule, STL10DataModule
 from pl_bolts.datasets.sr_mnist_dataset import SRMNIST
 from pl_bolts.models.gans import DCGAN, GAN, SRGAN, SRResNet
 
@@ -13,6 +13,7 @@
     [
         pytest.param(MNISTDataModule, id="mnist"),
         pytest.param(CIFAR10DataModule, id="cifar10"),
+        pytest.param(STL10DataModule, id="stl10"),
     ],
 )
 def test_gan(tmpdir, datadir, dm_cls):

tests/models/gans/unit/test_basic_components.py

@@ -0,0 +1,38 @@
+import pytest
+import torch
+from pytorch_lightning import seed_everything
+
+from pl_bolts.models.gans.basic.components import Discriminator, Generator
+
+
+@pytest.mark.parametrize(
+    "latent_dim, img_shape",
+    [
+        pytest.param(100, (3, 28, 28), id="100-multichannel"),
+        pytest.param(100, (1, 28, 28), id="100-singlechannel"),
+    ],
+)
+def test_generator(latent_dim, img_shape):
+    batch_dim = 10
+    seed_everything()
+    generator = Generator(latent_dim=latent_dim, img_shape=img_shape)
+    noise = torch.randn(batch_dim, latent_dim)
+    samples = generator(noise)
+    assert samples.shape == (batch_dim, *img_shape)
+
+
+@pytest.mark.parametrize(
+    "img_shape",
+    [
+        pytest.param((3, 28, 28), id="discriminator-multichannel"),
+        pytest.param((1, 28, 28), id="discriminator-singlechannel"),
+    ],
+)
+def test_discriminator(img_shape):
+    batch_dim = 10
+    seed_everything()
+    discriminator = Discriminator(img_shape=img_shape)
+    samples = torch.randn(batch_dim, *img_shape)
+    real_or_fake = discriminator(samples)
+    assert real_or_fake.shape == (batch_dim, 1)
+    assert (torch.clamp(real_or_fake.clone(), 0, 1) == real_or_fake).all()