Problem 174: Issue with test cases

[mavleo96]

The "Train a Simple GAN on 1D Gaussian Data" problem has incorrect test cases.

The GAN after training for 500-1500 epoch would already start resembling the real data but the test cases expect the test cases still expect the data to be zero mean, zero dev.

I think solution needs to reviewed based of which these test cases are generated.

[moe18]

will review the question, thank you for letting me know

[moe18]

would you mind sharing your solution?

[mavleo96]

would you mind sharing your solution?

import numpy as np

class Generator:
    def __init__(self, latent_dim, hidden_dim):
        self.w1 = np.random.uniform(-1, 1, (hidden_dim, latent_dim)) / np.sqrt(latent_dim)
        self.b1 = np.random.uniform(-1, 1, hidden_dim) / np.sqrt(latent_dim)
        self.w2 = np.random.uniform(-1, 1, (1, hidden_dim)) / np.sqrt(hidden_dim)
        self.b2 = np.random.uniform(-1, 1, 1) / np.sqrt(hidden_dim)

        self.reset_grad()

        self.cache = {
            "x": None,
            "h": None,
            "y": None,
        }

    def reset_grad(self):
        self.grad_w1 = np.zeros_like(self.w1)
        self.grad_b1 = np.zeros_like(self.b1)
        self.grad_w2 = np.zeros_like(self.w2)
        self.grad_b2 = np.zeros_like(self.b2)

    def cache_clear(self):
        self.cache = {k: None for k in self.cache}

    def forward(self, x):
        self.cache_clear()
        h = x @ self.w1.T + self.b1
        h = np.maximum(0, h)
        y = h @ self.w2.T + self.b2

        self.cache["x"] = x
        self.cache["h"] = h
        self.cache["y"] = y

        return y
    
    def backward(self, grad_out):
        x = self.cache["x"]
        h = self.cache["h"]
        y = self.cache["y"]

        self.grad_b2 += grad_out.sum(0)
        self.grad_w2 += grad_out.T @ h

        grad_h = grad_out @ self.w2
        grad_h = np.where(h > 0, grad_h, 0)

        self.grad_b1 += grad_h.sum(0)
        self.grad_w1 += grad_h.T @ x

        grad_in = grad_h @ self.w1
        return grad_in

    def update(self, lr):
        self.w1 -= lr * self.grad_w1
        self.b1 -= lr * self.grad_b1
        self.w2 -= lr * self.grad_w2
        self.b2 -= lr * self.grad_b2

class Discriminator:
    def __init__(self, input_dim, hidden_dim):
        self.w1 = np.random.uniform(-1, 1, (hidden_dim, input_dim)) / np.sqrt(input_dim)
        self.b1 = np.random.uniform(-1, 1, hidden_dim) / np.sqrt(input_dim)
        self.w2 = np.random.uniform(-1, 1, (1, hidden_dim)) / np.sqrt(hidden_dim)
        self.b2 = np.random.uniform(-1, 1, 1) / np.sqrt(hidden_dim)

        self.reset_grad()

        self.cache = {
            "h1": None,
            "y": None,
            "x": None,
        }

    def reset_grad(self):
        self.grad_w1 = np.zeros_like(self.w1)
        self.grad_b1 = np.zeros_like(self.b1)
        self.grad_w2 = np.zeros_like(self.w2)
        self.grad_b2 = np.zeros_like(self.b2)

    def cache_clear(self):
        self.cache = {k: None for k in self.cache}

    def forward(self, x):
        self.cache_clear()
        h1 = x @ self.w1.T + self.b1
        h1 = np.maximum(0, h1)
        h2 = h1 @ self.w2.T + self.b2
        y = 1 / (1 + np.exp(-h2))

        self.cache["x"] = x
        self.cache["h1"] = h1
        self.cache["y"] = y

        return y
    
    def backward(self, grad_out):
        x = self.cache["x"]
        h1 = self.cache["h1"]
        y = self.cache["y"]

        # grad_h2 = grad_out * y * (1 - y)
        grad_h2 = grad_out

        self.grad_b2 += grad_h2.sum(0)
        self.grad_w2 += grad_h2.T @ h1

        grad_h1 = grad_h2 @ self.w2
        grad_h1 = np.where(h1 > 0, grad_h1, 0)

        self.grad_b1 += grad_h1.sum(0)
        self.grad_w1 += grad_h1.T @ x

        grad_in = grad_h1 @ self.w1
        return grad_in

    def update(self, lr):
        self.w1 -= lr * self.grad_w1
        self.b1 -= lr * self.grad_b1
        self.w2 -= lr * self.grad_w2
        self.b2 -= lr * self.grad_b2

def train_gan(mean_real: float, std_real: float, latent_dim: int = 1, hidden_dim: int = 16, learning_rate: float = 0.001, epochs: int = 5000, batch_size: int = 128, seed: int = 42):
    """
    Train a simple GAN to learn a 1D Gaussian distribution.

    Args:
        mean_real: Mean of the target Gaussian
        std_real: Std of the target Gaussian
        latent_dim: Dimension of the noise input to the generator
        hidden_dim: Hidden layer size for both networks
        learning_rate: Learning rate for gradient descent
        epochs: Number of training epochs
        batch_size: Training batch size
        seed: Random seed for reproducibility

    Returns:
        gen_forward: A function that takes z and returns generated samples
    """
    np.random.seed(seed)
    g = Generator(latent_dim, hidden_dim)
    d = Discriminator(1, hidden_dim)

    losses = {"g": [], "d": []}
    for _ in range(epochs):
        d.reset_grad(); g.reset_grad()

        x_real = np.random.normal(mean_real, std_real, (batch_size, 1))
        prob_real = d.forward(x_real)
        target_real = np.ones_like(prob_real)
        grad_real = (prob_real - target_real) / batch_size
        d.backward(grad_real)

        x_fake = g.forward(np.random.randn(batch_size, latent_dim))
        prob_fake = d.forward(x_fake)
        target_fake = np.zeros_like(prob_fake)
        grad_fake = (prob_fake - target_fake) / batch_size
        d.backward(grad_fake)
        d.update(learning_rate)

        loss_d = -(np.log(prob_real + 1e-10) + np.log(1 - prob_fake + 1e-10)).mean()
        losses["d"].append(loss_d)

        d.reset_grad(); g.reset_grad()

        x_fake = g.forward(np.random.randn(batch_size, latent_dim))
        prob_fake = d.forward(x_fake)
        target_fake = np.ones_like(prob_fake)
        grad_fake = (prob_fake - target_fake) / batch_size
        g.backward(d.backward(grad_fake))
        g.update(learning_rate)

        loss_g = -np.log(prob_fake + 1e-10).mean()
        losses["g"].append(loss_g)

    return lambda z: (g.forward(z), None, None)

[mavleo96]

If you are able to update the repo with the solution for this question then i can take a look at this myself

[Open-Deep-ML]

Just uploaded the current solution
https://github.com/Open-Deep-ML/DML-OpenProblem/tree/main/questions/174_train-a-simple-gan-on-1d-gaussian-data

[mavleo96]

I just looked at the solution and what seems to be missing in the question is weights initialization strategy. I have used kaiming uniform init which is the torch default whereas the solution uses random normal (0, 0.01) for weights and zero init for bias.
I think the question should add details on weight initialization for clarity.

Additional comments:
The init strategy used in solution is not optimal and when you plot losses we can notice than the gan ends up learning nothing. We should definitely update the solution with better init (like xavier, lecun, kaiming) and test cases accordingly.