Create vae.py
This commit is contained in:
Morten Hjorth-Jensen
@@ -0,0 +1,173 @@
|
||||
import numpy as np
|
||||
|
||||
class VAE:
|
||||
def __init__(self, input_dim, hidden_dim, latent_dim, learning_rate=0.01):
|
||||
self.input_dim = input_dim
|
||||
self.hidden_dim = hidden_dim
|
||||
self.latent_dim = latent_dim
|
||||
self.learning_rate = learning_rate
|
||||
|
||||
# Encoder weights and biases
|
||||
self.W1 = np.random.randn(input_dim, hidden_dim) * 0.01
|
||||
self.b1 = np.zeros(hidden_dim)
|
||||
|
||||
self.W_mu = np.random.randn(hidden_dim, latent_dim) * 0.01
|
||||
self.b_mu = np.zeros(latent_dim)
|
||||
|
||||
self.W_logvar = np.random.randn(hidden_dim, latent_dim) * 0.01
|
||||
self.b_logvar = np.zeros(latent_dim)
|
||||
|
||||
# Decoder weights and biases
|
||||
self.W2 = np.random.randn(latent_dim, hidden_dim) * 0.01
|
||||
self.b2 = np.zeros(hidden_dim)
|
||||
|
||||
self.W_out = np.random.randn(hidden_dim, input_dim) * 0.01
|
||||
self.b_out = np.zeros(input_dim)
|
||||
|
||||
def sigmoid(self, x):
|
||||
return 1 / (1 + np.exp(-x))
|
||||
|
||||
def sigmoid_derivative(self, x):
|
||||
s = self.sigmoid(x)
|
||||
return s * (1 - s)
|
||||
|
||||
def binary_cross_entropy(self, recon_x, x):
|
||||
eps = 1e-8
|
||||
return -np.sum(x * np.log(recon_x + eps) + (1 - x) * np.log(1 - recon_x + eps))
|
||||
|
||||
def encode(self, x):
|
||||
h = self.sigmoid(np.dot(x, self.W1) + self.b1)
|
||||
mu = np.dot(h, self.W_mu) + self.b_mu
|
||||
logvar = np.dot(h, self.W_logvar) + self.b_logvar
|
||||
return mu, logvar, h
|
||||
|
||||
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_dec = self.sigmoid(np.dot(z, self.W2) + self.b2)
|
||||
x_recon = self.sigmoid(np.dot(h_dec, self.W_out) + self.b_out)
|
||||
return x_recon, h_dec
|
||||
|
||||
def compute_loss(self, x, x_recon, mu, logvar):
|
||||
bce = self.binary_cross_entropy(x_recon, x)
|
||||
kl = -0.5 * np.sum(1 + logvar - mu**2 - np.exp(logvar))
|
||||
return bce + kl
|
||||
|
||||
def train(self, data, epochs=100, batch_size=10):
|
||||
data = np.array(data)
|
||||
n_samples = data.shape[0]
|
||||
|
||||
for epoch in range(epochs):
|
||||
indices = np.random.permutation(n_samples)
|
||||
total_loss = 0
|
||||
|
||||
for i in range(0, n_samples, batch_size):
|
||||
batch = data[indices[i:i + batch_size]]
|
||||
grads = self._compute_gradients(batch)
|
||||
self._update_parameters(grads)
|
||||
total_loss += grads['loss']
|
||||
|
||||
avg_loss = total_loss / n_samples
|
||||
print(f"Epoch {epoch+1}/{epochs} - Loss: {avg_loss:.4f}")
|
||||
|
||||
def _compute_gradients(self, x):
|
||||
m = x.shape[0]
|
||||
|
||||
# Forward pass
|
||||
mu, logvar, h_enc = self.encode(x)
|
||||
z = self.reparameterize(mu, logvar)
|
||||
x_recon, h_dec = self.decode(z)
|
||||
|
||||
# Loss
|
||||
loss = self.compute_loss(x, x_recon, mu, logvar)
|
||||
|
||||
# Backpropagation (simplified SGD)
|
||||
# Output layer
|
||||
delta_out = x_recon - x # (m, input_dim)
|
||||
dW_out = np.dot(h_dec.T, delta_out) / m
|
||||
db_out = np.mean(delta_out, axis=0)
|
||||
|
||||
# Decoder hidden layer
|
||||
delta_dec = np.dot(delta_out, self.W_out.T) * self.sigmoid_derivative(np.dot(z, self.W2) + self.b2)
|
||||
dW2 = np.dot(z.T, delta_dec) / m
|
||||
db2 = np.mean(delta_dec, axis=0)
|
||||
|
||||
# Latent space gradients
|
||||
dz = np.dot(delta_dec, self.W2.T)
|
||||
|
||||
dmu = dz + mu / m
|
||||
dlogvar = 0.5 * dz * (np.exp(0.5 * logvar)) / m
|
||||
|
||||
# Encoder hidden layer
|
||||
dh = (np.dot(dmu, self.W_mu.T) + np.dot(dlogvar, self.W_logvar.T)) * self.sigmoid_derivative(np.dot(x, self.W1) + self.b1)
|
||||
|
||||
dW_mu = np.dot(h_enc.T, dmu) / m
|
||||
db_mu = np.mean(dmu, axis=0)
|
||||
|
||||
dW_logvar = np.dot(h_enc.T, dlogvar) / m
|
||||
db_logvar = np.mean(dlogvar, axis=0)
|
||||
|
||||
dW1 = np.dot(x.T, dh) / m
|
||||
db1 = np.mean(dh, axis=0)
|
||||
|
||||
return {
|
||||
'dW1': dW1, 'db1': db1,
|
||||
'dW_mu': dW_mu, 'db_mu': db_mu,
|
||||
'dW_logvar': dW_logvar, 'db_logvar': db_logvar,
|
||||
'dW2': dW2, 'db2': db2,
|
||||
'dW_out': dW_out, 'db_out': db_out,
|
||||
'loss': loss
|
||||
}
|
||||
|
||||
def _update_parameters(self, grads):
|
||||
self.W1 -= self.learning_rate * grads['dW1']
|
||||
self.b1 -= self.learning_rate * grads['db1']
|
||||
|
||||
self.W_mu -= self.learning_rate * grads['dW_mu']
|
||||
self.b_mu -= self.learning_rate * grads['db_mu']
|
||||
|
||||
self.W_logvar -= self.learning_rate * grads['dW_logvar']
|
||||
self.b_logvar -= self.learning_rate * grads['db_logvar']
|
||||
|
||||
self.W2 -= self.learning_rate * grads['dW2']
|
||||
self.b2 -= self.learning_rate * grads['db2']
|
||||
|
||||
self.W_out -= self.learning_rate * grads['dW_out']
|
||||
self.b_out -= self.learning_rate * grads['db_out']
|
||||
|
||||
def reconstruct(self, x):
|
||||
mu, logvar, _ = self.encode(x)
|
||||
z = self.reparameterize(mu, logvar)
|
||||
x_recon, _ = self.decode(z)
|
||||
return x_recon
|
||||
|
||||
def sample(self, n_samples=1):
|
||||
z = np.random.randn(n_samples, self.latent_dim)
|
||||
x_recon, _ = self.decode(z)
|
||||
return x_recon
|
||||
|
||||
|
||||
# Create synthetic binary data (patterns)
|
||||
np.random.seed(42)
|
||||
data = np.random.binomial(n=1, p=0.5, size=(100, 6))
|
||||
|
||||
# Initialize VAE: input=6, hidden=4, latent=2
|
||||
vae = VAE(input_dim=6, hidden_dim=4, latent_dim=2, learning_rate=0.1)
|
||||
|
||||
# Train
|
||||
vae.train(data, epochs=50, batch_size=10)
|
||||
|
||||
# Reconstruct
|
||||
x_test = data[0]
|
||||
x_recon = vae.reconstruct(x_test.reshape(1, -1))
|
||||
print("\nOriginal: ", x_test)
|
||||
print("Reconstructed:", np.round(x_recon[0], 2))
|
||||
|
||||
# Generate new samples
|
||||
samples = vae.sample(n_samples=3)
|
||||
print("\nGenerated samples:")
|
||||
print(np.round(samples, 2))
|
||||
|
||||
Reference in New Issue
Block a user