Upload cgan_architecture.py with huggingface_hub

cgan_architecture.py

@@ -0,0 +1,186 @@
+"""Conditional GAN for MNIST digit generation.
+
+Architecture:
+  Generator:     Embedding(y) → concat with z → ConvTranspose2d upsampling
+  Discriminator: Conv2d downsampling + Projection Discriminator (Miyato & Koyama, ICLR 2018)
+
+The Projection Discriminator incorporates the class label through an inner
+product between the discriminator's penultimate features and a class
+embedding, rather than via naive concatenation.  This is theoretically
+well-grounded and empirically superior.
+
+  D(x, y) = fc(h(x)) + ⟨embed(y), h(x)⟩
+
+where h(x) is the feature extractor output and fc gives the unconditional score.
+
+Reference: "cGANs with Projection Discriminator" (Miyato & Koyama, 2018)
+           https://arxiv.org/abs/1802.05637
+"""
+
+import torch
+import torch.nn as nn
+from torch.nn.utils import spectral_norm
+
+# ---------------------------------------------------------------------------
+# Generator
+# ---------------------------------------------------------------------------
+
+
+class Generator(nn.Module):
+    """Conditional generator: (z, y) → grayscale MNIST image.
+
+    Label embedding is concatenated to the latent vector before the first
+    linear projection, following the cGAN conditioning approach.
+
+    Args:
+        latent_dim:  Dimensionality of the noise vector z.
+        embed_dim:   Dimensionality of the class embedding.
+        base_channels: Channel width of the shallowest decoder stage.
+        num_classes: Number of digit classes (10 for MNIST).
+    """
+
+    def __init__(
+        self,
+        latent_dim: int = 100,
+        embed_dim: int = 100,
+        base_channels: int = 256,
+        num_classes: int = 10,
+    ) -> None:
+        super().__init__()
+        self.latent_dim = latent_dim
+        self.embed_dim = embed_dim
+        self.num_classes = num_classes
+
+        self.embed = nn.Embedding(num_classes, embed_dim)
+
+        in_dim = latent_dim + embed_dim
+        # Project to 7×7 feature map
+        self.project = nn.Sequential(
+            nn.Linear(in_dim, base_channels * 7 * 7, bias=False),
+            nn.BatchNorm1d(base_channels * 7 * 7),
+            nn.ReLU(inplace=True),
+        )
+        self.base_channels = base_channels
+
+        # 7×7 → 14×14 → 28×28
+        self.main = nn.Sequential(
+            nn.ConvTranspose2d(
+                base_channels, base_channels // 2, 4, stride=2, padding=1, bias=False
+            ),
+            nn.BatchNorm2d(base_channels // 2),
+            nn.ReLU(inplace=True),
+            nn.ConvTranspose2d(
+                base_channels // 2, base_channels // 4, 4, stride=2, padding=1, bias=False
+            ),
+            nn.BatchNorm2d(base_channels // 4),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(base_channels // 4, 1, 3, stride=1, padding=1),
+            nn.Tanh(),
+        )
+
+        self._init_weights()
+
+    def _init_weights(self) -> None:
+        for m in self.modules():
+            if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d)):
+                nn.init.normal_(m.weight, 0.0, 0.02)
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.normal_(m.weight, 1.0, 0.02)
+                nn.init.zeros_(m.bias)
+            elif isinstance(m, nn.Embedding):
+                nn.init.normal_(m.weight, 0.0, 0.02)
+
+    def forward(self, z: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            z: Noise tensor, shape (B, latent_dim).
+            y: Class labels, shape (B,), values in [0, num_classes).
+
+        Returns:
+            Generated images, shape (B, 1, 28, 28), range [-1, 1].
+        """
+        e = self.embed(y)  # B × embed_dim
+        inp = torch.cat([z, e], dim=1)  # B × (latent_dim + embed_dim)
+        h = self.project(inp).view(-1, self.base_channels, 7, 7)
+        return self.main(h)
+
+
+# ---------------------------------------------------------------------------
+# Discriminator (Projection)
+# ---------------------------------------------------------------------------
+
+
+class Discriminator(nn.Module):
+    """Projection discriminator for conditional image scoring.
+
+    D(x, y) = fc(h(x)) + ⟨embed(y), h(x)⟩
+
+    Spectral Normalization is applied to all linear and conv layers for
+    Lipschitz regularisation.  No BatchNorm (incompatible with SN).
+
+    Args:
+        base_channels: Channel width of the first downsampling stage.
+        num_classes:   Number of digit classes (10 for MNIST).
+        feature_dim:   Dimensionality of the penultimate feature vector.
+    """
+
+    def __init__(
+        self,
+        base_channels: int = 64,
+        num_classes: int = 10,
+        feature_dim: int = 512,
+    ) -> None:
+        super().__init__()
+        self.num_classes = num_classes
+
+        # 1×28×28 → 14×14 → 7×7 → 4×4 (after final conv)
+        self.features = nn.Sequential(
+            spectral_norm(nn.Conv2d(1, base_channels, 4, stride=2, padding=1, bias=False)),
+            nn.LeakyReLU(0.2, inplace=True),
+            spectral_norm(
+                nn.Conv2d(base_channels, base_channels * 2, 4, stride=2, padding=1, bias=False)
+            ),
+            nn.LeakyReLU(0.2, inplace=True),
+            spectral_norm(
+                nn.Conv2d(base_channels * 2, base_channels * 4, 3, stride=1, padding=1, bias=False)
+            ),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.AdaptiveAvgPool2d(1),
+            nn.Flatten(),
+        )
+        feat_ch = base_channels * 4
+
+        self.fc_feat = spectral_norm(nn.Linear(feat_ch, feature_dim, bias=False))
+
+        # Unconditional score
+        self.fc_out = spectral_norm(nn.Linear(feature_dim, 1, bias=True))
+
+        # Class embedding for projection
+        self.embed = spectral_norm(nn.Embedding(num_classes, feature_dim))
+
+        self._init_weights()
+
+    def _init_weights(self) -> None:
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.normal_(m.weight, 0.0, 0.02)
+            elif isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, 0.0, 0.02)
+            elif isinstance(m, nn.Embedding):
+                nn.init.normal_(m.weight, 0.0, 0.02)
+
+    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: Image tensor, shape (B, 1, 28, 28), range [-1, 1].
+            y: Class labels, shape (B,), values in [0, num_classes).
+
+        Returns:
+            Score tensor, shape (B, 1).
+        """
+        h = self.fc_feat(self.features(x))  # B × feature_dim
+
+        unconditional = self.fc_out(h)  # B × 1
+        projection = (h * self.embed(y)).sum(dim=1, keepdim=True)  # B × 1
+
+        return unconditional + projection