adding codes

vae does not work

Author: mhjensen

doc/Programs/DiffusionModels/vdiff.py

@@ -0,0 +1,141 @@
+import numpy as np
+import matplotlib.pyplot as plt
+import gzip, pickle, os, urllib.request
+
+# ===== Utility functions =====
+def load_mnist():
+    url = 'http://deeplearning.net/data/mnist/mnist.pkl.gz'
+    fname = 'mnist.pkl.gz'
+    if not os.path.exists(fname):
+        urllib.request.urlretrieve(url, fname)
+    with gzip.open(fname, 'rb') as f:
+        train_set, _, _ = pickle.load(f, encoding='latin1')
+    X, _ = train_set
+    return X.astype(np.float32)
+
+def linear_beta_schedule(timesteps, beta_start=1e-4, beta_end=0.02):
+    return np.linspace(beta_start, beta_end, timesteps)
+
+def sigmoid(x): return 1 / (1 + np.exp(-x))
+def relu(x): return np.maximum(0, x)
+
+# ===== Neural network for epsilon_theta =====
+class Dense:
+    def __init__(self, in_dim, out_dim, activation='relu'):
+        self.W = np.random.randn(in_dim, out_dim) * 0.01
+        self.b = np.zeros(out_dim)
+        self.activation = activation
+
+    def forward(self, x):
+        self.input = x
+        self.z = x @ self.W + self.b
+        if self.activation == 'relu':
+            self.out = relu(self.z)
+        elif self.activation == 'linear':
+            self.out = self.z
+        return self.out
+
+    def backward(self, grad_out, lr):
+        if self.activation == 'relu':
+            grad = grad_out * (self.z > 0).astype(float)
+        else:
+            grad = grad_out
+
+        dW = self.input.T @ grad
+        db = np.sum(grad, axis=0)
+        self.W -= lr * dW
+        self.b -= lr * db
+        return grad @ self.W.T
+
+class DenoiseMLP:
+    def __init__(self, input_dim, hidden_dims):
+        dims = [input_dim] + hidden_dims + [input_dim]
+        self.layers = [Dense(dims[i], dims[i+1], 'relu' if i < len(dims)-2 else 'linear') for i in range(len(dims)-1)]
+
+    def forward(self, x):
+        for layer in self.layers:
+            x = layer.forward(x)
+        return x
+
+    def backward(self, grad, lr):
+        for layer in reversed(self.layers):
+            grad = layer.backward(grad, lr)
+
+# ===== Variational Diffusion Model =====
+class DiffusionModel:
+    def __init__(self, img_dim, timesteps=1000, hidden_dims=[512, 256], lr=1e-3):
+        self.T = timesteps
+        self.beta = linear_beta_schedule(self.T)
+        self.alpha = 1.0 - self.beta
+        self.alpha_bar = np.cumprod(self.alpha)
+
+        self.model = DenoiseMLP(input_dim=img_dim, hidden_dims=hidden_dims)
+        self.lr = lr
+        self.img_dim = img_dim
+
+    def q_sample(self, x0, t, noise=None):
+        if noise is None:
+            noise = np.random.randn(*x0.shape)
+        sqrt_alpha_bar = np.sqrt(self.alpha_bar[t])[:, None]
+        sqrt_one_minus_alpha_bar = np.sqrt(1 - self.alpha_bar[t])[:, None]
+        return sqrt_alpha_bar * x0 + sqrt_one_minus_alpha_bar * noise
+
+    def train_step(self, x):
+        N = x.shape[0]
+        t = np.random.randint(0, self.T, size=N)
+        noise = np.random.randn(*x.shape)
+        xt = self.q_sample(x, t, noise)
+        pred_noise = self.model.forward(xt)
+
+        loss = np.mean((pred_noise - noise) ** 2)
+        grad = 2 * (pred_noise - noise) / N
+        self.model.backward(grad, self.lr)
+        return loss
+
+    def train(self, data, epochs=10, batch_size=128):
+        for epoch in range(epochs):
+            perm = np.random.permutation(len(data))
+            total_loss = 0
+            for i in range(0, len(data), batch_size):
+                x = data[perm[i:i+batch_size]]
+                total_loss += self.train_step(x)
+            print(f"Epoch {epoch+1} Loss: {total_loss:.4f}")
+
+    def p_sample(self, xt, t):
+        pred_noise = self.model.forward(xt)
+        alpha = self.alpha[t]
+        alpha_bar = self.alpha_bar[t]
+        beta = self.beta[t]
+
+        coef1 = 1 / np.sqrt(alpha)
+        coef2 = (1 - alpha) / np.sqrt(1 - alpha_bar)
+        mean = coef1 * (xt - coef2 * pred_noise)
+
+        if t > 0:
+            noise = np.random.randn(*xt.shape)
+        else:
+            noise = 0
+        return mean + np.sqrt(beta) * noise
+
+    def sample(self, n=16):
+        xt = np.random.randn(n, self.img_dim)
+        for t in reversed(range(self.T)):
+            xt = self.p_sample(xt, t)
+        return xt
+
+# ===== Visualization =====
+def plot_images(samples, n=8):
+    fig, axs = plt.subplots(1, n, figsize=(n, 1.5))
+    for i in range(n):
+        axs[i].imshow(samples[i].reshape(28, 28), cmap='gray')
+        axs[i].axis('off')
+    plt.suptitle("Generated Samples")
+    plt.show()
+
+# ===== Run full example =====
+if __name__ == "__main__":
+    X = load_mnist()[:5000]
+    model = DiffusionModel(img_dim=784, timesteps=100, hidden_dims=[256, 128], lr=1e-3)
+    model.train(X, epochs=10, batch_size=128)
+    samples = model.sample(n=8)
+    plot_images(samples)

doc/Programs/VAE/vae2.py

@@ -0,0 +1,172 @@
+import numpy as np
+import matplotlib.pyplot as plt
+import gzip
+import pickle
+import urllib.request
+import os
+
+# ----- Utility functions -----
+def load_mnist(normalize=True):
+    url = 'http://deeplearning.net/data/mnist/mnist.pkl.gz'
+    filename = 'mnist.pkl.gz'
+    if not os.path.exists(filename):
+        urllib.request.urlretrieve(url, filename)
+    with gzip.open(filename, 'rb') as f:
+        train_set, _, _ = pickle.load(f, encoding='latin1')
+    X, _ = train_set
+    if normalize:
+        X = X.astype(np.float32)
+    return X
+
+def sigmoid(x):
+    return 1 / (1 + np.exp(-x))
+
+def sigmoid_deriv(x):
+    s = sigmoid(x)
+    return s * (1 - s)
+
+# ----- Layer -----
+class Dense:
+    def __init__(self, in_dim, out_dim, activation='sigmoid'):
+        self.W = np.random.randn(in_dim, out_dim) * 0.01
+        self.b = np.zeros(out_dim)
+        self.activation = activation
+        self.input = None
+        self.z = None
+
+    def forward(self, x):
+        self.input = x
+        self.z = x @ self.W + self.b
+        if self.activation == 'sigmoid':
+            return sigmoid(self.z)
+        elif self.activation == 'linear':
+            return self.z
+        elif self.activation == 'relu':
+            return np.maximum(0, self.z)
+
+    def backward(self, grad_output, learning_rate):
+        if self.activation == 'sigmoid':
+            grad = grad_output * sigmoid_deriv(self.z)
+        elif self.activation == 'relu':
+            grad = grad_output * (self.z > 0).astype(float)
+        else:
+            grad = grad_output
+        grad_W = self.input.T @ grad
+        grad_b = np.sum(grad, axis=0)
+        self.W -= learning_rate * grad_W
+        self.b -= learning_rate * grad_b
+        return grad @ self.W.T
+
+# ----- VAE -----
+class VAE:
+    def __init__(self, input_dim=784, hidden_dims=[256], latent_dim=2, learning_rate=0.01):
+        self.encoder_layers = [Dense(input_dim, hidden_dims[0])]
+        for i in range(1, len(hidden_dims)):
+            self.encoder_layers.append(Dense(hidden_dims[i - 1], hidden_dims[i]))
+        self.W_mu = Dense(hidden_dims[-1], latent_dim, activation='linear')
+        self.W_logvar = Dense(hidden_dims[-1], latent_dim, activation='linear')
+
+        self.decoder_layers = [Dense(latent_dim, hidden_dims[-1])]
+        for i in range(len(hidden_dims) - 1, 0, -1):
+            self.decoder_layers.append(Dense(hidden_dims[i], hidden_dims[i - 1]))
+        self.decoder_layers.append(Dense(hidden_dims[0], input_dim, activation='sigmoid'))
+
+        self.learning_rate = learning_rate
+
+    def encode(self, x):
+        h = x
+        for layer in self.encoder_layers:
+            h = layer.forward(h)
+        mu = self.W_mu.forward(h)
+        logvar = self.W_logvar.forward(h)
+        return mu, logvar
+
+    def reparameterize(self, mu, logvar):
+        std = np.exp(0.5 * logvar)
+        eps = np.random.randn(*mu.shape)
+        return mu + eps * std
+
+    def decode(self, z):
+        h = z
+        for layer in self.decoder_layers:
+            h = layer.forward(h)
+        return h
+
+    def loss(self, recon_x, x, mu, logvar):
+        mse = np.mean((recon_x - x) ** 2)
+        kl = -0.5 * np.mean(1 + logvar - mu ** 2 - np.exp(logvar))
+        return mse + kl
+
+    def train_step(self, x):
+        # Forward
+        mu, logvar = self.encode(x)
+        z = self.reparameterize(mu, logvar)
+        x_recon = self.decode(z)
+        loss = self.loss(x_recon, x, mu, logvar)
+
+        # Backward
+        grad = 2 * (x_recon - x) / x.shape[0]
+        for layer in reversed(self.decoder_layers):
+            grad = layer.backward(grad, self.learning_rate)
+
+        # Gradients for latent
+        h = self.encoder_layers[-1].z
+        grad_mu = (mu / x.shape[0])
+        grad_logvar = 0.5 * (np.exp(logvar) - 1) / x.shape[0]
+
+        grad_latent = grad_mu + grad_logvar
+        self.W_mu.backward(grad_mu, self.learning_rate)
+        self.W_logvar.backward(grad_logvar, self.learning_rate)
+
+        for layer in reversed(self.encoder_layers):
+            grad_latent = layer.backward(grad_latent, self.learning_rate)
+
+        return loss
+
+    def train(self, X, epochs=10, batch_size=64):
+        for epoch in range(epochs):
+            perm = np.random.permutation(X.shape[0])
+            total_loss = 0
+            for i in range(0, X.shape[0], batch_size):
+                batch = X[perm[i:i+batch_size]]
+                total_loss += self.train_step(batch)
+            print(f"Epoch {epoch+1} Loss: {total_loss:.4f}")
+
+    def reconstruct(self, x):
+        mu, logvar = self.encode(x)
+        z = self.reparameterize(mu, logvar)
+        return self.decode(z)
+
+    def sample(self, n_samples=10):
+        z = np.random.randn(n_samples, self.W_mu.b.shape[0])
+        return self.decode(z)
+
+# ----- Visualize -----
+def plot_reconstructions(vae, X, n=10):
+    recon = vae.reconstruct(X[:n])
+    fig, axs = plt.subplots(2, n, figsize=(n, 2))
+    for i in range(n):
+        axs[0, i].imshow(X[i].reshape(28, 28), cmap='gray')
+        axs[0, i].axis('off')
+        axs[1, i].imshow(recon[i].reshape(28, 28), cmap='gray')
+        axs[1, i].axis('off')
+    axs[0, 0].set_title('Original')
+    axs[1, 0].set_title('Reconstructed')
+    plt.show()
+
+def plot_generated(vae, n=10):
+    samples = vae.sample(n)
+    fig, axs = plt.subplots(1, n, figsize=(n, 1.5))
+    for i in range(n):
+        axs[i].imshow(samples[i].reshape(28, 28), cmap='gray')
+        axs[i].axis('off')
+    plt.suptitle('Generated Samples')
+    plt.show()
+
+# ----- Run on MNIST -----
+if __name__ == "__main__":
+    X = load_mnist()[:10000]
+    vae = VAE(input_dim=784, hidden_dims=[128, 64], latent_dim=2, learning_rate=0.05)
+    vae.train(X, epochs=10, batch_size=128)
+    plot_reconstructions(vae, X)
+    plot_generated(vae)