Bugs in GANs

deeplay/initializers/constant.py

@@ -21,8 +21,5 @@ def __init__(
         self.weight = weight
         self.bias = bias
 
-    def initialize_weight(self, tensor):
+    def initialize_tensor(self, tensor, name):
         tensor.data.fill_(self.weight)
-
-    def initialize_bias(self, tensor):
-        tensor.data.fill_(self.bias)

deeplay/initializers/initializer.py

@@ -3,15 +3,11 @@ class Initializer:
     def __init__(self, targets):
         self.targets = targets
 
-    def initialize(self, module):
+    def initialize(self, module, tensors=("weight", "bias")):
         if isinstance(module, self.targets):
-            if hasattr(module, "weight") and module.weight is not None:
-                self.initialize_weight(module.weight)
-            if hasattr(module, "bias") and module.bias is not None:
-                self.initialize_bias(module.bias)
+            for tensor in tensors:
+                if hasattr(module, tensor) and getattr(module, tensor) is not None:
+                    self.initialize_tensor(getattr(module, tensor), name=tensor)
 
-    def initialize_weight(self, tensor):
-        pass
-
-    def initialize_bias(self, tensor):
-        pass
+    def initialize_tensor(self, tensor, name):
+        raise NotImplementedError

deeplay/initializers/kaiming.py

@@ -24,13 +24,20 @@ def __init__(
         targets: Tuple[Type[nn.Module], ...] = _kaiming_default_targets,
         mode: str = "fan_out",
         nonlinearity: str = "relu",
+        fill_bias: bool = True,
+        bias: float = 0.0,
     ):
         super().__init__(targets)
         self.mode = mode
         self.nonlinearity = nonlinearity
+        self.fill_bias = fill_bias
+        self.bias = bias
 
-    def initialize_weight(self, tensor):
-        nn.init.kaiming_normal_(tensor, mode=self.mode, nonlinearity=self.nonlinearity)
+    def initialize_tensor(self, tensor, name):
 
-    def initialize_bias(self, tensor):
-        tensor.data.fill_(0.0)
+        if name == "bias" and self.fill_bias:
+            tensor.data.fill_(self.bias)
+        else:
+            nn.init.kaiming_normal_(
+                tensor, mode=self.mode, nonlinearity=self.nonlinearity
+            )

deeplay/initializers/normal.py

@@ -28,8 +28,5 @@ def __init__(
         self.mean = mean
         self.std = std
 
-    def initialize_bias(self, tensor):
-        tensor.data.fill_(self.mean)
-
-    def initialize_weight(self, tensor):
-        tensor.data.normal_(self.mean, self.std)
+    def initialize_tensor(self, tensor, name):
+        tensor.data.normal_(mean=self.mean, std=self.std)

deeplay/models/discriminators/cyclegan.py

@@ -11,7 +11,10 @@
 @ConvolutionalEncoder2d.register_style
 def cyclegan_discriminator(encoder: ConvolutionalEncoder2d):
     encoder[..., "layer"].configure(kernel_size=4, padding=1)
-    encoder["blocks", 1:-1].all.normalized(nn.InstanceNorm2d, mode="insert", after="layer")
+    encoder["blocks", 1:-1].all.normalized(
+        nn.InstanceNorm2d, mode="insert", after="layer"
+    )
+    encoder["blocks", :].all.remove("pool", allow_missing=True)
     encoder["blocks", :-1].configure("activation", nn.LeakyReLU, negative_slope=0.2)
     encoder["blocks", :-2].configure(stride=2)
 

deeplay/models/discriminators/dcgan.py

@@ -13,19 +13,20 @@ def dcgan_discriminator(encoder: ConvolutionalEncoder2d):
     encoder.blocks[-1].configure("layer", padding=0)
     encoder["blocks", :].all.remove("pool", allow_missing=True)
     encoder["blocks", 1:-1].all.normalized()
-    encoder["block", :-1].all.configure("actication", nn.LeakyReLU, negative_slope=0.2)
+    encoder["blocks", :-1].all.configure("activation", nn.LeakyReLU, negative_slope=0.2)
     encoder.blocks[-1].activation.configure(nn.Sigmoid)
 
-    init = Normal(
+    initializer = Normal(
         targets=(
             nn.Conv2d,
             nn.BatchNorm2d,
+            nn.Embedding,
             nn.Linear,
         ),
         mean=0,
         std=0.02,
     )
-    encoder.initialize(init)
+    encoder.initialize(initializer, tensors="weight")
 
 
 class DCGANDiscriminator(ConvolutionalEncoder2d):
@@ -134,9 +135,10 @@ def forward(self, x, y=None):
             )
 
         if self.class_conditioned_model:
-            assert (
-                y is not None
-            ), "Class label y must be provided for class-conditional discriminator"
+            if y is None:
+                raise ValueError(
+                    "Class label y must be provided for class-conditional discriminator"
+                )
 
             y = self.label_embedding(y)
             y = y.view(-1, 1, 64, 64)

deeplay/models/generators/cyclegan.py

@@ -21,19 +21,23 @@ def cyclegan_resnet_encoder(encoder: ConvolutionalEncoder2d):
     encoder.strided(2)
     encoder.normalized(Layer(nn.InstanceNorm2d))
     encoder.blocks.configure(order=["layer", "normalization", "activation"])
-    encoder.blocks[0].prepend(Layer(nn.ReflectionPad2d, 3))
-    encoder.blocks[0].configure("layer", kernel_size=7, stride=1, padding=0)
+    encoder.blocks[0].configure(
+        "layer", kernel_size=7, stride=1, padding=3, padding_mode="reflect"
+    )
 
 
 @ConvolutionalDecoder2d.register_style
 def cyclegan_resnet_decoder(decoder: ConvolutionalDecoder2d):
-    decoder.normalized(Layer(nn.InstanceNorm2d))
+    decoder["blocks", :-1].all.normalized(
+        nn.InstanceNorm2d, mode="insert", after="layer"
+    )
     decoder.blocks.configure(order=["layer", "normalization", "activation"])
     decoder.blocks[:-1].configure(
         "layer", nn.ConvTranspose2d, stride=2, output_padding=1
     )
-    decoder.blocks[-1].configure(kernel_size=7, stride=1, padding=0)
-    decoder.blocks[-1].prepend(Layer(nn.ReflectionPad2d, 3))
+    decoder.blocks[-1].configure(
+        "layer", kernel_size=7, stride=1, padding=3, padding_mode="reflect"
+    )
 
 
 @ConvolutionalNeuralNetwork.register_style
@@ -67,7 +71,7 @@ class CycleGANResnetGenerator(ConvolutionalEncoderDecoder2d):
 
     Examples
     --------
-    >>> generator = CycleGANGenerator(in_channels=1, out_channels=3)
+    >>> generator = CycleGANResnetGenerator(in_channels=1, out_channels=3)
     >>> generator.build()
     >>> x = torch.randn(1, 1, 256, 256)
     >>> y = generator(x)

deeplay/models/generators/dcgan.py

@@ -13,17 +13,17 @@ def dcgan_generator(generator: ConvolutionalDecoder2d):
         "layer", nn.ConvTranspose2d, kernel_size=4, stride=2, padding=1
     ).remove("upsample", allow_missing=True)
     generator.blocks[0].layer.configure(stride=1, padding=0)
-    init = Normal(
+
+    initializer = Normal(
         targets=(
             nn.ConvTranspose2d,
             nn.BatchNorm2d,
             nn.Embedding,
-            nn.Linear,
         ),
         mean=0,
         std=0.02,
     )
-    generator.initialize(init)
+    generator.initialize(initializer, tensors="weight")
 
 
 class DCGANGenerator(ConvolutionalDecoder2d):
@@ -60,11 +60,11 @@ class DCGANGenerator(ConvolutionalDecoder2d):
 
     Examples
     --------
-    >>> generator = DCGAN_Generator(latent_dim=100, output_channels=1, class_conditioned_model=False)
+    >>> generator = DCGANGenerator(latent_dim=100, output_channels=1, class_conditioned_model=False)
     >>> generator.build()
     >>> batch_size = 16
     >>> input = torch.randn([batch_size, 100, 1, 1])
-    >>> output = generator(input)
+    >>> output = generator(x=input, y=None)
 
     Return Values
     -------------
@@ -81,16 +81,16 @@ class DCGANGenerator(ConvolutionalDecoder2d):
     def __init__(
         self,
         latent_dim: int = 100,
-        features_dim: int = 128,
+        features_dim: int = 64,
         out_channels: int = 1,
         class_conditioned_model: bool = False,
         embedding_dim: int = 100,
         num_classes: int = 10,
-        output_channels=None
+        output_channels=None,
     ):
         if output_channels is not None:
             out_channels = output_channels
-            
+
         self.latent_dim = latent_dim
         self.output_channels = out_channels
         self.class_conditioned_model = class_conditioned_model
@@ -104,10 +104,10 @@ def __init__(
         super().__init__(
             in_channels=in_channels,
             hidden_channels=[
+                features_dim * 16,
                 features_dim * 8,
                 features_dim * 4,
                 features_dim * 2,
-                features_dim * 1,
             ],
             out_channels=out_channels,
             out_activation=Layer(nn.Tanh),
@@ -122,9 +122,10 @@ def __init__(
 
     def forward(self, x, y=None):
         if self.class_conditioned_model:
-            assert (
-                y is not None
-            ), "Class label y must be provided for class-conditional generator"
+            if y is None:
+                raise ValueError(
+                    "Class label y must be provided for class-conditional generator"
+                )
 
             y = self.label_embedding(y)
             y = y.view(-1, self.embedding_dim, 1, 1)

deeplay/module.py

@@ -935,7 +935,12 @@ def predict(
                     if not isinstance(item, torch.Tensor):
                         if isinstance(item, np.ndarray):
                             batch[i] = torch.from_numpy(item).to(device)
-                            if batch[i].dtype in [torch.float64, torch.float32, torch.float16, torch.float]:
+                            if batch[i].dtype in [
+                                torch.float64,
+                                torch.float32,
+                                torch.float16,
+                                torch.float,
+                            ]:
                                 if hasattr(self, "dtype"):
                                     batch[i] = batch[i].to(self.dtype)
                                 else:
@@ -1017,16 +1022,20 @@ def log_tensor(self, name: str, tensor: torch.Tensor):
         """
         self.logs[name] = tensor
 
-    def initialize(self, initializer):
+    def initialize(
+        self, initializer, tensors: Union[str, Tuple[str, ...]] = ("weight", "bias")
+    ):
+        if isinstance(tensors, str):
+            tensors = (tensors,)
         for module in self.modules():
             if isinstance(module, DeeplayModule):
-                module._initialize_after_build(initializer)
+                module._initialize_after_build(initializer, tensors)
             else:
-                initializer.initialize(module)
+                initializer.initialize(module, tensors)
 
     @after_build
-    def _initialize_after_build(self, initializer):
-        initializer.initialize(self)
+    def _initialize_after_build(self, initializer, tensors: Tuple[str, ...]):
+        initializer.initialize(self, tensors)
 
     @after_build
     def _validate_after_build(self):
@@ -1208,13 +1217,10 @@ def _give_user_configuration(self, receiver: "DeeplayModule", name) -> bool:
         #                 break
         #         else:
         #             ...
-                # if config_before[key].value != config_after[key].value:
-                #     any_change = True
-                #     break
-
-
+        # if config_before[key].value != config_after[key].value:
+        #     any_change = True
+        #     break
 
-
         # self._user_config._detached_configurations += (
         #     receiver._user_config._detached_configurations
         # )
@@ -1515,7 +1521,9 @@ def filter(self, func: Callable[[str, nn.Module], bool]) -> "Selection":
 
         return Selection(self.model[0], new_selections)
 
-    def hasattr(self, attr: str, strict=True, include_layer_classtype: bool = True) -> "Selection":
+    def hasattr(
+        self, attr: str, strict=True, include_layer_classtype: bool = True
+    ) -> "Selection":
         """Filter the selection based on whether the modules have a certain attribute.
 
         Note, for layers, the attribute is checked in the layer's classtype
@@ -1545,14 +1553,14 @@ def _filter_fn(name: str, module: nn.Module):
                 if include_layer_classtype and isinstance(module, Layer):
                     return hasattr(module.classtype, attr)
                 return False
-
 
             if strict:
                 from deeplay.list import ReferringLayerList
+
                 if isinstance(getattr(module, attr), ReferringLayerList):
                     return False
             return True
- 
+
         return self.filter(_filter_fn)
 
     def isinstance(

deeplay/tests/models/discriminators/test_cyclegan.py

@@ -0,0 +1,33 @@
+import unittest
+
+import torch
+import torch.nn as nn
+
+from deeplay.models.discriminators.cyclegan import CycleGANDiscriminator
+
+
+class TestCycleGANDiscriminator(unittest.TestCase):
+
+    def test_discriminator_defaults(self):
+
+        discriminator = CycleGANDiscriminator().build()
+
+        self.assertEqual(len(discriminator.blocks), 5)
+        self.assertTrue(
+            all(
+                isinstance(discriminator.blocks.normalization[i], nn.InstanceNorm2d)
+                for i in range(1, 4)
+            )
+        )
+        self.assertTrue(
+            all(
+                isinstance(discriminator.blocks.activation[i], nn.LeakyReLU)
+                for i in range(4)
+            )
+        )
+        self.assertTrue(isinstance(discriminator.blocks[-1].activation, nn.Sigmoid))
+
+        # Test on a batch of 2
+        x = torch.rand(2, 1, 256, 256)
+        output = discriminator(x)
+        self.assertEqual(output.shape, (2, 1, 30, 30))