GAN Basics

z (noise) → G → fake image ─┐
                              ├─→ D → real / fake ?
real image  ─────────────────┘

min_G  max_D  E[log D(x)] + E[log(1 - D(G(z)))]

import torch, torch.nn as nn

class G(nn.Module):
    def __init__(self, z=100, ch=64):
        super().__init__()
        self.net = nn.Sequential(
            nn.ConvTranspose2d(z, ch*8, 4, 1, 0), nn.BatchNorm2d(ch*8), nn.ReLU(True),
            nn.ConvTranspose2d(ch*8, ch*4, 4, 2, 1), nn.BatchNorm2d(ch*4), nn.ReLU(True),
            nn.ConvTranspose2d(ch*4, ch*2, 4, 2, 1), nn.BatchNorm2d(ch*2), nn.ReLU(True),
            nn.ConvTranspose2d(ch*2, ch,   4, 2, 1), nn.BatchNorm2d(ch),   nn.ReLU(True),
            nn.ConvTranspose2d(ch, 3, 4, 2, 1), nn.Tanh())
    def forward(self, z): return self.net(z)

class D(nn.Module):
    def __init__(self, ch=64):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv2d(3, ch, 4, 2, 1), nn.LeakyReLU(0.2, True),
            nn.Conv2d(ch, ch*2, 4, 2, 1), nn.BatchNorm2d(ch*2), nn.LeakyReLU(0.2, True),
            nn.Conv2d(ch*2, ch*4, 4, 2, 1), nn.BatchNorm2d(ch*4), nn.LeakyReLU(0.2, True),
            nn.Conv2d(ch*4, 1, 4, 1, 0), nn.Sigmoid())
    def forward(self, x): return self.net(x).view(-1)

for real in loader:
    # ---- D ----
    z = torch.randn(b, 100, 1, 1, device=dev)
    fake = G(z).detach()
    loss_D = bce(D(real), ones) + bce(D(fake), zeros)
    opt_D.zero_grad(); loss_D.backward(); opt_D.step()

    # ---- G ----
    fake = G(torch.randn(b, 100, 1, 1, device=dev))
    loss_G = bce(D(fake), ones)   # G চায় D ভাবুক "real"
    opt_G.zero_grad(); loss_G.backward(); opt_G.step()