From 95d278910402760070042da7a52114de418536df Mon Sep 17 00:00:00 2001 From: Morten Hjorth-Jensen Date: Fri, 30 May 2025 10:24:27 +0200 Subject: [PATCH] adding codes vae does not work --- doc/Programs/DiffusionModels/vdiff.py | 141 +++++++++++++++++++++ doc/Programs/VAE/vae2.py | 172 ++++++++++++++++++++++++++ 2 files changed, 313 insertions(+) create mode 100644 doc/Programs/DiffusionModels/vdiff.py create mode 100644 doc/Programs/VAE/vae2.py diff --git a/doc/Programs/DiffusionModels/vdiff.py b/doc/Programs/DiffusionModels/vdiff.py new file mode 100644 index 000000000..64e6b0a0b --- /dev/null +++ b/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) diff --git a/doc/Programs/VAE/vae2.py b/doc/Programs/VAE/vae2.py new file mode 100644 index 000000000..2d955cabb --- /dev/null +++ b/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)