adding codes
vae does not work
This commit is contained in:
@@ -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)
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user