Morten Hjorth-Jensen
142 lines
4.6 KiB
Python
142 lines
4.6 KiB
Python
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)
|