adding codes

vae does not work
This commit is contained in:
Morten Hjorth-Jensen
2025-05-30 10:24:27 +02:00
parent 7ff326f990
commit 95d2789104
2 changed files with 313 additions and 0 deletions
+141
View File
@@ -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)
+172
View File
@@ -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)