adding codes

vae does not work

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)