mnist diffusion

### 使用 Diffusion Model 生成 MNIST 数据集图像 为了使用扩散模型 (Diffusion Model, DM) 来生成 MNIST 数据集中的手写数字图像，可以遵循以下方法。此过程涉及两个主要阶段：前向传播（即逐渐增加噪声）以及反向传播（即去噪）。具体来说： #### 前向传播（加噪） 在这一过程中，原始图像被逐步加入高斯白噪音直到变得完全随机化。这一步骤定义了一系列由 $\beta_t$ 控制的方差参数下的马尔可夫链状态转移[^1]。 #### 后向传播（去噪） 该部分旨在逆转上述加噪流程，从纯噪声恢复到清晰图像。网络学习预测每一步中应减去多少噪声以接近初始输入。整个训练目标是最小化均方误差损失函数，在这里用来衡量真实分布 $q(\mathbf{x}_t|\mathbf{x}_{t-1})$ 和估计分布 $p_\theta (\mathbf{x}_{t-1} | \mathbf{x}_t)$之间的差异。 下面是基于PyTorch的一个简化版实现例子，用于演示如何构建并应用一个简单的扩散模型来进行MNIST数据集上的图像生成任务: ```python import torch from torch import nn import torch.nn.functional as F from torchvision.datasets import MNIST from torchvision.transforms import Compose, ToTensor, Lambda from torch.utils.data import DataLoader class SimpleUnet(nn.Module): def __init__(self): super().__init__() self.encoder = nn.Sequential( nn.Conv2d(1, 64, kernel_size=3), nn.ReLU(), nn.BatchNorm2d(64), nn.Conv2d(64, 128, kernel_size=3), nn.MaxPool2d(kernel_size=2), nn.ReLU(), nn.BatchNorm2d(128), nn.Conv2d(128, 256, kernel_size=3), nn.MaxPool2d(kernel_size=2), nn.ReLU() ) self.decoder = nn.Sequential( nn.Upsample(scale_factor=2), nn.ConvTranspose2d(256, 128, kernel_size=3), nn.ReLU(), nn.BatchNorm2d(128), nn.Upsample(scale_factor=2), nn.ConvTranspose2d(128, 64, kernel_size=3), nn.ReLU(), nn.BatchNorm2d(64), nn.ConvTranspose2d(64, 1, kernel_size=3), nn.Sigmoid() ) def forward(self, x): encoded_x = self.encoder(x) decoded_x = self.decoder(encoded_x) return decoded_x def linear_beta_schedule(timesteps, start=0.0001, end=0.02): return torch.linspace(start, end, timesteps) if __name__ == "__main__": device = "cuda" if torch.cuda.is_available() else "cpu" T = 300 betas = linear_beta_schedule(T).to(device) alphas = 1 - betas alpha_bars = torch.cumprod(alphas, axis=0) unet_model = SimpleUnet().to(device) optimizer = torch.optim.Adam(unet_model.parameters(), lr=1e-3) transform = Compose([ToTensor(), Lambda(lambda x: (x * 2) - 1)]) dataset = MNIST(root='./data', train=True, download=True, transform=transform) dataloader = DataLoader(dataset, batch_size=128, shuffle=True) epochs = 10 for epoch in range(epochs): for step, (images, labels) in enumerate(dataloader): images = images.to(device) # Sample a random timestep from the distribution. ts = torch.randint(low=0, high=T, size=(len(images),)).long().to(device) noise = torch.randn_like(images).to(device) noisy_images = ( sqrt_alpha_bar[ts][:, None, None, None].sqrt()*images + (1-alpha_bars[ts])[:,None,None,None].sqrt()*noise ).clamp(-1., 1.) predicted_noise = unet_model(noisy_images, ts) loss = F.mse_loss(predicted_noise, noise) optimizer.zero_grad() loss.backward() optimizer.step() if not step % 100: print(f'Epoch {epoch}, Step {step}: Loss={loss.item()}') ``` 这段代码展示了怎样创建一个简易版本的U-net架构作为基础模型，并利用线性的$\beta$调度策略来控制噪声水平变化。此外还包含了完整的训练循环逻辑，其中涉及到采样时间步$t\in[T]$、计算带噪图片$x^{(t)}=\sqrt{\bar{\alpha}}_tx+(1-\bar{\alpha}_t)\epsilon,\;\epsilon∼N(0,I)$以及最小化预测噪声与实际添加噪声间的MSE距离[^2].

以下是一个基于PyTorch深度学习框架，使用Diffusion模型处理Fashion-MNIST数据集，并添加数据可视化的完整Python代码： ```python import torch import torch.nn as nn import torch.optim as optim from torchvision import datasets, transforms import matplotlib.pyplot as plt import numpy as np # 设置设备 device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # 超参数 batch_size = 128 epochs = 20 learning_rate = 1e-3 image_size = 28 channels = 1 n_steps = 1000 beta_start = 1e-4 beta_end = 0.02 # 数据预处理 transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,)) ]) # 加载数据集 train_dataset = datasets.FashionMNIST(root='./data', train=True, transform=transform, download=True) train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True) # 定义扩散模型 class DiffusionModel(nn.Module): def __init__(self): super(DiffusionModel, self).__init__() self.model = nn.Sequential( nn.Conv2d(1, 64, kernel_size=3, stride=1, padding=1), nn.ReLU(), nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1), nn.ReLU(), nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1), nn.ReLU(), nn.Flatten(), nn.Linear(128 * 7 * 7, 1024), nn.ReLU(), nn.Linear(1024, 128 * 7 * 7), nn.ReLU(), nn.Unflatten(1, (128, 7, 7)), nn.ConvTranspose2d(128, 64, kernel_size=3, stride=2, padding=1, output_padding=1), nn.ReLU(), nn.ConvTranspose2d(64, 64, kernel_size=3, stride=2, padding=1, output_padding=1), nn.ReLU(), nn.Conv2d(64, 1, kernel_size=3, stride=1, padding=1), nn.Tanh() ) def forward(self, x): return self.model(x) # 初始化模型、损失函数和优化器 model = DiffusionModel().to(device) criterion = nn.MSELoss() optimizer = optim.Adam(model.parameters(), lr=learning_rate) # 训练循环 for epoch in range(epochs): for i, (images, _) in enumerate(train_loader): images = images.to(device) # 生成时间步 t = torch.randint(0, n_steps, (batch_size,), device=device).long() # 计算beta betas = torch.linspace(beta_start, beta_end, n_steps).to(device) alpha = 1 - betas alpha_hat = torch.cumprod(alpha, dim=0) # 获取噪声 noise = torch.randn_like(images).to(device) # 计算noisy images sqrt_alpha_hat = torch.sqrt(alpha_hat[t].reshape(-1, 1, 1, 1)) sqrt_one_minus_alpha_hat = torch.sqrt(1 - alpha_hat[t].reshape(-1, 1, 1, 1)) noisy_images = sqrt_alpha_hat * images + sqrt_one_minus_alpha_hat * noise # 前向传播 outputs = model(noisy_images) loss = criterion(outputs, noise) # 反向传播和优化 optimizer.zero_grad() loss.backward() optimizer.step() print(f"Epoch [{epoch+1}/{epochs}], Loss: {loss.item():.4f}") # 可视化结果 with torch.no_grad(): # 生成随机噪声 sample = torch.randn(64, channels, image_size, image_size).to(device) # 反向扩散过程 for i in range(n_steps): betas_t = beta_end * torch.ones(batch_size, device=device) alpha_t = 1 - betas_t alpha_hat_t = torch.cumprod(alpha_t, dim=0)[-1] sqrt_recip_alpha_t = torch.sqrt(1 / alpha_t) sqrt_one_minus_alpha_hat_t = torch.sqrt(1 - alpha_hat_t) noise = torch.randn_like(sample).to(device) # 模型预测 model_input = sample model_output = model(model_input) # 更新样本 sample = sqrt_recip_alpha_t * (model_input - (1 - alpha_t) / torch.sqrt(1 - alpha_hat_t) * model_output) + torch.sqrt(betas_t) * noise # 显示生成的图片 grid = utils.make_grid(sample.cpu(), nrow=8) plt.figure(figsize=(8,8)) plt.imshow(grid.permute(1, 2, 0).squeeze(), cmap='gray') plt.axis('off') plt.show() ``` 这段代码实现了一个基本的Diffusion模型，用于处理Fashion-MNIST数据集，并包含了数据可视化功能。