GAN Variants

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163

import torch
import torch.nn as nn
import torch.nn.functional as F
 
class ConditionalGenerator(nn.Module):
    """
    Basic conditional generator with class embedding.
    
    Conditioning strategy: Concatenate one-hot class embedding with z
    """
    
    def __init__(
        self, 
        latent_dim=100, 
        num_classes=10, 
        img_channels=3,
        ngf=64
    ):
        super().__init__()
        
        # Embed class labels to same dimension as latent
        self.embed = nn.Embedding(num_classes, latent_dim)
        
        # Generator takes z + embedded label
        self.fc = nn.Linear(latent_dim * 2, ngf * 8 * 4 * 4)
        
        self.main = nn.Sequential(
            nn.BatchNorm2d(ngf * 8),
            nn.ReLU(True),
            
            nn.ConvTranspose2d(ngf * 8, ngf * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf * 4),
            nn.ReLU(True),
            
            nn.ConvTranspose2d(ngf * 4, ngf * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf * 2),
            nn.ReLU(True),
            
            nn.ConvTranspose2d(ngf * 2, ngf, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf),
            nn.ReLU(True),
            
            nn.ConvTranspose2d(ngf, img_channels, 4, 2, 1, bias=False),
            nn.Tanh()
        )
    
    def forward(self, z, labels):
        """
        z: [batch, latent_dim] - noise vector
        labels: [batch] - integer class labels
        """
        # Embed class labels
        label_embedding = self.embed(labels)  # [batch, latent_dim]
        
        # Concatenate z and label embedding
        combined = torch.cat([z, label_embedding], dim=1)  # [batch, latent_dim*2]
        
        # Project to spatial dimensions
        x = self.fc(combined)
        x = x.view(-1, 512, 4, 4)
        
        return self.main(x)
 
 
class ConditionalDiscriminator(nn.Module):
    """
    Conditional discriminator that judges reality AND class match.
    
    Conditioning strategy: Embed label, expand to image size, concatenate
    """
    
    def __init__(
        self, 
        num_classes=10, 
        img_channels=3, 
        img_size=64,
        ndf=64
    ):
        super().__init__()
        
        self.img_size = img_size
        
        # Embed class to image-sized tensor
        self.embed = nn.Embedding(num_classes, img_size * img_size)
        
        # Discriminator takes image + embedded label channel
        self.main = nn.Sequential(
            # Input: (img_channels + 1) x img_size x img_size
            nn.Conv2d(img_channels + 1, ndf, 4, 2, 1, bias=False),
            nn.LeakyReLU(0.2, inplace=True),
            
            nn.Conv2d(ndf, ndf * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf * 2),
            nn.LeakyReLU(0.2, inplace=True),
            
            nn.Conv2d(ndf * 2, ndf * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf * 4),
            nn.LeakyReLU(0.2, inplace=True),
            
            nn.Conv2d(ndf * 4, ndf * 8, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf * 8),
            nn.LeakyReLU(0.2, inplace=True),
            
            nn.Conv2d(ndf * 8, 1, 4, 1, 0, bias=False),
            nn.Sigmoid()
        )
    
    def forward(self, img, labels):
        """
        img: [batch, channels, H, W]
        labels: [batch] - integer class labels
        """
        # Create label image
        label_embedding = self.embed(labels)  # [batch, H*W]
        label_image = label_embedding.view(-1, 1, self.img_size, self.img_size)
        
        # Concatenate as additional channel
        combined = torch.cat([img, label_image], dim=1)
        
        return self.main(combined).view(-1)
 
 
# Training a conditional GAN
def train_cgan(G, D, dataloader, num_classes, epochs=100):
    """
    Conditional GAN training loop.
    
    Key difference from unconditional: both G and D receive class labels.
    """
    criterion = nn.BCELoss()
    optimG = torch.optim.Adam(G.parameters(), lr=0.0002, betas=(0.5, 0.999))
    optimD = torch.optim.Adam(D.parameters(), lr=0.0002, betas=(0.5, 0.999))
    
    for epoch in range(epochs):
        for real_images, real_labels in dataloader:
            batch_size = real_images.size(0)
            
            # ==== Train Discriminator ====
            D.zero_grad()
            
            # Real samples with correct labels
            real_output = D(real_images, real_labels)
            real_loss = criterion(real_output, torch.ones(batch_size))
            
            # Fake samples with conditioning
            z = torch.randn(batch_size, 100)
            # Generate with same labels as real batch (or random)
            fake_labels = real_labels  # Use same labels
            fake_images = G(z, fake_labels)
            fake_output = D(fake_images.detach(), fake_labels)
            fake_loss = criterion(fake_output, torch.zeros(batch_size))
            
            d_loss = real_loss + fake_loss
            d_loss.backward()
            optimD.step()
            
            # ==== Train Generator ====
            G.zero_grad()
            
            output = D(fake_images, fake_labels)
            g_loss = criterion(output, torch.ones(batch_size))
            g_loss.backward()
            optimG.step()

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102

import torch
import torch.nn as nn
 
class ACGANDiscriminator(nn.Module):
    """
    Discriminator with auxiliary classifier.
    
    Outputs:
    1. Real/fake probability (standard GAN output)
    2. Class probabilities (auxiliary classifier)
    """
    
    def __init__(self, num_classes, ndf=64):
        super().__init__()
        
        # Shared convolutional backbone
        self.features = nn.Sequential(
            nn.Conv2d(3, ndf, 4, 2, 1, bias=False),
            nn.LeakyReLU(0.2, inplace=True),
            
            nn.Conv2d(ndf, ndf * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf * 2),
            nn.LeakyReLU(0.2, inplace=True),
            
            nn.Conv2d(ndf * 2, ndf * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf * 4),
            nn.LeakyReLU(0.2, inplace=True),
            
            nn.Conv2d(ndf * 4, ndf * 8, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf * 8),
            nn.LeakyReLU(0.2, inplace=True),
        )
        
        # Head 1: Real/Fake discrimination
        self.disc_head = nn.Sequential(
            nn.Conv2d(ndf * 8, 1, 4, 1, 0, bias=False),
            nn.Sigmoid()
        )
        
        # Head 2: Class classification
        self.aux_head = nn.Sequential(
            nn.Conv2d(ndf * 8, num_classes, 4, 1, 0, bias=False),
            nn.Flatten()
            # No softmax here - use CrossEntropyLoss which includes it
        )
    
    def forward(self, x):
        features = self.features(x)
        
        disc_output = self.disc_head(features).view(-1)
        class_output = self.aux_head(features)
        
        return disc_output, class_output
 
 
def train_acgan(G, D, dataloader, num_classes, epochs=100):
    """
    ACGAN training with both adversarial and classification loss.
    """
    adv_loss = nn.BCELoss()
    aux_loss = nn.CrossEntropyLoss()
    
    optimG = torch.optim.Adam(G.parameters(), lr=0.0002, betas=(0.5, 0.999))
    optimD = torch.optim.Adam(D.parameters(), lr=0.0002, betas=(0.5, 0.999))
    
    for epoch in range(epochs):
        for real_images, real_labels in dataloader:
            batch_size = real_images.size(0)
            
            # ==== Train Discriminator ====
            D.zero_grad()
            
            # Real images
            disc_real, class_real = D(real_images)
            d_loss_real = adv_loss(disc_real, torch.ones(batch_size))
            d_loss_class_real = aux_loss(class_real, real_labels)
            
            # Fake images
            z = torch.randn(batch_size, 100)
            fake_labels = torch.randint(0, num_classes, (batch_size,))
            fake_images = G(z, fake_labels)
            
            disc_fake, class_fake = D(fake_images.detach())
            d_loss_fake = adv_loss(disc_fake, torch.zeros(batch_size))
            d_loss_class_fake = aux_loss(class_fake, fake_labels)
            
            # Total D loss: adversarial + classification for both real and fake
            d_loss = d_loss_real + d_loss_fake + d_loss_class_real + d_loss_class_fake
            d_loss.backward()
            optimD.step()
            
            # ==== Train Generator ====
            G.zero_grad()
            
            disc_output, class_output = D(fake_images)
            g_loss_adv = adv_loss(disc_output, torch.ones(batch_size))
            g_loss_aux = aux_loss(class_output, fake_labels)
            
            # G wants to fool D AND produce correct class
            g_loss = g_loss_adv + g_loss_aux
            g_loss.backward()
            optimG.step()

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159

import torch
import torch.nn as nn
 
class UNetGenerator(nn.Module):
    """
    U-Net generator for pix2pix.
    
    Encoder-decoder architecture with skip connections.
    Skip connections help preserve spatial details from input.
    
    Structure:
    Input → Enc1 → Enc2 → ... → EncN → DecN → ... → Dec2 → Dec1 → Output
              ↓       ↓              ↑        ↑
              └───────┴── skips ────┴────────┘
    """
    
    def __init__(self, in_channels=3, out_channels=3, ngf=64):
        super().__init__()
        
        # Encoder
        self.enc1 = self._encoder_block(in_channels, ngf, bn=False)  # 64
        self.enc2 = self._encoder_block(ngf, ngf * 2)                 # 128
        self.enc3 = self._encoder_block(ngf * 2, ngf * 4)             # 256
        self.enc4 = self._encoder_block(ngf * 4, ngf * 8)             # 512
        self.enc5 = self._encoder_block(ngf * 8, ngf * 8)             # 512
        self.enc6 = self._encoder_block(ngf * 8, ngf * 8)             # 512
        self.enc7 = self._encoder_block(ngf * 8, ngf * 8)             # 512
        self.enc8 = self._encoder_block(ngf * 8, ngf * 8, bn=False)   # 512 (bottleneck)
        
        # Decoder with skip connections
        self.dec1 = self._decoder_block(ngf * 8, ngf * 8, dropout=True)
        self.dec2 = self._decoder_block(ngf * 16, ngf * 8, dropout=True)  # *2 for skip
        self.dec3 = self._decoder_block(ngf * 16, ngf * 8, dropout=True)
        self.dec4 = self._decoder_block(ngf * 16, ngf * 8)
        self.dec5 = self._decoder_block(ngf * 16, ngf * 4)
        self.dec6 = self._decoder_block(ngf * 8, ngf * 2)
        self.dec7 = self._decoder_block(ngf * 4, ngf)
        
        self.final = nn.Sequential(
            nn.ConvTranspose2d(ngf * 2, out_channels, 4, 2, 1),
            nn.Tanh()
        )
    
    def _encoder_block(self, in_ch, out_ch, bn=True):
        layers = [nn.Conv2d(in_ch, out_ch, 4, 2, 1, bias=False)]
        if bn:
            layers.append(nn.BatchNorm2d(out_ch))
        layers.append(nn.LeakyReLU(0.2, inplace=True))
        return nn.Sequential(*layers)
    
    def _decoder_block(self, in_ch, out_ch, dropout=False):
        layers = [
            nn.ConvTranspose2d(in_ch, out_ch, 4, 2, 1, bias=False),
            nn.BatchNorm2d(out_ch),
        ]
        if dropout:
            layers.append(nn.Dropout(0.5))
        layers.append(nn.ReLU(inplace=True))
        return nn.Sequential(*layers)
    
    def forward(self, x):
        # Encoder pass - save for skip connections
        e1 = self.enc1(x)
        e2 = self.enc2(e1)
        e3 = self.enc3(e2)
        e4 = self.enc4(e3)
        e5 = self.enc5(e4)
        e6 = self.enc6(e5)
        e7 = self.enc7(e6)
        e8 = self.enc8(e7)
        
        # Decoder pass with skip connections
        d1 = self.dec1(e8)
        d2 = self.dec2(torch.cat([d1, e7], dim=1))  # Skip from e7
        d3 = self.dec3(torch.cat([d2, e6], dim=1))
        d4 = self.dec4(torch.cat([d3, e5], dim=1))
        d5 = self.dec5(torch.cat([d4, e4], dim=1))
        d6 = self.dec6(torch.cat([d5, e3], dim=1))
        d7 = self.dec7(torch.cat([d6, e2], dim=1))
        
        return self.final(torch.cat([d7, e1], dim=1))
 
 
class PatchGANDiscriminator(nn.Module):
    """
    PatchGAN discriminator for pix2pix.
    
    Instead of outputting a single real/fake score, outputs a grid
    where each cell judges whether a 70x70 patch is real or fake.
    
    This focuses the discriminator on local texture/style,
    leaving global structure to the L1 loss.
    """
    
    def __init__(self, in_channels=6, ndf=64):  # 6 = input + target
        super().__init__()
        
        self.model = nn.Sequential(
            # 70x70 PatchGAN
            nn.Conv2d(in_channels, ndf, 4, 2, 1),
            nn.LeakyReLU(0.2, inplace=True),
            
            nn.Conv2d(ndf, ndf * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf * 2),
            nn.LeakyReLU(0.2, inplace=True),
            
            nn.Conv2d(ndf * 2, ndf * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf * 4),
            nn.LeakyReLU(0.2, inplace=True),
            
            nn.Conv2d(ndf * 4, ndf * 8, 4, 1, 1, bias=False),
            nn.BatchNorm2d(ndf * 8),
            nn.LeakyReLU(0.2, inplace=True),
            
            nn.Conv2d(ndf * 8, 1, 4, 1, 1),
            # No sigmoid - use BCEWithLogitsLoss
        )
    
    def forward(self, input_img, target_img):
        # Concatenate input and target
        x = torch.cat([input_img, target_img], dim=1)
        return self.model(x)
 
 
# pix2pix loss
def pix2pix_loss(G, D, input_img, target_img, lambda_l1=100):
    """
    pix2pix objective: adversarial + L1.
    
    cGAN Loss: The adversarial loss
    L1 Loss: Encourages output to be similar to target
    
    The L1 term is crucial - without it, outputs are realistic
    but may not match the target content.
    """
    fake_img = G(input_img)
    
    # Discriminator loss
    pred_real = D(input_img, target_img)
    pred_fake = D(input_img, fake_img.detach())
    
    loss_D_real = F.binary_cross_entropy_with_logits(
        pred_real, torch.ones_like(pred_real)
    )
    loss_D_fake = F.binary_cross_entropy_with_logits(
        pred_fake, torch.zeros_like(pred_fake)
    )
    loss_D = (loss_D_real + loss_D_fake) * 0.5
    
    # Generator loss
    pred_fake_for_G = D(input_img, fake_img)
    loss_G_adv = F.binary_cross_entropy_with_logits(
        pred_fake_for_G, torch.ones_like(pred_fake_for_G)
    )
    loss_G_l1 = F.l1_loss(fake_img, target_img)
    
    loss_G = loss_G_adv + lambda_l1 * loss_G_l1
    
    return loss_D, loss_G

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115

import torch
import torch.nn as nn
import torch.nn.functional as F
 
class CycleGAN:
    """
    CycleGAN for unpaired image-to-image translation.
    
    Key losses:
    1. Adversarial: G(x) should look like Y, F(y) should look like X
    2. Cycle consistency: F(G(x)) ≈ x and G(F(y)) ≈ y
    3. Identity (optional): G(y) ≈ y and F(x) ≈ x
    """
    
    def __init__(self, G, F, D_X, D_Y, lambda_cycle=10, lambda_identity=0.5):
        """
        G: Generator X → Y
        F: Generator Y → X
        D_X: Discriminator for domain X
        D_Y: Discriminator for domain Y
        """
        self.G = G
        self.F = F
        self.D_X = D_X
        self.D_Y = D_Y
        self.lambda_cycle = lambda_cycle
        self.lambda_identity = lambda_identity
    
    def generator_loss(self, real_X, real_Y):
        """
        Generator loss for both G and F.
        """
        # ----- Forward cycle: X → Y → X -----
        fake_Y = self.G(real_X)           # G(X)
        reconstructed_X = self.F(fake_Y)  # F(G(X))
        
        # ----- Backward cycle: Y → X → Y -----
        fake_X = self.F(real_Y)           # F(Y)
        reconstructed_Y = self.G(fake_X)  # G(F(Y))
        
        # ----- Adversarial loss -----
        # G wants D_Y to think fake_Y is real
        pred_fake_Y = self.D_Y(fake_Y)
        loss_G_adv = F.mse_loss(pred_fake_Y, torch.ones_like(pred_fake_Y))
        
        # F wants D_X to think fake_X is real
        pred_fake_X = self.D_X(fake_X)
        loss_F_adv = F.mse_loss(pred_fake_X, torch.ones_like(pred_fake_X))
        
        # ----- Cycle consistency loss -----
        loss_cycle_X = F.l1_loss(reconstructed_X, real_X)
        loss_cycle_Y = F.l1_loss(reconstructed_Y, real_Y)
        loss_cycle = loss_cycle_X + loss_cycle_Y
        
        # ----- Identity loss (optional, helps preserve color) -----
        # G should be identity on Y, F should be identity on X
        identity_Y = self.G(real_Y)
        identity_X = self.F(real_X)
        loss_identity = F.l1_loss(identity_Y, real_Y) + F.l1_loss(identity_X, real_X)
        
        # Total generator loss
        loss_G = (loss_G_adv + loss_F_adv + 
                 self.lambda_cycle * loss_cycle + 
                 self.lambda_identity * self.lambda_cycle * loss_identity)
        
        return loss_G, {
            'adv_G': loss_G_adv.item(),
            'adv_F': loss_F_adv.item(),
            'cycle': loss_cycle.item()
        }
    
    def discriminator_loss(self, real_X, real_Y):
        """
        Discriminator loss for both D_X and D_Y.
        """
        # Generate fakes
        fake_Y = self.G(real_X).detach()
        fake_X = self.F(real_Y).detach()
        
        # D_Y: distinguish real_Y from fake_Y
        pred_real_Y = self.D_Y(real_Y)
        pred_fake_Y = self.D_Y(fake_Y)
        loss_D_Y = 0.5 * (
            F.mse_loss(pred_real_Y, torch.ones_like(pred_real_Y)) +
            F.mse_loss(pred_fake_Y, torch.zeros_like(pred_fake_Y))
        )
        
        # D_X: distinguish real_X from fake_X
        pred_real_X = self.D_X(real_X)
        pred_fake_X = self.D_X(fake_X)
        loss_D_X = 0.5 * (
            F.mse_loss(pred_real_X, torch.ones_like(pred_real_X)) +
            F.mse_loss(pred_fake_X, torch.zeros_like(pred_fake_X))
        )
        
        return loss_D_X + loss_D_Y
 
 
# The cycle consistency constraint is crucial:
"""
Without cycle consistency:
- G could map ALL horses to the SAME zebra
- The zebra would be realistic (fooling D_Y)
- But the mapping would lose all input information
 
With cycle consistency:
- If G(horse_1) = zebra_A and G(horse_2) = zebra_A
- Then F(zebra_A) must equal both horse_1 AND horse_2
- This is impossible! F can only output one thing
- So G must preserve enough information to reconstruct
 
The cycle constraint forces meaningful correspondence:
- Each horse maps to a DIFFERENT, corresponding zebra
- The transformation preserves structure/identity
"""