目录
13.1.1 VAE的ELBO推导与重参数化技巧（Reparameterization Trick），后验坍缩（Posterior Collapse）与Cyclic Annealing策略
13.1.2 层次VAE（HVAE）的马尔可夫先验与Skip-VAE的确定性路径，β-VAE的解耦表示学习与容量超参数
13.1.3 VQ-VAE的向量量化与码本（Codebook）优化，Gumbel-Softmax松弛与直通估计器（Straight-Through Estimator）
13.2.3 训练稳定性：谱归一化（Spectral Normalization）的幂迭代实现，DiffAugment与AdaAugment的数据增强策略
13.3.1 可逆神经网络（Invertible NN）的雅可比行列式计算与变量变换公式，RealNVP的耦合层（Coupling Layers）与棋盘掩码
13.3.2 Glow的ActNorm初始化、1x1可逆卷积（LU分解参数化）与多尺度架构，Flow++的变分去量化（Variational Dequantization）
13.4.2 连续时间扩散：随机微分方程（SDE）视角（VP/VE/Sub-VP）与概率流ODE，概率密度变换的Fokker-Planck方程
13.4.4 引导生成：Classifier Guidance的梯度缩放与无分类器引导（CFG）的隐式分类器，PNDM/DDIM的确定性采样与步数加速
第13章 深度生成模型
基于CMU 10-715 2025年秋季扩散模型专题与MIT流匹配课程，本章系统阐述深度生成建模的核心理论与实现技术，涵盖变分推断、对抗训练、可逆流与扩散概率模型等范式。
13.1 变分自编码器与潜空间模型
13.1.1 VAE的ELBO推导与重参数化技巧（Reparameterization Trick），后验坍缩（Posterior Collapse）与Cyclic Annealing策略
变分自编码器通过引入潜变量扩展自编码器框架的概率解释，利用变分推断近似难以计算的后验分布。证据下界将边际对数似然分解为重建期望与相对熵正则项，前者度量输入重构保真度，后者约束近似后验与标准高斯先验的统计距离。重参数化技巧通过从标准正态分布采样噪声并经由学习到的仿射变换转换为潜变量，实现随机梯度估计的方差缩减，使端到端反向传播训练成为可能。
后验坍缩现象指解码器能力过强时，编码器退化为与输入无关的先验分布，导致潜空间丧失语义结构。循环退火策略在训练过程中周期性调节正则项权重，初期强制编码器传递信息，中期允许模型优化重建质量，后期恢复完整目标函数，通过周期性重启避免优化陷入局部极小值，维持编码器的活跃信息传递能力。
Python
"""
Script: vae_elbo_reparameterization.py
Content: VAE implementation with ELBO optimization, reparameterization trick,
posterior collapse detection, and cyclic annealing scheduling.
Usage: python vae_elbo_reparameterization.py --dataset mnist --epochs 50 --anneal_cycles 4
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import numpy as np
import matplotlib.pyplot as plt
from typing import Tuple, Dict, List
import argparse
class ReparameterizationSampler(nn.Module):
"""
Reparameterization trick implementation for differentiable sampling.
z = mu + sigma * epsilon, where epsilon ~ N(0, I)
"""
def forward(self, mu: torch.Tensor, log_var: torch.Tensor) -> torch.Tensor:
std = torch.exp(0.5 * log_var)
eps = torch.randn_like(std)
return mu + std * eps
class Encoder(nn.Module):
"""
Probabilistic encoder mapping observations to latent distribution parameters.
Outputs mean and log-variance for diagonal Gaussian posterior.
"""
def __init__(self, input_dim: int = 784, hidden_dim: int = 400, latent_dim: int = 20):
super().__init__()
self.fc1 = nn.Linear(input_dim, hidden_dim)
self.fc2 = nn.Linear(hidden_dim, hidden_dim)
self.fc_mu = nn.Linear(hidden_dim, latent_dim)
self.fc_logvar = nn.Linear(hidden_dim, latent_dim)
self.sampler = ReparameterizationSampler()
# Initialize for stable training
nn.init.xavier_uniform_(self.fc1.weight)
nn.init.xavier_uniform_(self.fc2.weight)
nn.init.xavier_uniform_(self.fc_mu.weight)
nn.init.xavier_uniform_(self.fc_logvar.weight)
def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
h = F.relu(self.fc1(x))
h = F.relu(self.fc2(h))
mu = self.fc_mu(h)
log_var = self.fc_logvar(h)
z = self.sampler(mu, log_var)
return z, mu, log_var
class Decoder(nn.Module):
"""
Probabilistic decoder reconstructing observations from latent samples.
"""
def __init__(self, latent_dim: int = 20, hidden_dim: int = 400, output_dim: int = 784):
super().__init__()
self.fc1 = nn.Linear(latent_dim, hidden_dim)
self.fc2 = nn.Linear(hidden_dim, hidden_dim)
self.fc3 = nn.Linear(hidden_dim, output_dim)
nn.init.xavier_uniform_(self.fc1.weight)
nn.init.xavier_uniform_(self.fc2.weight)
nn.init.xavier_uniform_(self.fc3.weight)
def forward(self, z: torch.Tensor) -> torch.Tensor:
h = F.relu(self.fc1(z))
h = F.relu(self.fc2(h))
# Sigmoid output for binary cross-entropy reconstruction
return torch.sigmoid(self.fc3(h))
class VAE(nn.Module):
"""
Variational Autoencoder with cyclic annealing and posterior collapse monitoring.
"""
def __init__(self, input_dim: int = 784, hidden_dim: int = 400,
latent_dim: int = 20, anneal_cycles: int = 4,
anneal_ratio: float = 0.5):
super().__init__()
self.encoder = Encoder(input_dim, hidden_dim, latent_dim)
self.decoder = Decoder(latent_dim, hidden_dim, input_dim)
self.latent_dim = latent_dim
self.anneal_cycles = anneal_cycles
self.anneal_ratio = anneal_ratio
# Metrics tracking for posterior collapse detection
self.active_units_history = []
self.kl_divergence_history = []
self.reconstruction_history = []
def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
batch_size = x.size(0)
x_flat = x.view(batch_size, -1)
# Encode
z, mu, log_var = self.encoder(x_flat)
# Decode
x_recon = self.decoder(z)
# Compute ELBO components
recon_loss = F.binary_cross_entropy(x_recon, x_flat, reduction='sum') / batch_size
kl_div = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp()) / batch_size
return {
'x_recon': x_recon,
'mu': mu,
'log_var': log_var,
'z': z,
'reconstruction_loss': recon_loss,
'kl_divergence': kl_div
}
def compute_elbo(self, x: torch.Tensor, beta: float = 1.0) -> torch.Tensor:
"""
Compute Evidence Lower Bound with adjustable beta parameter.
ELBO = E[log p(x|z)] - beta * KL(q(z|x) || p(z))
"""
outputs = self.forward(x)
elbo = -(outputs['reconstruction_loss'] + beta * outputs['kl_divergence'])
return -elbo # Return negative ELBO as loss
def cyclic_annealing_step(self, epoch: int, total_epochs: int) -> float:
"""
Compute cyclic annealing weight following Fu et al. (2019).
Monotonic increase within each cycle, periodic reset.
"""
period = total_epochs // self.anneal_cycles
epoch_in_cycle = epoch % period
beta = min(1.0, (epoch_in_cycle / (period * self.anneal_ratio)))
return beta
def detect_posterior_collapse(self, mu: torch.Tensor, threshold: float = 0.01) -> int:
"""
Detect number of active latent units based on variance in mu across batch.
Low variance indicates collapsed dimension (independent of input).
"""
# Compute variance across batch for each latent dimension
mu_var = torch.var(mu, dim=0)
active_units = (mu_var > threshold).sum().item()
self.active_units_history.append(active_units)
return active_units
def train_vae_with_annealing(dataset: str = 'mnist', epochs: int = 50,
latent_dim: int = 20, anneal_cycles: int = 4,
batch_size: int = 128):
"""
Complete training pipeline with cyclic annealing and posterior collapse monitoring.
"""
# Data preparation
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Lambda(lambda x: x.view(-1)) # Flatten
])
if dataset == 'mnist':
train_data = datasets.MNIST('./data', train=True, download=True, transform=transform)
test_data = datasets.MNIST('./data', train=False, transform=transform)
else:
raise ValueError(f"Unsupported dataset: {dataset}")
train_loader = DataLoader(train_data, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_data, batch_size=batch_size, shuffle=False)
# Model initialization
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = VAE(input_dim=784, hidden_dim=400, latent_dim=latent_dim,
anneal_cycles=anneal_cycles).to(device)
optimizer = optim.Adam(model.parameters(), lr=1e-3)
# Training history
history = {
'train_loss': [],
'recon_loss': [],
'kl_loss': [],
'beta_values': [],
'active_units': []
}
for epoch in range(epochs):
model.train()
epoch_loss = 0
epoch_recon = 0
epoch_kl = 0
epoch_active = []
# Get cyclic beta for this epoch
beta = model.cyclic_annealing_step(epoch, epochs)
for batch_idx, (data, _) in enumerate(train_loader):
data = data.to(device)
optimizer.zero_grad()
# Forward pass
outputs = model(data)
loss = outputs['reconstruction_loss'] + beta * outputs['kl_divergence']
# Backward pass
loss.backward()
optimizer.step()
# Statistics
epoch_loss += loss.item()
epoch_recon += outputs['reconstruction_loss'].item()
epoch_kl += outputs['kl_divergence'].item()
# Monitor active units
active = model.detect_posterior_collapse(outputs['mu'])
epoch_active.append(active)
# Record epoch statistics
avg_loss = epoch_loss / len(train_loader)
avg_recon = epoch_recon / len(train_loader)
avg_kl = epoch_kl / len(train_loader)
avg_active = np.mean(epoch_active)
history['train_loss'].append(avg_loss)
history['recon_loss'].append(avg_recon)
history['kl_loss'].append(avg_kl)
history['beta_values'].append(beta)
history['active_units'].append(avg_active)
if epoch % 5 == 0:
print(f"Epoch {epoch}/{epochs} | Beta: {beta:.3f} | "
f"Loss: {avg_loss:.4f} | Recon: {avg_recon:.4f} | "
f"KL: {avg_kl:.4f} | Active Units: {avg_active:.0f}/{latent_dim}")
# Visualization
visualize_vae_results(model, test_loader, device, history, latent_dim)
return model, history
def visualize_vae_results(model: VAE, test_loader: DataLoader, device: torch.device,
history: Dict, latent_dim: int):
"""
Comprehensive visualization of VAE training dynamics and latent space structure.
"""
fig, axes = plt.subplots(2, 3, figsize=(18, 12))
# 1. Training loss decomposition
ax = axes[0, 0]
epochs = range(len(history['train_loss']))
ax.plot(epochs, history['recon_loss'], label='Reconstruction', color='blue')
ax.plot(epochs, history['kl_loss'], label='KL Divergence', color='red')
ax.set_xlabel('Epoch')
ax.set_ylabel('Loss Component')
ax.set_title('ELBO Decomposition Over Training')
ax.legend()
ax.grid(True, alpha=0.3)
# 2. Cyclic annealing schedule visualization
ax = axes[0, 1]
ax.plot(epochs, history['beta_values'], label='Beta Weight', color='green', linewidth=2)
ax.set_xlabel('Epoch')
ax.set_ylabel('Beta Value')
ax.set_title('Cyclic Annealing Schedule')
ax.grid(True, alpha=0.3)
ax.set_ylim(0, 1.1)
# 3. Posterior collapse monitoring
ax = axes[0, 2]
ax.plot(epochs, history['active_units'], label='Active Units', color='purple', marker='o')
ax.axhline(y=latent_dim, color='red', linestyle='--', label='Total Units')
ax.set_xlabel('Epoch')
ax.set_ylabel('Number of Active Units')
ax.set_title('Posterior Collapse Detection')
ax.legend()
ax.grid(True, alpha=0.3)
# 4. Latent space visualization (2D projection if dim > 2)
model.eval()
z_samples = []
labels_samples = []
with torch.no_grad():
for data, labels in test_loader:
data = data.to(device)
_, mu, _ = model.encoder(data)
z_samples.append(mu.cpu())
labels_samples.append(labels)
if len(z_samples) > 10: # Limit samples for speed
break
z_samples = torch.cat(z_samples, dim=0).numpy()
labels_samples = torch.cat(labels_samples, dim=0).numpy()
ax = axes[1, 0]
if latent_dim >= 2:
# Use first two dimensions or PCA if dim > 2
if latent_dim > 2:
from sklearn.decomposition import PCA
pca = PCA(n_components=2)
z_2d = pca.fit_transform(z_samples)
else:
z_2d = z_samples
scatter = ax.scatter(z_2d[:, 0], z_2d[:, 1], c=labels_samples,
cmap='tab10', alpha=0.6, s=20)
plt.colorbar(scatter, ax=ax)
ax.set_xlabel('z[0]' if latent_dim == 2 else 'PC1')
ax.set_ylabel('z[1]' if latent_dim == 2 else 'PC2')
ax.set_title('Latent Space Structure')
ax.grid(True, alpha=0.3)
# 5. Reconstruction quality
ax = axes[1, 1]
with torch.no_grad():
data, _ = next(iter(test_loader))
data = data[:8].to(device)
recon = model(data)['x_recon'].cpu().view(-1, 28, 28)
data = data.cpu().view(-1, 28, 28)
# Plot original and reconstructed
for i in range(8):
plt.subplot(2, 8, i + 1)
plt.imshow(data[i], cmap='gray')
plt.axis('off')
if i == 0:
plt.title('Original')
plt.subplot(2, 8, i + 9)
plt.imshow(recon[i], cmap='gray')
plt.axis('off')
if i == 0:
plt.title('Reconstructed')
# 6. Latent traversal
ax = axes[1, 2]
model.eval()
with torch.no_grad():
# Generate grid in latent space
n = 15
digit_size = 28
figure = np.zeros((digit_size * n, digit_size * n))
# Linearly spaced coordinates on the unit square
grid_x = np.linspace(-3, 3, n)
grid_y = np.linspace(-3, 3, n)
for i, yi in enumerate(grid_x):
for j, xi in enumerate(grid_y):
if latent_dim >= 2:
z_sample = torch.zeros(1, latent_dim).to(device)
z_sample[0, 0] = xi
z_sample[0, 1] = yi
x_decoded = model.decoder(z_sample).cpu().numpy().reshape(digit_size, digit_size)
figure[i * digit_size: (i + 1) * digit_size,
j * digit_size: (j + 1) * digit_size] = x_decoded
ax.imshow(figure, cmap='gray')
ax.set_title('Latent Space Traversal')
ax.axis('off')
plt.tight_layout()
plt.savefig('vae_training_analysis.png', dpi=300, bbox_inches='tight')
plt.show()
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--dataset', type=str, default='mnist')
parser.add_argument('--epochs', type=int, default=50)
parser.add_argument('--latent_dim', type=int, default=20)
parser.add_argument('--anneal_cycles', type=int, default=4)
parser.add_argument('--batch_size', type=int, default=128)
args = parser.parse_args()
train_vae_with_annealing(args.dataset, args.epochs, args.latent_dim,
args.anneal_cycles, args.batch_size)13.1.2 层次VAE（HVAE）的马尔可夫先验与Skip-VAE的确定性路径，β-VAE的解耦表示学习与容量超参数
层次变分自编码器扩展标准架构以支持多尺度潜变量，构建深度潜在变量模型捕捉数据的分层抽象特征。马尔可夫先验假设相邻层次间满足条件独立性，上层潜变量通过复杂先验分布约束下层变量，形成自上而下的生成过程与自下而上的推断路径。Skip-VAE引入跨层确定性连接，在概率路径之外保留信息流动的直接通道，缓解深层变分后验的梯度消失问题，增强高维潜空间中的信号传播效率。
β-变分自编码器通过调节证据下界中的正则项权重控制潜空间的解耦程度，超参数严格大于一时增强独立性约束，促使单个潜维度对应单一生成因子，学习可解释的数据变化轴。容量机制动态调整正则项的约束阈值，随着训练进展逐步增加允许的编码信息量，防止优化初期过度正则化导致的重建质量退化，在表示解耦与重建保真之间实现渐进平衡。
Python
"""
Script: hierarchical_vae_beta_vae.py
Content: Hierarchical VAE with Markovian prior, Skip-VAE with deterministic shortcuts,
and β-VAE with disentanglement metrics and capacity annealing.
Usage: python hierarchical_vae_beta_vae.py --model_type beta_vae --beta 4.0 --epochs 100
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import numpy as np
import matplotlib.pyplot as plt
from typing import List, Tuple, Dict
import argparse
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score
class HierarchicalEncoder(nn.Module):
"""
Hierarchical encoder with bottom-up inference path.
Produces latent variables at multiple scales.
"""
def __init__(self, input_dim: int = 784, hidden_dims: List[int] = [400, 200],
latent_dims: List[int] = [20, 10]):
super().__init__()
self.latent_dims = latent_dims
self.num_levels = len(latent_dims)
# Bottom-up layers
self.inference_layers = nn.ModuleList()
prev_dim = input_dim
for i, h_dim in enumerate(hidden_dims):
self.inference_layers.append(nn.Sequential(
nn.Linear(prev_dim, h_dim),
nn.ReLU()
))
prev_dim = h_dim
# Latent parameter layers (one per level)
self.mu_layers = nn.ModuleList()
self.logvar_layers = nn.ModuleList()
for i, lat_dim in enumerate(latent_dims):
self.mu_layers.append(nn.Linear(hidden_dims[-1], lat_dim))
self.logvar_layers.append(nn.Linear(hidden_dims[-1], lat_dim))
def forward(self, x: torch.Tensor) -> Tuple[List[torch.Tensor], List[torch.Tensor], List[torch.Tensor]]:
"""
Returns list of z, mu, log_var for each hierarchical level.
"""
h = x
for layer in self.inference_layers:
h = layer(h)
mus = []
logvars = []
zs = []
for mu_layer, logvar_layer in zip(self.mu_layers, self.logvar_layers):
mu = mu_layer(h)
logvar = logvar_layer(h)
std = torch.exp(0.5 * logvar)
eps = torch.randn_like(std)
z = mu + std * eps
mus.append(mu)
logvars.append(logvar)
zs.append(z)
return zs, mus, logvars
class HierarchicalDecoder(nn.Module):
"""
Hierarchical decoder with top-down generative path.
Incorporates skip connections for improved gradient flow.
"""
def __init__(self, latent_dims: List[int] = [20, 10],
hidden_dims: List[int] = [200, 400],
output_dim: int = 784, use_skip: bool = True):
super().__init__()
self.latent_dims = latent_dims
self.num_levels = len(latent_dims)
self.use_skip = use_skip
# Top-down generation with skip connections
self.generative_layers = nn.ModuleList()
self.skip_projections = nn.ModuleList() if use_skip else None
total_latent = sum(latent_dims)
prev_dim = latent_dims[-1] # Start from deepest level
for i, h_dim in enumerate(hidden_dims):
input_dim = prev_dim
if use_skip and i > 0:
# Add skip connection from appropriate latent level
skip_dim = latent_dims[-(i+1)] if i < len(latent_dims) else 0
input_dim += skip_dim
self.skip_projections.append(nn.Linear(skip_dim, prev_dim))
elif use_skip:
self.skip_projections.append(None)
self.generative_layers.append(nn.Sequential(
nn.Linear(input_dim, h_dim),
nn.ReLU()
))
prev_dim = h_dim
self.output_layer = nn.Linear(hidden_dims[-1], output_dim)
def forward(self, zs: List[torch.Tensor]) -> torch.Tensor:
"""
Decode from hierarchical latents with skip connections.
"""
# Start with deepest latent (last in list)
h = zs[-1]
for i, layer in enumerate(self.generative_layers):
if self.use_skip and i > 0 and i < len(zs):
# Add skip from corresponding latent level
skip_z = zs[-(i+1)]
skip_proj = self.skip_projections[i]
if skip_proj is not None:
h = h + skip_proj(skip_z)
h = layer(h)
return torch.sigmoid(self.output_layer(h))
class BetaVAE(nn.Module):
"""
β-VAE with adjustable disentanglement strength and capacity annealing.
"""
def __init__(self, input_dim: int = 784, hidden_dim: int = 400,
latent_dim: int = 20, beta: float = 4.0,
max_capacity: float = 25.0, capacity_iters: int = 100000):
super().__init__()
self.latent_dim = latent_dim
self.beta = beta
self.max_capacity = max_capacity
self.capacity_iters = capacity_iters
self.iteration = 0
# Standard encoder/decoder
self.encoder_fc1 = nn.Linear(input_dim, hidden_dim)
self.encoder_fc2 = nn.Linear(hidden_dim, hidden_dim)
self.fc_mu = nn.Linear(hidden_dim, latent_dim)
self.fc_logvar = nn.Linear(hidden_dim, latent_dim)
self.decoder_fc1 = nn.Linear(latent_dim, hidden_dim)
self.decoder_fc2 = nn.Linear(hidden_dim, hidden_dim)
self.decoder_out = nn.Linear(hidden_dim, input_dim)
def encode(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
h = F.relu(self.encoder_fc1(x))
h = F.relu(self.encoder_fc2(h))
mu = self.fc_mu(h)
logvar = self.fc_logvar(h)
std = torch.exp(0.5 * logvar)
eps = torch.randn_like(std)
z = mu + std * eps
return z, mu, logvar
def decode(self, z: torch.Tensor) -> torch.Tensor:
h = F.relu(self.decoder_fc1(z))
h = F.relu(self.decoder_fc2(h))
return torch.sigmoid(self.decoder_out(h))
def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
z, mu, logvar = self.encode(x)
x_recon = self.decode(z)
recon_loss = F.binary_cross_entropy(x_recon, x, reduction='sum') / x.size(0)
kl_loss = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp()) / x.size(0)
# Capacity annealing: gradually increase allowed information
capacity = min(self.max_capacity, self.max_capacity * self.iteration / self.capacity_iters)
self.iteration += 1
# β-VAE loss with capacity constraint
# Loss = Recon + beta * |KL - C|
constrained_kl = torch.abs(kl_loss - capacity)
total_loss = recon_loss + self.beta * constrained_kl
return {
'loss': total_loss,
'recon_loss': recon_loss,
'kl_loss': kl_loss,
'capacity': capacity,
'z': z,
'mu': mu,
'x_recon': x_recon
}
class DisentanglementMetrics:
"""
Metrics for evaluating disentanglement quality.
"""
def __init__(self):
self.metrics_history = {
'mutual_information_gap': [],
'separated_attribute_predictability': [],
'modularity': []
}
def compute_mig(self, latent_codes: np.ndarray, factors: np.ndarray) -> float:
"""
Mutual Information Gap (MIG) metric.
Measures gap between top two latent dimensions in terms of mutual information
with ground truth factors.
"""
from sklearn.feature_selection import mutual_info_regression
num_factors = factors.shape[1]
num_latents = latent_codes.shape[1]
mig_scores = []
for factor_idx in range(num_factors):
factor = factors[:, factor_idx]
# Compute MI between each latent dimension and this factor
mi_scores = []
for lat_idx in range(num_latents):
mi = mutual_info_regression(latent_codes[:, lat_idx:lat_idx+1], factor)[0]
mi_scores.append(mi)
mi_scores = sorted(mi_scores, reverse=True)
if len(mi_scores) >= 2:
mig = mi_scores[0] - mi_scores[1]
mig_scores.append(mig)
return np.mean(mig_scores) if mig_scores else 0.0
def compute_sap(self, latent_codes: np.ndarray, factors: np.ndarray) -> float:
"""
Separated Attribute Predictability (SAP) score.
Measures linear separability of factors using latent codes.
"""
num_factors = factors.shape[1]
sap_scores = []
for factor_idx in range(num_factors):
factor = factors[:, factor_idx]
# Train classifier to predict factor from latents
clf = LogisticRegression(max_iter=1000)
clf.fit(latent_codes, factor)
pred = clf.predict(latent_codes)
acc = accuracy_score(factor, pred)
sap_scores.append(acc)
return np.mean(sap_scores)
def train_beta_vae(dataset: str = 'mnist', beta: float = 4.0,
epochs: int = 100, latent_dim: int = 20):
"""
Training pipeline for β-VAE with disentanglement evaluation.
"""
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Lambda(lambda x: x.view(-1))
])
train_data = datasets.MNIST('./data', train=True, download=True, transform=transform)
train_loader = DataLoader(train_data, batch_size=128, shuffle=True)
model = BetaVAE(input_dim=784, hidden_dim=400, latent_dim=latent_dim,
beta=beta, max_capacity=25.0).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
history = {'loss': [], 'recon': [], 'kl': [], 'capacity': []}
metrics = DisentanglementMetrics()
for epoch in range(epochs):
model.train()
epoch_loss = 0
for batch_idx, (data, _) in enumerate(train_loader):
data = data.to(device)
optimizer.zero_grad()
outputs = model(data)
loss = outputs['loss']
loss.backward()
optimizer.step()
epoch_loss += loss.item()
avg_loss = epoch_loss / len(train_loader)
history['loss'].append(avg_loss)
history['recon'].append(outputs['recon_loss'].item())
history['kl'].append(outputs['kl_loss'].item())
history['capacity'].append(outputs['capacity'])
if epoch % 10 == 0:
print(f"Epoch {epoch}/{epochs} | Loss: {avg_loss:.4f} | "
f"Recon: {outputs['recon_loss']:.4f} | KL: {outputs['kl_loss']:.4f} | "
f"Capacity: {outputs['capacity']:.2f}")
# Visualization
visualize_beta_vae(model, train_loader, device, history, latent_dim, beta)
return model, history
def visualize_beta_vae(model: BetaVAE, data_loader: DataLoader, device: torch.device,
history: Dict, latent_dim: int, beta: float):
"""
Comprehensive visualization for β-VAE including disentanglement analysis.
"""
fig, axes = plt.subplots(2, 3, figsize=(18, 12))
# 1. Training curves
ax = axes[0, 0]
epochs = range(len(history['loss']))
ax.plot(epochs, history['recon'], label='Reconstruction', color='blue')
ax.plot(epochs, history['kl'], label='KL Divergence', color='red')
ax.plot(epochs, history['capacity'], label='Capacity', color='green', linestyle='--')
ax.set_xlabel('Epoch')
ax.set_ylabel('Loss / Capacity')
ax.set_title(f'β-VAE Training (β={beta})')
ax.legend()
ax.grid(True, alpha=0.3)
# 2. Latent traversal (showing disentanglement)
model.eval()
with torch.no_grad():
n = 10
digit_size = 28
figure = np.zeros((digit_size * latent_dim, digit_size * n))
# Sample base latent
z_base = torch.zeros(1, latent_dim).to(device)
for i in range(latent_dim):
for j, val in enumerate(np.linspace(-3, 3, n)):
z = z_base.clone()
z[0, i] = val
recon = model.decode(z).cpu().numpy().reshape(digit_size, digit_size)
figure[i * digit_size: (i + 1) * digit_size,
j * digit_size: (j + 1) * digit_size] = recon
ax = axes[0, 1]
ax.imshow(figure, cmap='gray')
ax.set_title('Latent Traversal (Disentanglement Check)')
ax.axis('off')
# 3. KL per dimension (showing which dimensions are active)
model.eval()
kl_per_dim = []
with torch.no_grad():
for data, _ in data_loader:
data = data.to(device)
_, mu, logvar = model.encode(data)
# KL per dimension: -0.5 * sum(1 + logvar - mu^2 - exp(logvar))
kl_dims = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp())
kl_per_dim.append(kl_dims.mean(dim=0).cpu().numpy())
if len(kl_per_dim) > 5:
break
kl_per_dim = np.mean(kl_per_dim, axis=0)
ax = axes[0, 2]
ax.bar(range(latent_dim), kl_per_dim, color='purple', alpha=0.7)
ax.set_xlabel('Latent Dimension')
ax.set_ylabel('Average KL Divergence')
ax.set_title('Per-Dimension KL (Information Usage)')
ax.grid(True, alpha=0.3, axis='y')
# 4. Reconstruction quality comparison
ax = axes[1, 0]
with torch.no_grad():
data, _ = next(iter(data_loader))
data = data[:8].to(device)
recon = model(data)['x_recon'].cpu().view(-1, 28, 28)
data = data.cpu().view(-1, 28, 28)
for i in range(8):
plt.subplot(2, 8, i + 1)
plt.imshow(data[i], cmap='gray')
plt.axis('off')
plt.subplot(2, 8, i + 9)
plt.imshow(recon[i], cmap='gray')
plt.axis('off')
# 5. Interpolation in latent space
ax = axes[1, 1]
with torch.no_grad():
# Get two random samples
z1, _, _ = model.encode(data[0:1].to(device).view(1, -1))
z2, _, _ = model.encode(data[1:2].to(device).view(1, -1))
interpolations = []
for alpha in np.linspace(0, 1, 8):
z_interp = (1 - alpha) * z1 + alpha * z2
recon = model.decode(z_interp).cpu().view(28, 28)
interpolations.append(recon)
# Plot interpolation
interp_grid = np.concatenate(interpolations, axis=1)
ax.imshow(interp_grid, cmap='gray')
ax.set_title('Linear Interpolation in Latent Space')
ax.axis('off')
# 6. Information bottleneck visualization
ax = axes[1, 2]
# Plot relationship between KL and reconstruction error
kls = np.array(history['kl'])
recons = np.array(history['recon'])
ax.scatter(kls, recons, c=range(len(kls)), cmap='viridis', alpha=0.6)
ax.set_xlabel('KL Divergence (Information)')
ax.set_ylabel('Reconstruction Error')
ax.set_title('Information vs Reconstruction Trade-off')
cbar = plt.colorbar(ax.collections[0], ax=ax)
cbar.set_label('Epoch')
ax.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig(f'beta_vae_analysis_beta{beta}.png', dpi=300)
plt.show()
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--model_type', type=str, default='beta_vae',
choices=['hvae', 'skip_vae', 'beta_vae'])
parser.add_argument('--beta', type=float, default=4.0)
parser.add_argument('--epochs', type=int, default=100)
parser.add_argument('--latent_dim', type=int, default=20)
args = parser.parse_args()
train_beta_vae(beta=args.beta, epochs=args.epochs, latent_dim=args.latent_dim)13.1.3 VQ-VAE的向量量化与码本（Codebook）优化，Gumbel-Softmax松弛与直通估计器（Straight-Through Estimator）
向量量化变分自编码器通过离散潜表示桥接自回归模型与连续潜空间，码本学习将编码器输出映射到有限嵌入向量的最近邻元素。离散瓶颈迫使模型捕获高层语义结构而非细粒度噪声，通过指数移动平均更新码本向量或承诺损失优化编码器与码本的相对位置，缓解梯度在不可微量化算子处的截断问题。
Gumbel-Softmax重参数化在连续近似与离散采样间建立可微桥梁，通过Gumbel分布采样与温度退火策略生成类别分布的软近似。直通估计器在前向传播采用离散argmax决策，反向传播直接将梯度复制到编码器输出，实现离散选择的梯度下降优化。两种技术分别适用于需要概率 relax 或确定性离散化的场景，支撑高分辨率图像与音频波形的离散潜空间建模。
Python
"""
Script: vqvae_gumbel_softmax.py
Content: VQ-VAE with vector quantization, codebook optimization via EMA and commitment loss,
and Gumbel-Softmax VAE with straight-through estimator.
Usage: python vqvae_gumbel_softmax.py --model_type vqvae --num_embeddings 512 --embedding_dim 64
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import numpy as np
import matplotlib.pyplot as plt
from typing import Tuple, Dict, Optional
import argparse
class VectorQuantizer(nn.Module):
"""
Vector Quantization layer with EMA codebook update and commitment loss.
Following van den Oord et al. (2017) and Razavi et al. (2019).
"""
def __init__(self, num_embeddings: int = 512, embedding_dim: int = 64,
commitment_cost: float = 0.25, decay: float = 0.99, epsilon: float = 1e-5):
super().__init__()
self.num_embeddings = num_embeddings
self.embedding_dim = embedding_dim
self.commitment_cost = commitment_cost
self.decay = decay
self.epsilon = epsilon
# Codebook embeddings
self.embedding = nn.Embedding(num_embeddings, embedding_dim)
self.embedding.weight.data.uniform_(-1/num_embeddings, 1/num_embeddings)
# EMA parameters
self.register_buffer('ema_cluster_size', torch.zeros(num_embeddings))
self.register_buffer('ema_w', self.embedding.weight.data.clone())
def forward(self, inputs: torch.Tensor) -> Tuple[torch.Tensor, Dict]:
"""
Quantize continuous inputs to nearest codebook vectors.
"""
# Convert inputs from BCHW -> BHWC
input_shape = inputs.shape
flat_input = inputs.view(-1, self.embedding_dim)
# Calculate distances to codebook vectors
distances = (torch.sum(flat_input**2, dim=1, keepdim=True)
+ torch.sum(self.embedding.weight**2, dim=1)
- 2 * torch.matmul(flat_input, self.embedding.weight.t()))
# Encoding
encoding_indices = torch.argmin(distances, dim=1)
encodings = F.one_hot(encoding_indices, self.num_embeddings).float()
quantized = self.embedding(encoding_indices).view(input_shape)
# Loss
e_latent_loss = F.mse_loss(quantized.detach(), inputs)
q_latent_loss = F.mse_loss(quantized, inputs.detach())
loss = q_latent_loss + self.commitment_cost * e_latent_loss
# Straight-through estimator
quantized = inputs + (quantized - inputs).detach()
# EMA update during training
if self.training:
self.ema_cluster_size = self.ema_cluster_size * self.decay + \
(1 - self.decay) * torch.sum(encodings, dim=0)
# Laplace smoothing
n = torch.sum(self.ema_cluster_size.data)
self.ema_cluster_size = (
(self.ema_cluster_size + self.epsilon) / (n + self.num_embeddings * self.epsilon) * n
)
dw = torch.matmul(encodings.t(), flat_input)
self.ema_w = nn.Parameter(self.ema_w * self.decay + (1 - self.decay) * dw)
self.embedding.weight = nn.Parameter(self.ema_w / self.ema_cluster_size.unsqueeze(1))
# Perplexity (metric for codebook usage)
avg_probs = torch.mean(encodings, dim=0)
perplexity = torch.exp(-torch.sum(avg_probs * torch.log(avg_probs + 1e-10)))
return quantized, {
'quantizer_loss': loss,
'encoding_indices': encoding_indices.view(input_shape[0], input_shape[2], input_shape[3]),
'perplexity': perplexity,
'encodings': encodings
}
class VQVAE(nn.Module):
"""
Vector Quantized VAE with encoder-decoder architecture for discrete representation learning.
"""
def __init__(self, in_channels: int = 1, embedding_dim: int = 64,
num_embeddings: int = 512, hidden_dim: int = 128,
num_res_blocks: int = 2):
super().__init__()
# Encoder: maps image to latent space
self.encoder = nn.Sequential(
nn.Conv2d(in_channels, hidden_dim, 4, 2, 1), # 28 -> 14
nn.ReLU(),
nn.Conv2d(hidden_dim, hidden_dim, 4, 2, 1), # 14 -> 7
nn.ReLU(),
nn.Conv2d(hidden_dim, embedding_dim, 3, 1, 1), # 7 -> 7
)
# Pre-quantization convolution
self.pre_quant_conv = nn.Conv2d(embedding_dim, embedding_dim, 1)
# Vector quantization
self.quantizer = VectorQuantizer(num_embeddings, embedding_dim)
# Decoder: reconstructs from quantized latents
self.decoder = nn.Sequential(
nn.ConvTranspose2d(embedding_dim, hidden_dim, 4, 2, 1), # 7 -> 14
nn.ReLU(),
nn.ConvTranspose2d(hidden_dim, hidden_dim, 4, 2, 1), # 14 -> 28
nn.ReLU(),
nn.Conv2d(hidden_dim, in_channels, 3, 1, 1),
nn.Sigmoid()
)
def forward(self, x: torch.Tensor) -> Dict:
# Encode
z = self.encoder(x)
z = self.pre_quant_conv(z)
# Quantize
quantized, quant_info = self.quantizer(z)
# Decode
x_recon = self.decoder(quantized)
# Reconstruction loss
recon_loss = F.mse_loss(x_recon, x)
total_loss = recon_loss + quant_info['quantizer_loss']
return {
'loss': total_loss,
'recon_loss': recon_loss,
'quantizer_loss': quant_info['quantizer_loss'],
'x_recon': x_recon,
'quantized': quantized,
'perplexity': quant_info['perplexity'],
'encoding_indices': quant_info['encoding_indices']
}
def encode(self, x: torch.Tensor) -> torch.Tensor:
"""Get discrete codes for input."""
z = self.encoder(x)
z = self.pre_quant_conv(z)
_, info = self.quantizer(z)
return info['encoding_indices']
def decode(self, encoding_indices: torch.Tensor) -> torch.Tensor:
"""Reconstruct from discrete codes."""
quantized = self.quantizer.embedding(encoding_indices)
# Reshape if needed (assuming square spatial dimensions)
# This is simplified - actual implementation needs proper reshaping
return self.decoder(quantized.permute(0, 3, 1, 2))
class GumbelSoftmaxVAE(nn.Module):
"""
VAE with discrete latent using Gumbel-Softmax reparameterization.
Alternative to VQ-VAE for categorical latent variables.
"""
def __init__(self, input_dim: int = 784, hidden_dim: int = 400,
num_categories: int = 10, num_samples: int = 30, # K categories, N samples
temperature: float = 0.5, hard: bool = True):
super().__init__()
self.num_categories = num_categories
self.num_samples = num_samples
self.temperature = temperature
self.hard = hard
# Encoder
self.encoder = nn.Sequential(
nn.Linear(input_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, num_categories * num_samples) # Logits for Gumbel-Softmax
)
# Decoder
self.decoder = nn.Sequential(
nn.Linear(num_categories * num_samples, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, input_dim),
nn.Sigmoid()
)
def sample_gumbel_softmax(self, logits: torch.Tensor, temperature: float) -> torch.Tensor:
"""
Gumbel-Softmax sampling with straight-through estimator for hard samples.
"""
# Gumbel noise
u = torch.rand_like(logits)
gumbel_noise = -torch.log(-torch.log(u + 1e-10) + 1e-10)
# Add noise and apply temperature
y = logits + gumbel_noise
y_soft = F.softmax(y / temperature, dim=-1)
if self.hard:
# Straight-through estimator
# Forward: argmax (hard one-hot)
# Backward: soft gradients
index = y_soft.max(dim=-1, keepdim=True)[1]
y_hard = torch.zeros_like(y_soft).scatter_(-1, index, 1.0)
y = y_hard - y_soft.detach() + y_soft
else:
y = y_soft
return y
def forward(self, x: torch.Tensor) -> Dict:
batch_size = x.size(0)
# Get logits
logits = self.encoder(x)
logits = logits.view(batch_size, self.num_samples, self.num_categories)
# Sample using Gumbel-Softmax
samples = self.sample_gumbel_softmax(logits, self.temperature)
# Flatten for decoder
z = samples.view(batch_size, -1)
x_recon = self.decoder(z)
# Loss
recon_loss = F.binary_cross_entropy(x_recon, x, reduction='sum') / batch_size
# KL divergence for categorical (uniform prior)
q = F.softmax(logits, dim=-1)
log_q = F.log_softmax(logits, dim=-1)
kl_loss = torch.sum(q * (log_q - torch.log(torch.tensor(1.0 / self.num_categories)))) / batch_size
total_loss = recon_loss + kl_loss
return {
'loss': total_loss,
'recon_loss': recon_loss,
'kl_loss': kl_loss,
'x_recon': x_recon,
'samples': samples,
'logits': logits
}
def anneal_temperature(self, epoch: int, total_epochs: int, min_temp: float = 0.1):
"""Anneal temperature for sharper discrete samples."""
self.temperature = max(min_temp, 0.5 * np.exp(-epoch / total_epochs * 3))
def train_vqvae(dataset: str = 'mnist', epochs: int = 50,
num_embeddings: int = 512, embedding_dim: int = 64):
"""
Training pipeline for VQ-VAE with codebook usage monitoring.
"""
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
transform = transforms.Compose([
transforms.ToTensor(),
])
train_data = datasets.MNIST('./data', train=True, download=True, transform=transform)
train_loader = DataLoader(train_data, batch_size=128, shuffle=True)
model = VQVAE(in_channels=1, embedding_dim=embedding_dim,
num_embeddings=num_embeddings).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=2e-4)
history = {
'loss': [],
'recon': [],
'vq_loss': [],
'perplexity': [],
'codebook_usage': []
}
for epoch in range(epochs):
model.train()
epoch_stats = {'loss': 0, 'recon': 0, 'vq': 0, 'perplexity': 0}
for batch_idx, (data, _) in enumerate(train_loader):
data = data.to(device)
optimizer.zero_grad()
outputs = model(data)
loss = outputs['loss']
loss.backward()
optimizer.step()
epoch_stats['loss'] += loss.item()
epoch_stats['recon'] += outputs['recon_loss'].item()
epoch_stats['vq'] += outputs['quantizer_loss'].item()
epoch_stats['perplexity'] += outputs['perplexity'].item()
for k in epoch_stats:
epoch_stats[k] /= len(train_loader)
history[k].append(epoch_stats[k])
# Calculate codebook usage percentage
with torch.no_grad():
sample_data = next(iter(train_loader))[0].to(device)
codes = model.encode(sample_data)
unique_codes = torch.unique(codes).numel()
usage_pct = unique_codes / num_embeddings * 100
history['codebook_usage'].append(usage_pct)
if epoch % 5 == 0:
print(f"Epoch {epoch}/{epochs} | Loss: {epoch_stats['loss']:.4f} | "
f"Recon: {epoch_stats['recon']:.4f} | VQ: {epoch_stats['vq']:.4f} | "
f"Perplexity: {epoch_stats['perplexity']:.2f} | "
f"Codebook Usage: {usage_pct:.1f}%")
visualize_vqvae(model, train_loader, device, history, num_embeddings)
return model, history
def visualize_vqvae(model: VQVAE, data_loader: DataLoader, device: torch.device,
history: Dict, num_embeddings: int):
"""
Visualization for VQ-VAE including codebook analysis and reconstructions.
"""
fig, axes = plt.subplots(2, 3, figsize=(18, 12))
# 1. Training curves
ax = axes[0, 0]
epochs = range(len(history['loss']))
ax.plot(epochs, history['recon'], label='Reconstruction', color='blue')
ax.plot(epochs, history['vq'], label='Vector Quantization', color='red')
ax.set_xlabel('Epoch')
ax.set_ylabel('Loss')
ax.set_title('VQ-VAE Training Loss')
ax.legend()
ax.grid(True, alpha=0.3)
# 2. Codebook perplexity (usage metric)
ax = axes[0, 1]
ax.plot(epochs, history['perplexity'], color='purple', marker='o')
ax.set_xlabel('Epoch')
ax.set_ylabel('Perplexity')
ax.set_title(f'Codebook Perplexity (max={num_embeddings})')
ax.grid(True, alpha=0.3)
# 3. Codebook usage percentage
ax = axes[0, 2]
ax.plot(epochs, history['codebook_usage'], color='green', marker='s')
ax.axhline(y=100, color='red', linestyle='--', alpha=0.5)
ax.set_xlabel('Epoch')
ax.set_ylabel('Usage %')
ax.set_title('Codebook Utilization')
ax.set_ylim(0, 100)
ax.grid(True, alpha=0.3)
# 4. Reconstruction comparison
ax = axes[1, 0]
model.eval()
with torch.no_grad():
data, _ = next(iter(data_loader))
data = data[:8].to(device)
recon = model(data)['x_recon'].cpu()
# Show original and reconstructed
grid = torch.cat([data.cpu(), recon], dim=0)
grid = grid.view(-1, 1, 28, 28)
from torchvision.utils import make_grid
grid_img = make_grid(grid, nrow=8, padding=2, normalize=True)
ax.imshow(grid_img.permute(1, 2, 0).numpy())
ax.set_title('Original (top) vs Reconstruction (bottom)')
ax.axis('off')
# 5. Embedding space visualization (PCA of codebook)
ax = axes[1, 1]
embeddings = model.quantizer.embedding.weight.data.cpu().numpy()
from sklearn.decomposition import PCA
pca = PCA(n_components=2)
emb_2d = pca.fit_transform(embeddings)
scatter = ax.scatter(emb_2d[:, 0], emb_2d[:, 1], alpha=0.6, s=20, c=range(len(emb_2d)), cmap='tab20')
ax.set_title('Codebook Embeddings (PCA)')
ax.set_xlabel('PC1')
ax.set_ylabel('PC2')
plt.colorbar(scatter, ax=ax, label='Embedding Index')
# 6. Code frequency histogram
ax = axes[1, 2]
model.eval()
code_counts = torch.zeros(num_embeddings)
with torch.no_grad():
for i, (data, _) in enumerate(data_loader):
if i > 10: # Sample subset
break
data = data.to(device)
codes = model.encode(data)
for c in range(num_embeddings):
code_counts[c] += (codes == c).sum().item()
top_codes = torch.argsort(code_counts, descending=True)[:50]
ax.bar(range(50), code_counts[top_codes].numpy(), color='orange', alpha=0.7)
ax.set_xlabel('Top 50 Codes (sorted by frequency)')
ax.set_ylabel('Usage Count')
ax.set_title('Code Frequency Distribution')
ax.grid(True, alpha=0.3, axis='y')
plt.tight_layout()
plt.savefig('vqvae_analysis.png', dpi=300)
plt.show()
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--model_type', type=str, default='vqvae',
choices=['vqvae', 'gumbel'])
parser.add_argument('--num_embeddings', type=int, default=512)
parser.add_argument('--embedding_dim', type=int, default=64)
parser.add_argument('--epochs', type=int, default=50)
args = parser.parse_args()
train_vqvae(num_embeddings=args.num_embeddings,
embedding_dim=args.embedding_dim,
epochs=args.epochs)13.2 生成对抗网络
13.2.1 GAN的Minimax优化与JS散度缺陷，Wasserstein GAN的EM距离与Lipschitz约束（Weight Clipping/RGP/Spectral Normalization）
生成对抗网络通过生成器与判别器的双人博弈学习数据分布，生成器合成样本欺骗判别器，判别器区分真实与伪造数据，最小最大优化目标推动生成分布逼近真实分布。原始框架的散度最小化存在梯度消失缺陷，当分布支撑集不交叠时，Jensen-Shannon散度饱和导致生成器梯度稀疏，无法有效更新网络参数。
Wasserstein生成对抗网络引入地球移动距离替代散度度量，通过最优传输理论衡量分布间的最小质量转移成本，在连续空间上提供处处定义的梯度。Lipschitz约束确保判别器函数类满足平滑性要求，权重裁剪将参数限制在紧集内实现弱Lipschitz约束，梯度惩罚直接在目标函数添加正则项强制梯度范数逼近单位值，谱归一化通过幂迭代控制每一层的最大奇异值，从架构层面保证全局Lipschitz连续性。
Python
"""
Script: gan_wasserstein_improvements.py
Content: Standard GAN with Minimax objective, WGAN with weight clipping,
WGAN-GP with gradient penalty, and Spectral Normalization.
Usage: python gan_wasserstein_improvements.py --method wgan_gp --dataset mnist --epochs 50
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import numpy as np
import matplotlib.pyplot as plt
from typing import Tuple, Dict, List
import argparse
class MinimaxGAN:
"""
Standard GAN with Minimax optimization and binary cross-entropy.
Demonstrates JS divergence limitations (mode collapse, vanishing gradients).
"""
def __init__(self, latent_dim: int = 100, hidden_dim: int = 128,
image_dim: int = 784):
self.latent_dim = latent_dim
# Generator
self.G = nn.Sequential(
nn.Linear(latent_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, image_dim),
nn.Sigmoid()
)
# Discriminator (outputs probability)
self.D = nn.Sequential(
nn.Linear(image_dim, hidden_dim),
nn.LeakyReLU(0.2),
nn.Linear(hidden_dim, hidden_dim),
nn.LeakyReLU(0.2),
nn.Linear(hidden_dim, 1),
nn.Sigmoid()
)
def train_step(self, real_data: torch.Tensor,
g_optimizer: optim.Optimizer,
d_optimizer: optim.Optimizer) -> Dict[str, float]:
batch_size = real_data.size(0)
device = real_data.device
# Labels
real_labels = torch.ones(batch_size, 1).to(device)
fake_labels = torch.zeros(batch_size, 1).to(device)
# Train Discriminator
d_optimizer.zero_grad()
# Real data
real_output = self.D(real_data)
d_loss_real = F.binary_cross_entropy(real_output, real_labels)
# Fake data
z = torch.randn(batch_size, self.latent_dim).to(device)
fake_data = self.G(z)
fake_output = self.D(fake_data.detach())
d_loss_fake = F.binary_cross_entropy(fake_output, fake_labels)
d_loss = d_loss_real + d_loss_fake
d_loss.backward()
d_optimizer.step()
# Train Generator (maximize log(D(G(z))))
g_optimizer.zero_grad()
fake_output = self.D(fake_data)
g_loss = F.binary_cross_entropy(fake_output, real_labels) # Fool discriminator
g_loss.backward()
g_optimizer.step()
return {
'd_loss': d_loss.item(),
'g_loss': g_loss.item(),
'd_real_acc': real_output.mean().item(),
'd_fake_acc': (1 - fake_output).mean().item()
}
class WGAN:
"""
Wasserstein GAN with weight clipping for Lipschitz constraint.
Uses Wasserstein distance (Earth Mover's distance) instead of JS divergence.
"""
def __init__(self, latent_dim: int = 100, hidden_dim: int = 128,
image_dim: int = 784, clip_value: float = 0.01):
self.latent_dim = latent_dim
self.clip_value = clip_value
# Generator
self.G = nn.Sequential(
nn.Linear(latent_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, image_dim),
nn.Tanh() # WGAN typically uses tanh
)
# Critic (not discriminator - outputs unbounded score)
self.C = nn.Sequential(
nn.Linear(image_dim, hidden_dim),
nn.LeakyReLU(0.2),
nn.Linear(hidden_dim, hidden_dim),
nn.LeakyReLU(0.2),
nn.Linear(hidden_dim, 1) # No sigmoid!
)
def clip_weights(self):
"""Enforce Lipschitz constraint via weight clipping."""
for param in self.C.parameters():
param.data.clamp_(-self.clip_value, self.clip_value)
def train_step(self, real_data: torch.Tensor,
g_optimizer: optim.Optimizer,
c_optimizer: optim.Optimizer,
n_critic: int = 5) -> Dict[str, float]:
batch_size = real_data.size(0)
device = real_data.device
# Train Critic multiple times
c_losses = []
for _ in range(n_critic):
c_optimizer.zero_grad()
# Sample noise
z = torch.randn(batch_size, self.latent_dim).to(device)
fake_data = self.G(z).detach()
# Critic scores
real_score = self.C(real_data)
fake_score = self.C(fake_data)
# Wasserstein distance: maximize E[C(real)] - E[C(fake)]
# (equivalent to minimize negative)
c_loss = -(real_score.mean() - fake_score.mean())
c_loss.backward()
c_optimizer.step()
# Clip weights
self.clip_weights()
c_losses.append(c_loss.item())
# Train Generator
g_optimizer.zero_grad()
z = torch.randn(batch_size, self.latent_dim).to(device)
fake_data = self.G(z)
fake_score = self.C(fake_data)
# Maximize E[C(fake)] (make fake samples score high)
g_loss = -fake_score.mean()
g_loss.backward()
g_optimizer.step()
return {
'c_loss': np.mean(c_losses),
'g_loss': g_loss.item(),
'wasserstein_dist': -np.mean(c_losses) # Approximate W-dist
}
class WGANGP:
"""
WGAN with Gradient Penalty (Gulrajani et al., 2017).
Uses gradient penalty instead of weight clipping for better Lipschitz enforcement.
"""
def __init__(self, latent_dim: int = 100, hidden_dim: int = 128,
image_dim: int = 784, lambda_gp: float = 10.0):
self.latent_dim = latent_dim
self.lambda_gp = lambda_gp
self.G = nn.Sequential(
nn.Linear(latent_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, image_dim),
nn.Tanh()
)
self.C = nn.Sequential(
nn.Linear(image_dim, hidden_dim),
nn.LeakyReLU(0.2),
nn.Linear(hidden_dim, hidden_dim),
nn.LeakyReLU(0.2),
nn.Linear(hidden_dim, 1)
)
def compute_gradient_penalty(self, real_data: torch.Tensor,
fake_data: torch.Tensor) -> torch.Tensor:
"""Compute gradient penalty: (||∇_x C(x)|| - 1)^2"""
batch_size = real_data.size(0)
# Random interpolation
alpha = torch.rand(batch_size, 1).to(real_data.device)
interpolates = (alpha * real_data + ((1 - alpha) * fake_data)).requires_grad_(True)
critic_interpolates = self.C(interpolates)
# Compute gradients
gradients = torch.autograd.grad(
outputs=critic_interpolates,
inputs=interpolates,
grad_outputs=torch.ones_like(critic_interpolates),
create_graph=True,
retain_graph=True,
only_inputs=True
)[0]
gradients = gradients.view(batch_size, -1)
gradient_norm = gradients.norm(2, dim=1)
penalty = ((gradient_norm - 1) ** 2).mean()
return penalty
def train_step(self, real_data: torch.Tensor,
g_optimizer: optim.Optimizer,
c_optimizer: optim.Optimizer,
n_critic: int = 5) -> Dict[str, float]:
batch_size = real_data.size(0)
device = real_data.device
c_losses = []
for _ in range(n_critic):
c_optimizer.zero_grad()
z = torch.randn(batch_size, self.latent_dim).to(device)
fake_data = self.G(z).detach()
real_score = self.C(real_data)
fake_score = self.C(fake_data)
# Gradient penalty
gp = self.compute_gradient_penalty(real_data, fake_data)
# Total critic loss
c_loss = -(real_score.mean() - fake_score.mean()) + self.lambda_gp * gp
c_loss.backward()
c_optimizer.step()
c_losses.append(c_loss.item())
# Train Generator
g_optimizer.zero_grad()
z = torch.randn(batch_size, self.latent_dim).to(device)
fake_data = self.G(z)
fake_score = self.C(fake_data)
g_loss = -fake_score.mean()
g_loss.backward()
g_optimizer.step()
return {
'c_loss': np.mean(c_losses),
'g_loss': g_loss.item(),
'gradient_penalty': gp.item()
}
class SpectralNorm:
"""
Spectral Normalization wrapper for enforcing Lipschitz constraint via power iteration.
More stable than weight clipping, used in SAGAN and BigGAN.
"""
def __init__(self, module: nn.Module, n_power_iterations: int = 1):
self.module = module
self.n_power_iterations = n_power_iterations
# Initialize u vector for power iteration
weight = self.module.weight
height = weight.size(0)
u = F.normalize(weight.new_empty(height).normal_(0, 1), dim=0)
self.module.register_buffer('spectral_norm_u', u)
def compute_weight(self) -> torch.Tensor:
"""Compute spectrally normalized weight using power iteration."""
weight = self.module.weight
u = self.module.spectral_norm_u
# Power iteration to approximate largest singular value
with torch.no_grad():
for _ in range(self.n_power_iterations):
v = F.normalize(torch.matmul(weight.t(), u), dim=0)
u = F.normalize(torch.matmul(weight, v), dim=0)
# Update buffer
self.module.spectral_norm_u.copy_(u)
# Compute spectral norm
sigma = torch.matmul(torch.matmul(u, weight), v)
return weight / sigma
def __call__(self, x: torch.Tensor) -> torch.Tensor:
weight = self.compute_weight()
return F.linear(x, weight, self.module.bias)
class SNGAN:
"""
GAN with Spectral Normalization for Lipschitz constraint.
"""
def __init__(self, latent_dim: int = 100, hidden_dim: int = 128, image_dim: int = 784):
self.latent_dim = latent_dim
# Generator (no spectral norm needed)
self.G = nn.Sequential(
nn.Linear(latent_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, image_dim),
nn.Tanh()
)
# Discriminator with spectral normalization
self.D = nn.Sequential(
self._sn_layer(image_dim, hidden_dim),
nn.LeakyReLU(0.2),
self._sn_layer(hidden_dim, hidden_dim),
nn.LeakyReLU(0.2),
self._sn_layer(hidden_dim, 1)
)
def _sn_layer(self, in_dim: int, out_dim: int) -> nn.Module:
"""Create linear layer with spectral normalization."""
layer = nn.Linear(in_dim, out_dim)
return SpectralNorm(layer)
def train_step(self, real_data: torch.Tensor,
g_optimizer: optim.Optimizer,
d_optimizer: optim.Optimizer) -> Dict[str, float]:
batch_size = real_data.size(0)
device = real_data.device
# Train Discriminator
d_optimizer.zero_grad()
real_output = self.D(real_data)
z = torch.randn(batch_size, self.latent_dim).to(device)
fake_data = self.G(z).detach()
fake_output = self.D(fake_data)
# Hinge loss (common with spectral norm)
d_loss = torch.mean(F.relu(1.0 - real_output)) + \
torch.mean(F.relu(1.0 + fake_output))
d_loss.backward()
d_optimizer.step()
# Train Generator
g_optimizer.zero_grad()
fake_output = self.D(self.G(z))
g_loss = -fake_output.mean()
g_loss.backward()
g_optimizer.step()
return {
'd_loss': d_loss.item(),
'g_loss': g_loss.item()
}
def train_gan(method: str = 'wgan_gp', dataset: str = 'mnist', epochs: int = 50):
"""
Training pipeline comparing different GAN variants.
"""
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,)) # Scale to [-1, 1] for WGAN
])
train_data = datasets.MNIST('./data', train=True, download=True, transform=transform)
train_loader = DataLoader(train_data, batch_size=64, shuffle=True)
# Initialize model
if method == 'minimax':
model = MinimaxGAN()
g_opt = optim.Adam(model.G.parameters(), lr=0.0002, betas=(0.5, 0.999))
d_opt = optim.Adam(model.D.parameters(), lr=0.0002, betas=(0.5, 0.999))
elif method == 'wgan':
model = WGAN(clip_value=0.01)
g_opt = optim.RMSprop(model.G.parameters(), lr=0.00005)
d_opt = optim.RMSprop(model.C.parameters(), lr=0.00005)
elif method == 'wgan_gp':
model = WGANGP(lambda_gp=10.0)
g_opt = optim.Adam(model.G.parameters(), lr=0.0001, betas=(0.0, 0.9))
d_opt = optim.Adam(model.C.parameters(), lr=0.0001, betas=(0.0, 0.9))
elif method == 'spectral':
model = SNGAN()
g_opt = optim.Adam(model.G.parameters(), lr=0.0001, betas=(0.0, 0.9))
d_opt = optim.Adam(model.D.parameters(), lr=0.0004, betas=(0.0, 0.9))
model.G.to(device)
if hasattr(model, 'D'):
model.D.to(device)
if hasattr(model, 'C'):
model.C.to(device)
history = {'g_loss': [], 'd_loss': [], 'extra': []}
for epoch in range(epochs):
epoch_g_loss = 0
epoch_d_loss = 0
for batch_idx, (data, _) in enumerate(train_loader):
data = data.view(-1, 784).to(device)
if method == 'minimax':
losses = model.train_step(data, g_opt, d_opt)
elif method == 'wgan':
losses = model.train_step(data, g_opt, d_opt, n_critic=5)
history['extra'].append(losses['wasserstein_dist'])
elif method == 'wgan_gp':
losses = model.train_step(data, g_opt, d_opt, n_critic=5)
history['extra'].append(losses['gradient_penalty'])
elif method == 'spectral':
losses = model.train_step(data, g_opt, d_opt)
epoch_g_loss += losses['g_loss']
epoch_d_loss += losses['d_loss']
avg_g = epoch_g_loss / len(train_loader)
avg_d = epoch_d_loss / len(train_loader)
history['g_loss'].append(avg_g)
history['d_loss'].append(avg_d)
if epoch % 5 == 0:
print(f"Epoch {epoch}/{epochs} | G Loss: {avg_g:.4f} | D Loss: {avg_d:.4f}")
visualize_gan_results(model, method, history, device)
return model, history
def visualize_gan_results(model, method: str, history: Dict, device: torch.device):
"""
Visualization comparing GAN training stability and sample quality.
"""
fig, axes = plt.subplots(2, 3, figsize=(18, 12))
# 1. Generator and Discriminator/Critic loss curves
ax = axes[0, 0]
epochs = range(len(history['g_loss']))
ax.plot(epochs, history['g_loss'], label='Generator', color='blue')
ax.plot(epochs, history['d_loss'], label='Discriminator/Critic', color='red')
ax.set_xlabel('Epoch')
ax.set_ylabel('Loss')
ax.set_title(f'{method.upper()} Training Loss')
ax.legend()
ax.grid(True, alpha=0.3)
# 2. Additional metric (W-distance or Gradient Penalty)
if 'extra' in history and history['extra']:
ax = axes[0, 1]
ax.plot(history['extra'], color='green')
if method == 'wgan':
ax.set_ylabel('Wasserstein Distance')
elif method == 'wgan_gp':
ax.set_ylabel('Gradient Penalty')
ax.set_title('Training Dynamics')
ax.grid(True, alpha=0.3)
# 3. Generated samples grid
ax = axes[0, 2]
model.G.eval()
with torch.no_grad():
z = torch.randn(64, model.latent_dim).to(device)
samples = model.G(z).cpu().view(-1, 28, 28)
grid = torch.zeros(8*28, 8*28)
for i in range(8):
for j in range(8):
grid[i*28:(i+1)*28, j*28:(j+1)*28] = samples[i*8 + j]
ax.imshow(grid.numpy(), cmap='gray')
ax.set_title('Generated Samples')
ax.axis('off')
# 4. Interpolation in latent space
ax = axes[1, 0]
with torch.no_grad():
z1 = torch.randn(1, model.latent_dim).to(device)
z2 = torch.randn(1, model.latent_dim).to(device)
interpolations = []
for alpha in np.linspace(0, 1, 8):
z_interp = (1 - alpha) * z1 + alpha * z2
sample = model.G(z_interp).cpu().view(28, 28)
interpolations.append(sample)
interp_grid = np.concatenate(interpolations, axis=1)
ax.imshow(interp_grid.numpy(), cmap='gray')
ax.set_title('Latent Space Interpolation')
ax.axis('off')
# 5. Mode coverage analysis (simplified)
ax = axes[1, 1]
# Generate many samples and check digit diversity
with torch.no_grad():
z = torch.randn(1000, model.latent_dim).to(device)
samples = model.G(z).cpu().view(-1, 784)
# Compute sample variance as proxy for mode coverage
sample_vars = torch.var(samples, dim=0).numpy()
ax.hist(sample_vars, bins=50, alpha=0.7, color='purple')
ax.set_xlabel('Pixel Variance Across Samples')
ax.set_ylabel('Count')
ax.set_title('Sample Diversity (Mode Coverage Proxy)')
ax.grid(True, alpha=0.3)
# 6. Comparison of methods (if multiple trained)
ax = axes[1, 2]
methods_comparison = {
'JS Divergence (Minimax)': {'stability': 0.3, 'sample_quality': 0.7},
'Weight Clipping (WGAN)': {'stability': 0.6, 'sample_quality': 0.8},
'Gradient Penalty (WGAN-GP)': {'stability': 0.9, 'sample_quality': 0.9},
'Spectral Norm (SNGAN)': {'stability': 0.85, 'sample_quality': 0.85}
}
if method in [m.lower().replace(' ', '_').replace('(', '').replace(')', '')
for m in methods_comparison.keys()]:
names = list(methods_comparison.keys())
stability = [v['stability'] for v in methods_comparison.values()]
quality = [v['sample_quality'] for v in methods_comparison.values()]
x = np.arange(len(names))
width = 0.35
ax.bar(x - width/2, stability, width, label='Training Stability', alpha=0.8)
ax.bar(x + width/2, quality, width, label='Sample Quality', alpha=0.8)
ax.set_ylabel('Score')
ax.set_title('GAN Variant Comparison')
ax.set_xticks(x)
ax.set_xticklabels(names, rotation=15, ha='right')
ax.legend()
ax.set_ylim(0, 1)
plt.tight_layout()
plt.savefig(f'gan_analysis_{method}.png', dpi=300)
plt.show()
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--method', type=str, default='wgan_gp',
choices=['minimax', 'wgan', 'wgan_gp', 'spectral'])
parser.add_argument('--epochs', type=int, default=50)
args = parser.parse_args()
train_gan(args.method, epochs=args.epochs)13.2.2 条件GAN（cGAN）的投影判别器（Projection Discriminator）与辅助分类器（AC-GAN），Progressive GAN与StyleGAN的自适应实例归一化（AdaIN）与风格映射网络
条件生成对抗网络通过辅助信息引导生成过程，类别标签或属性向量作为额外输入约束生成分布的特定维度。投影判别器将条件嵌入与中间特征内积作为输出分量的计算基础，避免条件信息在早期层被忽视，通过点积运算显式建模类别与特征的相容性。辅助分类器生成器同时优化对抗目标与分类准确率，强制生成样本具备明确的类别归属，提升多类别场景下的模式覆盖率。
渐进式增长策略从低分辨率开始训练，逐步添加新层增加生成细节，平滑的层间过渡避免 sudden 结构变化导致的训练震荡。风格生成网络解耦潜在代码与图像生成过程，通过映射网络将简单分布转换为解耦的中间表示，自适应实例归一化将风格信息注入各层特征，通过学习得到的尺度与偏移参数控制局部统计特性，实现细粒度的视觉属性操控而不影响整体内容结构。
Python
"""
Script: conditional_gan_stylegan.py
Content: Conditional GAN with Projection Discriminator and AC-GAN,
Progressive GAN growing, and StyleGAN with AdaIN and Mapping Network.
Usage: python conditional_gan_stylegan.py --model_type stylegan --dataset mnist --epochs 100
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import numpy as np
import matplotlib.pyplot as plt
from typing import Tuple, Dict, List
import math
import argparse
class ProjectionDiscriminator(nn.Module):
"""
Projection Discriminator (Miyato & Koyama, 2018).
Uses inner product between conditional embedding and feature.
"""
def __init__(self, num_classes: int = 10, image_dim: int = 784,
hidden_dim: int = 256, embed_dim: int = 256):
super().__init__()
self.num_classes = num_classes
# Feature extractor
self.feature_extractor = nn.Sequential(
nn.Linear(image_dim, hidden_dim),
nn.LeakyReLU(0.2),
nn.Linear(hidden_dim, hidden_dim),
nn.LeakyReLU(0.2),
nn.Linear(hidden_dim, hidden_dim)
)
# Conditional embedding
self.embed = nn.Embedding(num_classes, embed_dim)
# Unconditional output
self.output_layer = nn.Linear(hidden_dim, 1)
def forward(self, x: torch.Tensor, labels: torch.Tensor = None) -> torch.Tensor:
features = self.feature_extractor(x)
# Unconditional score
output = self.output_layer(features)
# Conditional projection (inner product)
if labels is not None:
embed = self.embed(labels) # (batch, embed_dim)
# Projection term: sum(embed * features)
projection = torch.sum(embed * features, dim=1, keepdim=True)
output = output + projection
return output
class ACGANGenerator(nn.Module):
"""
AC-GAN Generator with auxiliary classification support.
Generates class-conditional samples and predicts class.
"""
def __init__(self, latent_dim: int = 100, num_classes: int = 10,
image_dim: int = 784, hidden_dim: int = 256):
super().__init__()
self.latent_dim = latent_dim
self.num_classes = num_classes
# Class embedding
self.class_embed = nn.Embedding(num_classes, latent_dim)
# Generator network
self.net = nn.Sequential(
nn.Linear(latent_dim * 2, hidden_dim), # z + class_embed
nn.ReLU(),
nn.Linear(hidden_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, image_dim),
nn.Tanh()
)
def forward(self, z: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
class_vec = self.class_embed(labels)
x = torch.cat([z, class_vec], dim=1)
return self.net(x)
class ACGANDiscriminator(nn.Module):
"""
AC-GAN Discriminator with real/fake and auxiliary classifier outputs.
"""
def __init__(self, image_dim: int = 784, num_classes: int = 10, hidden_dim: int = 256):
super().__init__()
self.feature_extractor = nn.Sequential(
nn.Linear(image_dim, hidden_dim),
nn.LeakyReLU(0.2),
nn.Linear(hidden_dim, hidden_dim),
nn.LeakyReLU(0.2)
)
# Real/fake classifier
self.discriminator = nn.Linear(hidden_dim, 1)
# Auxiliary classifier
self.classifier = nn.Linear(hidden_dim, num_classes)
def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
features = self.feature_extractor(x)
real_fake = self.discriminator(features)
classes = self.classifier(features)
return real_fake, classes
class AdaptiveInstanceNorm(nn.Module):
"""
Adaptive Instance Normalization (AdaIN) for StyleGAN.
Transforms content feature statistics based on style input.
"""
def __init__(self, num_features: int, style_dim: int):
super().__init__()
self.norm = nn.InstanceNorm2d(num_features, affine=False)
# Learn to predict scale and shift from style
self.style_scale = nn.Linear(style_dim, num_features)
self.style_shift = nn.Linear(style_dim, num_features)
def forward(self, content: torch.Tensor, style: torch.Tensor) -> torch.Tensor:
# Normalize content
normalized = self.norm(content)
# Get style parameters
scale = self.style_scale(style).unsqueeze(2).unsqueeze(3)
shift = self.style_shift(style).unsqueeze(2).unsqueeze(3)
# Apply affine transformation
return normalized * (1 + scale) + shift
class StyleGANMapping(nn.Module):
"""
Mapping Network for StyleGAN (Karras et al., 2019).
Transforms Z-space to W-space for disentangled representations.
"""
def __init__(self, z_dim: int = 512, w_dim: int = 512, num_layers: int = 8):
super().__init__()
layers = []
for i in range(num_layers):
in_dim = z_dim if i == 0 else w_dim
out_dim = w_dim
layers.append(nn.Linear(in_dim, out_dim))
layers.append(nn.LeakyReLU(0.2))
self.net = nn.Sequential(*layers)
def forward(self, z: torch.Tensor) -> torch.Tensor:
# Normalize latent code
z = F.normalize(z, dim=1)
return self.net(z)
class StyleGANGenerator(nn.Module):
"""
Simplified StyleGAN Generator for MNIST (28x28).
Demonstrates AdaIN injection and constant input.
"""
def __init__(self, z_dim: int = 512, w_dim: int = 512,
image_size: int = 28, channels: int = 1):
super().__init__()
self.mapping = StyleGANMapping(z_dim, w_dim)
# Learnable constant input (4x4 spatial)
self.const_input = nn.Parameter(torch.randn(1, 256, 4, 4))
# Style-modulated layers
self.layer1 = self._make_style_layer(256, 128, w_dim)
self.layer2 = self._make_style_layer(128, 64, w_dim)
self.layer3 = self._make_style_layer(64, 32, w_dim)
# Upsampling to target size
self.to_rgb = nn.Conv2d(32, channels, 3, 1, 1)
self.upsample = nn.Upsample(size=(image_size, image_size), mode='bilinear')
# Broadcasted noise (simplified)
self.noise1 = nn.Parameter(torch.zeros(1, 128, 8, 8))
self.noise2 = nn.Parameter(torch.zeros(1, 64, 14, 14))
def _make_style_layer(self, in_channels: int, out_channels: int, style_dim: int):
return nn.ModuleDict({
'conv': nn.Conv2d(in_channels, out_channels, 3, 1, 1),
'adain': AdaptiveInstanceNorm(out_channels, style_dim),
'activation': nn.LeakyReLU(0.2)
})
def forward(self, z: torch.Tensor) -> torch.Tensor:
# Map to W space
w = self.mapping(z)
# Start with constant input
x = self.const_input.repeat(z.size(0), 1, 1, 1)
# Layer 1: 4x4 -> 8x8
x = F.interpolate(x, scale_factor=2)
x = self.layer1['conv'](x)
x = self.layer1['adain'](x, w)
x = self.layer1['activation'](x + self.noise1.expand(x.size(0), -1, -1, -1))
# Layer 2: 8x8 -> 14x14 (approximate for 28x28 target)
x = F.interpolate(x, size=(14, 14))
x = self.layer2['conv'](x)
x = self.layer2['adain'](x, w)
x = self.layer2['activation'](x + self.noise2.expand(x.size(0), -1, -1, -1))
# Layer 3: 14x14 -> 28x28
x = F.interpolate(x, size=(28, 28))
x = self.layer3['conv'](x)
x = self.layer3['adain'](x, w)
x = self.layer3['activation'](x)
# To RGB
x = torch.tanh(self.to_rgb(x))
return x.view(-1, 28*28) # Flatten for consistency with other models
class ProgressiveGAN:
"""
Progressive GAN training with fade-in mechanism.
Grows from 4x4 to 28x28 progressively.
"""
def __init__(self, max_size: int = 28, latent_dim: int = 100):
self.max_size = max_size
self.current_size = 4
self.latent_dim = latent_dim
self.transition_alpha = 0.0 # For fade-in
self.generator = self._build_generator(4)
self.discriminator = self._build_discriminator(4)
def _build_generator(self, size: int) -> nn.Module:
"""Build generator for specific resolution."""
layers = []
log_size = int(math.log2(size))
in_ch = 512
for i in range(log_size - 2):
out_ch = 512 if i < 2 else 256
layers.extend([
nn.ConvTranspose2d(in_ch, out_ch, 4, 2, 1),
nn.BatchNorm2d(out_ch),
nn.ReLU()
])
in_ch = out_ch
layers.extend([
nn.Conv2d(in_ch, 1, 3, 1, 1),
nn.Tanh()
])
return nn.Sequential(*layers)
def grow_network(self, new_size: int):
"""Add new layer and fade in."""
# This is a simplified version - real implementation requires
# careful handling of state dicts and smooth transition
self.current_size = new_size
self.transition_alpha = 0.0
def get_alpha(self) -> float:
"""Get current fade-in alpha."""
return min(1.0, self.transition_alpha + 0.01)
def train_conditional_gan(model_type: str = 'acgan', dataset: str = 'mnist', epochs: int = 100):
"""
Training pipeline for conditional GAN variants.
"""
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,))
])
train_data = datasets.MNIST('./data', train=True, download=True, transform=transform)
train_loader = DataLoader(train_data, batch_size=64, shuffle=True)
num_classes = 10
if model_type == 'acgan':
G = ACGANGenerator(latent_dim=100, num_classes=num_classes).to(device)
D = ACGANDiscriminator(num_classes=num_classes).to(device)
g_opt = torch.optim.Adam(G.parameters(), lr=0.0002, betas=(0.5, 0.999))
d_opt = torch.optim.Adam(D.parameters(), lr=0.0002, betas=(0.5, 0.999))
elif model_type == 'projection':
G = ACGANGenerator(latent_dim=100, num_classes=num_classes).to(device) # Reuse architecture
D = ProjectionDiscriminator(num_classes=num_classes).to(device)
g_opt = torch.optim.Adam(G.parameters(), lr=0.0001, betas=(0.0, 0.9))
d_opt = torch.optim.Adam(D.parameters(), lr=0.0004, betas=(0.0, 0.9))
elif model_type == 'stylegan':
G = StyleGANGenerator(z_dim=512, w_dim=512).to(device)
D = ACGANDiscriminator(num_classes=num_classes).to(device) # Simple disc for demo
g_opt = torch.optim.Adam(G.parameters(), lr=0.001, betas=(0.0, 0.99))
d_opt = torch.optim.Adam(D.parameters(), lr=0.001, betas=(0.0, 0.99))
else:
raise ValueError(f"Unknown model type: {model_type}")
history = {'g_loss': [], 'd_loss': [], 'aux_acc': []}
for epoch in range(epochs):
epoch_g_loss = 0
epoch_d_loss = 0
epoch_aux_acc = 0
for batch_idx, (data, labels) in enumerate(train_loader):
data = data.view(-1, 784).to(device)
labels = labels.to(device)
batch_size = data.size(0)
# Train Discriminator
d_opt.zero_grad()
if model_type == 'acgan':
real_pred, real_class = D(data)
d_real_loss = F.binary_cross_entropy_with_logits(real_pred, torch.ones_like(real_pred))
z = torch.randn(batch_size, 100).to(device)
fake_labels = torch.randint(0, num_classes, (batch_size,)).to(device)
fake_data = G(z, fake_labels)
fake_pred, fake_class = D(fake_data.detach())
d_fake_loss = F.binary_cross_entropy_with_logits(fake_pred, torch.zeros_like(fake_pred))
# Auxiliary classification loss
aux_loss = F.cross_entropy(real_class, labels) + F.cross_entropy(fake_class, fake_labels)
d_loss = d_real_loss + d_fake_loss + aux_loss
d_loss.backward()
d_opt.step()
# Accuracy
_, pred_class = torch.max(real_class, 1)
acc = (pred_class == labels).float().mean().item()
elif model_type == 'projection':
real_pred = D(data, labels)
d_real_loss = F.relu(1 - real_pred).mean()
z = torch.randn(batch_size, 100).to(device)
fake_labels = torch.randint(0, num_classes, (batch_size,)).to(device)
fake_data = G(z, fake_labels)
fake_pred = D(fake_data, fake_labels)
d_fake_loss = F.relu(1 + fake_pred).mean()
d_loss = d_real_loss + d_fake_loss
d_loss.backward()
d_opt.step()
acc = 0 # No auxiliary classifier
# Train Generator
g_opt.zero_grad()
if model_type == 'acgan':
z = torch.randn(batch_size, 100).to(device)
fake_labels = torch.randint(0, num_classes, (batch_size,)).to(device)
fake_data = G(z, fake_labels)
fake_pred, fake_class = D(fake_data)
g_loss = F.binary_cross_entropy_with_logits(fake_pred, torch.ones_like(fake_pred))
g_loss += F.cross_entropy(fake_class, fake_labels) # Auxiliary
elif model_type == 'projection':
z = torch.randn(batch_size, 100).to(device)
fake_labels = torch.randint(0, num_classes, (batch_size,)).to(device)
fake_data = G(z, fake_labels)
fake_pred = D(fake_data, fake_labels)
g_loss = -fake_pred.mean()
elif model_type == 'stylegan':
z = torch.randn(batch_size, 512).to(device)
fake_data = G(z)
fake_pred, _ = D(fake_data)
g_loss = -fake_pred.mean()
g_loss.backward()
g_opt.step()
epoch_g_loss += g_loss.item()
epoch_d_loss += d_loss.item()
epoch_aux_acc += acc
avg_g = epoch_g_loss / len(train_loader)
avg_d = epoch_d_loss / len(train_loader)
avg_acc = epoch_aux_acc / len(train_loader)
history['g_loss'].append(avg_g)
history['d_loss'].append(avg_d)
history['aux_acc'].append(avg_acc)
if epoch % 10 == 0:
print(f"Epoch {epoch}/{epochs} | G: {avg_g:.4f} | D: {avg_d:.4f} | Acc: {avg_acc:.2%}")
visualize_conditional_gan(G, D, model_type, history, device, num_classes)
return G, D, history
def visualize_conditional_gan(G, D, model_type: str, history: Dict,
device: torch.device, num_classes: int):
"""
Visualization for conditional GAN showing class-conditional generation.
"""
fig, axes = plt.subplots(2, 3, figsize=(18, 12))
# 1. Training curves
ax = axes[0, 0]
epochs = range(len(history['g_loss']))
ax.plot(epochs, history['g_loss'], label='Generator', color='blue')
ax.plot(epochs, history['d_loss'], label='Discriminator', color='red')
ax.set_xlabel('Epoch')
ax.set_ylabel('Loss')
ax.set_title(f'{model_type.upper()} Training Loss')
ax.legend()
ax.grid(True, alpha=0.3)
# 2. Auxiliary classifier accuracy (if applicable)
if model_type == 'acgan' and history['aux_acc']:
ax = axes[0, 1]
ax.plot(epochs, history['aux_acc'], color='green', marker='o')
ax.set_xlabel('Epoch')
ax.set_ylabel('Classification Accuracy')
ax.set_title('Auxiliary Classifier Accuracy')
ax.grid(True, alpha=0.3)
ax.set_ylim(0, 1)
# 3. Class-conditional generation grid
ax = axes[0, 2]
G.eval()
with torch.no_grad():
# Generate one sample per class
samples_per_class = 8
z = torch.randn(num_classes * samples_per_class, 100 if model_type != 'stylegan' else 512).to(device)
labels = torch.repeat_interleave(torch.arange(num_classes), samples_per_class).to(device)
if model_type == 'stylegan':
fake_data = G(z)
else:
fake_data = G(z, labels)
fake_data = fake_data.cpu().view(-1, 28, 28)
# Arrange in grid: rows = classes, cols = samples
grid = torch.zeros(num_classes * 28, samples_per_class * 28)
for i in range(num_classes):
for j in range(samples_per_class):
idx = i * samples_per_class + j
grid[i*28:(i+1)*28, j*28:(j+1)*28] = fake_data[idx]
ax.imshow(grid.numpy(), cmap='gray')
ax.set_title('Class-Conditional Generation\n(Rows: classes 0-9)')
ax.axis('off')
# 4. Style mixing (StyleGAN specific)
if model_type == 'stylegan':
ax = axes[1, 0]
with torch.no_grad():
# Generate with two different styles at different layers
z1 = torch.randn(1, 512).to(device)
z2 = torch.randn(1, 512).to(device)
# Simplified style mixing visualization
samples = []
for alpha in [0.0, 0.25, 0.5, 0.75, 1.0]:
z_mixed = (1 - alpha) * z1 + alpha * z2
sample = G(z_mixed).cpu().view(28, 28)
samples.append(sample)
mixed_grid = np.concatenate(samples, axis=1)
ax.imshow(mixed_grid.numpy(), cmap='gray')
ax.set_title('Style Interpolation in W-space')
ax.axis('off')
# 5. Per-class generation quality (diversity within class)
ax = axes[1, 1]
with torch.no_grad():
class_vars = []
for c in range(num_classes):
z = torch.randn(100, 100 if model_type != 'stylegan' else 512).to(device)
labels = torch.full((100,), c, dtype=torch.long).to(device)
if model_type == 'stylegan':
samples = G(z)
else:
samples = G(z, labels)
# Measure variance within class
var = torch.var(samples, dim=0).mean().item()
class_vars.append(var)
ax.bar(range(num_classes), class_vars, color='purple', alpha=0.7)
ax.set_xlabel('Digit Class')
ax.set_ylabel('Intra-class Variance')
ax.set_title('Per-Class Sample Diversity')
ax.grid(True, alpha=0.3, axis='y')
# 6. Projection discriminator analysis (if applicable)
if model_type == 'projection':
ax = axes[1, 2]
# Visualize embedding space
if hasattr(D, 'embed'):
embed_weights = D.embed.weight.data.cpu().numpy()
from sklearn.decomposition import PCA
pca = PCA(n_components=2)
embed_2d = pca.fit_transform(embed_weights)
scatter = ax.scatter(embed_2d[:, 0], embed_2d[:, 1],
c=range(num_classes), cmap='tab10', s=200)
for i, (x, y) in enumerate(embed_2d):
ax.annotate(str(i), (x, y), ha='center', va='center')
ax.set_title('Projection Discriminator\nClass Embeddings (PCA)')
plt.colorbar(scatter, ax=ax)
plt.tight_layout()
plt.savefig(f'conditional_gan_{model_type}.png', dpi=300)
plt.show()
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--model_type', type=str, default='acgan',
choices=['acgan', 'projection', 'stylegan'])
parser.add_argument('--epochs', type=int, default=100)
args = parser.parse_args()
train_conditional_gan(args.model_type, epochs=args.epochs)13.2.3 训练稳定性：谱归一化（Spectral Normalization）的幂迭代实现，DiffAugment与AdaAugment的数据增强策略
谱归一化通过限制网络层谱范数保证判别器满足Lipschitz连续性，幂迭代算法高效估计最大奇异值而不需完整矩阵分解。前向传播中归一化权重矩阵，反向传播通过幂迭代向量传播梯度，实现单步近似奇异值计算。该方法相比梯度惩罚无需额外超参数调节，相比权重裁剪避免容量削减，在大规模图像生成任务中展现出稳定的训练动态与优异的样本质量。
数据增强策略DiffAugment与AdaAugment针对GAN训练中的过拟合问题，在生成样本与真实数据上应用可微分几何变换。DiffAugment确保增强操作对判别器可微分，允许梯度流经增强层回传至生成器。AdaAugment通过强化学习或密度匹配搜索最优增强策略分布，在保持变换可逆性的前提下最大化增强样本的信息熵，提升判别器的泛化能力并防止模式崩溃。
Python
"""
Script: gan_stability_spectral_diffaugment.py
Content: Spectral Normalization with power iteration, DiffAugment and AdaAugment
for training stability in GANs.
Usage: python gan_stability_spectral_diffaugment.py --augment_type diffaugment --epochs 50
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import numpy as np
import matplotlib.pyplot as plt
from typing import Tuple, Dict, List
import argparse
class SpectralNormLayer(nn.Module):
"""
Spectral Normalization implementation with power iteration.
Wraps any linear/conv layer and enforces Lipschitz constraint via σ(W) = 1.
"""
def __init__(self, module: nn.Module, n_power_iterations: int = 1):
super().__init__()
self.module = module
self.n_power_iterations = n_power_iterations
# Get weight dimensions
if isinstance(module, nn.Linear):
height = module.weight.size(0)
elif isinstance(module, nn.Conv2d):
height = module.weight.size(0) * module.weight.size(2) * module.weight.size(3)
else:
raise NotImplementedError
# Initialize u vector for power iteration
u = F.normalize(module.weight.new_empty(height).normal_(0, 1), dim=0)
self.register_buffer('u', u)
self.register_buffer('v', None)
# Hook for applying spectral norm
self.module._forward_pre_hooks.clear()
self.module.register_forward_pre_hook(self._compute_weight)
def _compute_weight(self, module, input):
"""Compute spectrally normalized weight using power iteration."""
weight = module.weight
if isinstance(module, nn.Conv2d):
# Reshape conv weight for matmul: (out_ch, in_ch*k*k)
weight_mat = weight.view(weight.size(0), -1)
else:
weight_mat = weight
# Power iteration to compute left and right singular vectors
with torch.no_grad():
for _ in range(self.n_power_iterations):
# v = W^T u / ||W^T u||
self.v = F.normalize(torch.matmul(weight_mat.t(), self.u), dim=0)
# u = W v / ||W v||
self.u = F.normalize(torch.matmul(weight_mat, self.v), dim=0)
# Compute spectral norm (largest singular value)
sigma = torch.matmul(torch.matmul(self.u, weight_mat), self.v)
# Normalize weight
module.weight.data = weight / sigma
def forward(self, *args):
return self.module(*args)
class DiffAugment:
"""
Differentiable Augmentation for GANs (Zhao et al., 2020).
Applies augmentations that are differentiable w.r.t. input.
"""
def __init__(self, policy: str = 'color,translation,cutout'):
self.policy = policy.split(',')
def __call__(self, x: torch.Tensor) -> torch.Tensor:
if 'color' in self.policy:
x = self._rand_brightness(x)
x = self._rand_saturation(x)
x = self._rand_contrast(x)
if 'translation' in self.policy:
x = self._rand_translation(x)
if 'cutout' in self.policy:
x = self._rand_cutout(x)
return x
def _rand_brightness(self, x: torch.Tensor) -> torch.Tensor:
"""Random brightness shift (differentiable)."""
return x + (torch.rand(x.size(0), 1, 1, 1, device=x.device) - 0.5) * 0.2
def _rand_saturation(self, x: torch.Tensor) -> torch.Tensor:
"""Random saturation adjustment."""
x_mean = x.mean(dim=1, keepdim=True)
return (x - x_mean) * (torch.rand(x.size(0), 1, 1, 1, device=x.device) * 2) + x_mean
def _rand_contrast(self, x: torch.Tensor) -> torch.Tensor:
"""Random contrast adjustment."""
x_mean = x.mean(dim=[1, 2, 3], keepdim=True)
return (x - x_mean) * (torch.rand(x.size(0), 1, 1, 1, device=x.device) + 0.5) + x_mean
def _rand_translation(self, x: torch.Tensor, ratio: float = 0.125) -> torch.Tensor:
"""Random translation (shift) using roll (differentiable via interpolation)."""
shift_x = int(x.size(2) * ratio * (torch.rand(1).item() - 0.5))
shift_y = int(x.size(3) * ratio * (torch.rand(1).item() - 0.5))
return torch.roll(x, shifts=(shift_x, shift_y), dims=(2, 3))
def _rand_cutout(self, x: torch.Tensor, ratio: float = 0.5) -> torch.Tensor:
"""Random cutout (zero masking)."""
cutout_size = int(x.size(2) * ratio)
if cutout_size <= 0:
return x
# Random position
cx = torch.randint(0, x.size(2) - cutout_size + 1, (1,)).item()
cy = torch.randint(0, x.size(3) - cutout_size + 1, (1,)).item()
x[:, :, cx:cx+cutout_size, cy:cy+cutout_size] = 0
return x
class AdaAugment:
"""
Adaptive Augmentation (Adaptive Discriminator Augmentation).
Adjusts augmentation probability based on overfitting metric.
"""
def __init__(self, initial_p: float = 0.0, target_rt: float = 0.6,
speed: float = 500.0, max_p: float = 0.8):
self.p = initial_p
self.target_rt = target_rt # Target for r_t (sign of overfitting)
self.speed = speed
self.max_p = max_p
# DiffAugment as base
self.augment = DiffAugment('color,translation,cutout')
# History for tracking
self.rt_history = []
def adapt_probability(self, real_pred: torch.Tensor, fake_pred: torch.Tensor):
"""
Adapt augmentation probability based on discriminator outputs.
r_t = sign(real_pred) indicating overfitting to real data.
"""
# Compute r_t (portion of real samples with positive prediction)
r_t = (real_pred > 0).float().mean().item()
self.rt_history.append(r_t)
# Adjust p: increase if overfitting (r_t > target), decrease if underfitting
if r_t > self.target_rt:
self.p += (r_t - self.target_rt) / self.speed
else:
self.p -= (self.target_rt - r_t) / self.speed
self.p = np.clip(self.p, 0.0, self.max_p)
def __call__(self, x: torch.Tensor, is_real: bool = True) -> Tuple[torch.Tensor, bool]:
"""Apply augmentation with probability p."""
if is_real and torch.rand(1).item() < self.p:
return self.augment(x), True
return x, False
class StableSNGAN:
"""
GAN with Spectral Normalization and Adaptive Augmentation.
Combines architectural and data-based stabilization.
"""
def __init__(self, latent_dim: int = 100, image_channels: int = 1,
use_spectral_norm: bool = True, use_adaptive_aug: bool = True):
self.latent_dim = latent_dim
# Generator (no spectral norm)
self.G = nn.Sequential(
nn.ConvTranspose2d(latent_dim, 256, 7, 1, 0), # 1x1 -> 7x7
nn.BatchNorm2d(256),
nn.ReLU(),
nn.ConvTranspose2d(256, 128, 4, 2, 1), # 7x7 -> 14x14
nn.BatchNorm2d(128),
nn.ReLU(),
nn.ConvTranspose2d(128, image_channels, 4, 2, 1), # 14x14 -> 28x28
nn.Tanh()
)
# Discriminator with spectral norm
if use_spectral_norm:
self.D = nn.Sequential(
SpectralNormLayer(nn.Conv2d(image_channels, 64, 4, 2, 1)), # 28->14
nn.LeakyReLU(0.2),
SpectralNormLayer(nn.Conv2d(64, 128, 4, 2, 1)), # 14->7
nn.LeakyReLU(0.2),
SpectralNormLayer(nn.Conv2d(128, 256, 3, 1, 1)),
nn.LeakyReLU(0.2),
SpectralNormLayer(nn.Conv2d(256, 1, 7, 1, 0)) # 7->1
)
else:
self.D = nn.Sequential(
nn.Conv2d(image_channels, 64, 4, 2, 1),
nn.LeakyReLU(0.2),
nn.Conv2d(64, 128, 4, 2, 1),
nn.LeakyReLU(0.2),
nn.Conv2d(128, 1, 7, 1, 0)
)
self.use_adaptive_aug = use_adaptive_aug
if use_adaptive_aug:
self.ada_aug = AdaAugment()
else:
self.diff_aug = DiffAugment()
def train_step(self, real_data: torch.Tensor,
g_optimizer: torch.optim.Optimizer,
d_optimizer: torch.optim.Optimizer) -> Dict[str, float]:
device = real_data.device
batch_size = real_data.size(0)
# Discriminator training
d_optimizer.zero_grad()
# Augment real data if enabled
if self.use_adaptive_aug:
real_augmented, _ = self.ada_aug(real_data, is_real=True)
else:
real_augmented = self.diff_aug(real_data)
real_pred = self.D(real_augmented)
# Generate fake data
z = torch.randn(batch_size, self.latent_dim, 1, 1).to(device)
fake_data = self.G(z)
# Augment fake data
if self.use_adaptive_aug:
fake_augmented, _ = self.ada_aug(fake_data, is_real=False)
else:
fake_augmented = self.diff_aug(fake_data)
fake_pred = self.D(fake_augmented.detach())
# Hinge loss
d_loss = torch.mean(F.relu(1.0 - real_pred)) + \
torch.mean(F.relu(1.0 + fake_pred))
d_loss.backward()
d_optimizer.step()
# Adapt augmentation probability
if self.use_adaptive_aug:
with torch.no_grad():
real_pred_raw = self.D(real_data)
self.ada_aug.adapt_probability(real_pred_raw.view(-1), fake_pred.view(-1))
# Generator training
g_optimizer.zero_grad()
fake_pred = self.D(fake_data)
g_loss = -fake_pred.mean()
g_loss.backward()
g_optimizer.step()
return {
'd_loss': d_loss.item(),
'g_loss': g_loss.item(),
'aug_p': self.ada_aug.p if self.use_adaptive_aug else 0.5
}
def train_stable_gan(augment_type: str = 'ada', epochs: int = 50):
"""
Training pipeline with stability techniques.
"""
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,))
])
train_data = datasets.MNIST('./data', train=True, download=True, transform=transform)
train_loader = DataLoader(train_data, batch_size=64, shuffle=True)
use_ada = (augment_type == 'ada')
model = StableSNGAN(use_spectral_norm=True, use_adaptive_aug=use_ada).to(device)
g_opt = torch.optim.Adam(model.G.parameters(), lr=0.0001, betas=(0.0, 0.99))
d_opt = torch.optim.Adam(model.D.parameters(), lr=0.0004, betas=(0.0, 0.99))
history = {
'd_loss': [], 'g_loss': [], 'aug_p': [],
'spectral_norms': [] # Track spectral norms over training
}
for epoch in range(epochs):
epoch_stats = {'d_loss': 0, 'g_loss': 0, 'aug_p': 0}
for batch_idx, (data, _) in enumerate(train_loader):
data = data.to(device)
losses = model.train_step(data, g_opt, d_opt)
for k in ['d_loss', 'g_loss', 'aug_p']:
if k in losses:
epoch_stats[k] += losses[k]
# Average
for k in epoch_stats:
epoch_stats[k] /= len(train_loader)
history[k].append(epoch_stats[k])
# Monitor spectral norms (first layer of D)
if hasattr(model.D[0], 'u'):
with torch.no_grad():
W = model.D[0].module.weight
if isinstance(model.D[0].module, nn.Conv2d):
W = W.view(W.size(0), -1)
sigma = torch.linalg.norm(W, ord=2).item()
history['spectral_norms'].append(sigma)
if epoch % 5 == 0:
print(f"Epoch {epoch}/{epochs} | D: {epoch_stats['d_loss']:.4f} | "
f"G: {epoch_stats['g_loss']:.4f} | Aug p: {epoch_stats['aug_p']:.3f}")
visualize_stability_analysis(model, history, augment_type)
return model, history
def visualize_stability_analysis(model, history: Dict, augment_type: str):
"""
Comprehensive visualization of training stability metrics.
"""
fig, axes = plt.subplots(2, 3, figsize=(18, 12))
# 1. Loss curves with stability analysis
ax = axes[0, 0]
epochs = range(len(history['g_loss']))
ax.plot(epochs, history['g_loss'], label='Generator', color='blue', linewidth=2)
ax.plot(epochs, history['d_loss'], label='Discriminator', color='red', linewidth=2)
# Add moving average for trend
window = 5
if len(history['g_loss']) > window:
g_ma = np.convolve(history['g_loss'], np.ones(window)/window, mode='valid')
ax.plot(range(window-1, len(history['g_loss'])), g_ma, '--', color='blue', alpha=0.5)
ax.set_xlabel('Epoch')
ax.set_ylabel('Loss')
ax.set_title('Training Stability (Loss Curves)')
ax.legend()
ax.grid(True, alpha=0.3)
# 2. Adaptive augmentation probability evolution
if augment_type == 'ada' and history['aug_p']:
ax = axes[0, 1]
ax.plot(epochs, history['aug_p'], color='green', marker='o', linewidth=2)
ax.axhline(y=0.6, color='red', linestyle='--', label='Target overfitting level')
ax.set_xlabel('Epoch')
ax.set_ylabel('Augmentation Probability')
ax.set_title('Adaptive Augmentation Schedule')
ax.legend()
ax.grid(True, alpha=0.3)
ax.set_ylim(0, 1)
# 3. Spectral norm evolution (Lipschitz constant monitoring)
if history['spectral_norms']:
ax = axes[0, 2]
ax.plot(epochs[:len(history['spectral_norms'])], history['spectral_norms'],
color='purple', marker='s', linewidth=2)
ax.axhline(y=1.0, color='red', linestyle='--', label='Target Lipschitz constant')
ax.set_xlabel('Epoch')
ax.set_ylabel('Spectral Norm σ(W)')
ax.set_title('Lipschitz Constant Enforcement')
ax.legend()
ax.grid(True, alpha=0.3)
# 4. Gradient penalty estimation (proxy for smoothness)
ax = axes[1, 0]
# Simulate gradient norm over training (would be computed during training in practice)
grad_norms = [1.0 / (1 + 0.1 * e) + np.random.normal(0, 0.05) for e in epochs]
ax.plot(epochs, grad_norms, color='orange', alpha=0.7)
ax.fill_between(epochs, [max(0, g - 0.1) for g in grad_norms],
[g + 0.1 for g in grad_norms], alpha=0.3, color='orange')
ax.set_xlabel('Epoch')
ax.set_ylabel('Gradient Norm')
ax.set_title('Gradient Dynamics (Smoothness)')
ax.grid(True, alpha=0.3)
# 5. Sample quality evolution
ax = axes[1, 1]
model.G.eval()
with torch.no_grad():
z = torch.randn(64, 100, 1, 1).to(next(model.G.parameters()).device)
samples = model.G(z).cpu().view(-1, 28, 28)
# Compute statistics
sample_mean = samples.mean(dim=0)
sample_std = samples.std(dim=0)
# Show mean image
ax.imshow(sample_mean.numpy(), cmap='gray')
ax.set_title('Mean Generated Sample\n(Should converge to data mean)')
ax.axis('off')
# 6. Augmentation effectiveness
ax = axes[1, 2]
if augment_type == 'diff':
# Visualize DiffAugment operations
sample_img = torch.randn(1, 1, 28, 28) * 0.5 + 0.5
aug_ops = ['original', 'brightness', 'translation', 'cutout']
aug_results = [sample_img.squeeze().numpy()]
diff_aug = DiffAugment()
aug_results.append(diff_aug._rand_brightness(sample_img).squeeze().numpy())
aug_results.append(diff_aug._rand_translation(sample_img).squeeze().numpy())
aug_results.append(diff_aug._rand_cutout(sample_img.clone()).squeeze().numpy())
for i, (op, img) in enumerate(zip(aug_ops, aug_results)):
plt.subplot(2, 2, i + 1)
plt.imshow(img, cmap='gray')
plt.title(op)
plt.axis('off')
plt.tight_layout()
plt.savefig(f'gan_stability_{augment_type}.png', dpi=300)
plt.show()
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--augment_type', type=str, default='ada',
choices=['none', 'diff', 'ada'])
parser.add_argument('--epochs', type=int, default=50)
args = parser.parse_args()
train_stable_gan(args.augment_type, args.epochs)13.3 标准化流与可逆网络
13.3.1 可逆神经网络（Invertible NN）的雅可比行列式计算与变量变换公式，RealNVP的耦合层（Coupling Layers）与棋盘掩码
标准化流通过可逆神经网络构建复杂分布与先验分布间的双射映射，变量变换公式要求计算网络雅可比行列式对数以修正概率测度。可逆架构设计确保雅可比矩阵具有三角形结构或分解形式，使行列式计算简化为对角元素乘积或对数行列式迹运算。RealNVP架构将输入维度划分为两半，一半直接通过，另一半通过以通过部分为条件的神经网络进行仿射变换，棋盘掩码与通道掩码交替实现所有维度的耦合交互，保持前向生成与反向推断的解析可逆性。
Python
"""
Script: realnvp_flow.py
Content: Invertible Neural Networks with Jacobian computation, RealNVP coupling layers,
checkerboard and channel-wise masking.
Usage: python realnvp_flow.py --num_coupling_layers 8 --dataset mnist --epochs 50
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import numpy as np
import matplotlib.pyplot as plt
from typing import Tuple, List, Dict
import argparse
class InvertibleLayer(nn.Module):
"""
Base class for invertible layers with log determinant Jacobian computation.
"""
def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Returns:
y: Transformed tensor
log_det: Log determinant of Jacobian
"""
raise NotImplementedError
def inverse(self, y: torch.Tensor) -> torch.Tensor:
"""Inverse transformation."""
raise NotImplementedError
class AffineCouplingLayer(InvertibleLayer):
"""
RealNVP-style affine coupling layer.
Splits input into two parts: x1 (unchanged) and x2 (transformed).
"""
def __init__(self, num_features: int, hidden_dim: int = 256, mask_type: str = 'checkerboard'):
super().__init__()
self.num_features = num_features
self.mask_type = mask_type
# Mask: 1 means pass-through, 0 means transform
if mask_type == 'checkerboard':
# For images, alternate pixels
self.register_buffer('mask', self._create_checkerboard_mask(28, 28))
elif mask_type == 'channel':
# Split by channel (for multi-channel images)
self.register_buffer('mask', torch.cat([
torch.ones(num_features // 2),
torch.zeros(num_features - num_features // 2)
]))
else:
raise ValueError(f"Unknown mask type: {mask_type}")
# Networks for scale and translation
# s(x1) and t(x1) where x1 is the masked (pass-through) part
input_dim = num_features // 2 if mask_type == 'channel' else num_features
self.scale_net = nn.Sequential(
nn.Linear(input_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, input_dim),
nn.Tanh() # Bounded scale for stability
)
self.translate_net = nn.Sequential(
nn.Linear(input_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, input_dim)
)
def _create_checkerboard_mask(self, h: int, w: int) -> torch.Tensor:
"""Create checkerboard mask for images."""
mask = torch.zeros(h, w)
for i in range(h):
for j in range(w):
mask[i, j] = (i + j) % 2
return mask.view(-1)
def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
batch_size = x.size(0)
# Apply mask
x_masked = x * self.mask
# Get x1 (pass-through) and prepare input for networks
if self.mask_type == 'checkerboard':
x1 = x_masked # Pass-through part
x1_flat = x_masked
else:
x1 = x_masked[:, :self.num_features // 2]
x1_flat = x1
# Compute scale and translation
s = self.scale_net(x1_flat)
t = self.translate_net(x1_flat)
# Prepare full-dimension scale and translation
if self.mask_type == 'checkerboard':
s_full = s * (1 - self.mask) # Apply to unmasked positions
t_full = t * (1 - self.mask)
else:
s_full = torch.cat([torch.zeros_like(x1), s], dim=1)
t_full = torch.cat([torch.zeros_like(x1), t], dim=1)
# Affine transformation: x2' = x2 * exp(s) + t
y = x_masked + (1 - self.mask) * (x * torch.exp(s_full) + t_full)
# Log determinant: sum of scale terms for transformed dimensions
log_det = torch.sum(s_full * (1 - self.mask), dim=1)
return y, log_det
def inverse(self, y: torch.Tensor) -> torch.Tensor:
"""Inverse: x2 = (y2 - t) * exp(-s)"""
y_masked = y * self.mask
if self.mask_type == 'checkerboard':
y1_flat = y_masked
else:
y1 = y_masked[:, :self.num_features // 2]
y1_flat = y1
s = self.scale_net(y1_flat)
t = self.translate_net(y1_flat)
if self.mask_type == 'checkerboard':
s_full = s * (1 - self.mask)
t_full = t * (1 - self.mask)
else:
zeros = torch.zeros_like(y[:, :self.num_features // 2])
s_full = torch.cat([zeros, s], dim=1)
t_full = torch.cat([zeros, t], dim=1)
x = y_masked + (1 - self.mask) * ((y - t_full) * torch.exp(-s_full))
return x
class RealNVP(nn.Module):
"""
Real-valued Non-Volume Preserving (RealNVP) flow.
Composes multiple affine coupling layers with alternating masks.
"""
def __init__(self, num_features: int = 784, num_coupling_layers: int = 8,
hidden_dim: int = 256):
super().__init__()
self.num_features = num_features
self.num_coupling_layers = num_coupling_layers
self.prior = torch.distributions.Normal(0, 1)
# Create coupling layers with alternating masks
self.coupling_layers = nn.ModuleList()
for i in range(num_coupling_layers):
if i % 2 == 0:
mask_type = 'checkerboard' if num_features == 784 else 'channel'
else:
# Invert mask
mask_type = 'checkerboard_inverted' if num_features == 784 else 'channel_inverted'
layer = AffineCouplingLayer(num_features, hidden_dim, 'checkerboard')
self.coupling_layers.append(layer)
# For inverted, we just flip the mask in the layer
if 'inverted' in mask_type:
layer.mask = 1 - layer.mask
def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
"""Forward pass (encoding/inference): x -> z."""
log_det_sum = 0
z = x
for layer in self.coupling_layers:
z, log_det = layer(z)
log_det_sum += log_det
return z, log_det_sum
def inverse(self, z: torch.Tensor) -> torch.Tensor:
"""Inverse pass (decoding/generation): z -> x."""
x = z
for layer in reversed(self.coupling_layers):
x = layer.inverse(x)
return x
def log_prob(self, x: torch.Tensor) -> torch.Tensor:
"""Compute log probability log p(x)."""
z, log_det = self.forward(x)
# Log probability under prior (standard normal)
log_pz = self.prior.log_prob(z).sum(dim=1)
# Change of variables formula: log p(x) = log p(z) + log |det J|
log_px = log_pz + log_det
return log_px
def sample(self, num_samples: int, device: str = 'cpu') -> torch.Tensor:
"""Generate samples from the model."""
z = self.prior.sample((num_samples, self.num_features)).to(device)
x = self.inverse(z)
return x
def train_realnvp(dataset: str = 'mnist', num_coupling_layers: int = 8,
epochs: int = 50):
"""
Training RealNVP flow via maximum likelihood.
"""
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Lambda(lambda x: x.view(-1) * 2 - 1) # Flatten and scale to [-1, 1]
])
train_data = datasets.MNIST('./data', train=True, download=True, transform=transform)
train_loader = DataLoader(train_data, batch_size=128, shuffle=True)
model = RealNVP(num_features=784, num_coupling_layers=num_coupling_layers).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4, weight_decay=1e-5)
history = {'loss': [], 'log_prob': [], 'jacobian_mean': []}
for epoch in range(epochs):
epoch_loss = 0
epoch_log_prob = 0
for batch_idx, (data, _) in enumerate(train_loader):
data = data.to(device)
optimizer.zero_grad()
# Compute negative log likelihood
log_px = model.log_prob(data)
loss = -log_px.mean()
loss.backward()
# Gradient clipping for stability
torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)
optimizer.step()
epoch_loss += loss.item()
epoch_log_prob += log_px.mean().item()
avg_loss = epoch_loss / len(train_loader)
avg_log_prob = epoch_log_prob / len(train_loader)
history['loss'].append(avg_loss)
history['log_prob'].append(avg_log_prob)
if epoch % 5 == 0:
print(f"Epoch {epoch}/{epochs} | Loss: {avg_loss:.4f} | "
f"Avg Log Prob: {avg_log_prob:.4f}")
visualize_realnvp(model, train_loader, device, history)
return model, history
def visualize_realnvp(model: RealNVP, data_loader: DataLoader,
device: torch.device, history: Dict):
"""
Visualization for RealNVP including latent space analysis and invertibility check.
"""
fig, axes = plt.subplots(2, 3, figsize=(18, 12))
# 1. Training loss
ax = axes[0, 0]
epochs = range(len(history['loss']))
ax.plot(epochs, history['loss'], color='blue', linewidth=2)
ax.set_xlabel('Epoch')
ax.set_ylabel('Negative Log Likelihood')
ax.set_title('RealNVP Training Loss')
ax.grid(True, alpha=0.3)
# 2. Latent space distribution (should be standard normal)
ax = axes[0, 1]
model.eval()
with torch.no_grad():
data, _ = next(iter(data_loader))
data = data.to(device)
z, _ = model.forward(data)
z = z.cpu().numpy()
# Plot histogram of latent dimensions
ax.hist(z.flatten(), bins=50, density=True, alpha=0.7, color='green', label='Empirical')
# Plot standard normal for comparison
x_range = np.linspace(z.min(), z.max(), 100)
ax.plot(x_range, (1/np.sqrt(2*np.pi)) * np.exp(-0.5 * x_range**2),
'r-', linewidth=2, label='N(0,1)')
ax.set_title('Latent Space Distribution')
ax.legend()
ax.grid(True, alpha=0.3)
# 3. Jacobian determinant distribution
ax = axes[0, 2]
with torch.no_grad():
_, log_dets = model.forward(data)
log_dets = log_dets.cpu().numpy()
ax.hist(log_dets, bins=30, alpha=0.7, color='purple')
ax.set_xlabel('Log Determinant of Jacobian')
ax.set_ylabel('Count')
ax.set_title('Jacobian Determinant Distribution')
ax.grid(True, alpha=0.3)
# 4. Invertibility check (reconstruction error)
ax = axes[1, 0]
with torch.no_grad():
# Encode and decode
z, _ = model.forward(data)
x_recon = model.inverse(z)
# Compute reconstruction error
recon_error = torch.mean((data - x_recon)**2, dim=1).cpu().numpy()
ax.hist(recon_error, bins=30, alpha=0.7, color='orange')
ax.set_xlabel('MSE Reconstruction Error')
ax.set_ylabel('Count')
ax.set_title(f'Invertibility Check (Mean Error: {recon_error.mean():.2e})')
ax.set_yscale('log')
ax.grid(True, alpha=0.3)
# 5. Generated samples
ax = axes[1, 1]
with torch.no_grad():
samples = model.sample(64, device).cpu().view(-1, 28, 28)
# Grid
grid = torch.zeros(8*28, 8*28)
for i in range(8):
for j in range(8):
idx = i * 8 + j
grid[i*28:(i+1)*28, j*28:(j+1)*28] = samples[idx]
ax.imshow(grid.numpy(), cmap='gray')
ax.set_title('Generated Samples (Inverse Transform)')
ax.axis('off')
# 6. Latent interpolation
ax = axes[1, 2]
with torch.no_grad():
# Get two data points
x1 = data[0:1]
x2 = data[1:2]
# Encode
z1, _ = model.forward(x1)
z2, _ = model.forward(x2)
# Interpolate in latent space and decode
interpolations = []
for alpha in np.linspace(0, 1, 8):
z_interp = (1 - alpha) * z1 + alpha * z2
x_interp = model.inverse(z_interp)
interpolations.append(x_interp.cpu().view(28, 28))
interp_grid = np.concatenate(interpolations, axis=1)
ax.imshow(interp_grid.numpy(), cmap='gray')
ax.set_title('Latent Space Interpolation')
ax.axis('off')
plt.tight_layout()
plt.savefig('realnvp_analysis.png', dpi=300)
plt.show()
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--num_coupling_layers', type=int, default=8)
parser.add_argument('--epochs', type=int, default=50)
args = parser.parse_args()
train_realnvp(num_coupling_layers=args.num_coupling_layers, epochs=args.epochs)13.3.2 Glow的ActNorm初始化、1x1可逆卷积（LU分解参数化）与多尺度架构，Flow++的变分去量化（Variational Dequantization）
Glow架构引入可逆分步标准化流，通过actnorm层补偿批归一化在可逆网络中的不可行性，基于数据依赖的初始化缩放平移激活统计量至标准分布。可逆卷积层采用LU分解参数化或QR分解，以可训练参数替换随机正交矩阵，在保持行列式易于计算的同时实现通道间的信息混合。多尺度架构在特定深度将部分维度直接输出为潜变量，剩余维度继续深层变换，降低计算成本并构建层次化生成表示。
Flow++针对离散数据的连续流建模缺陷，引入变分去量化模块学习随机去量化映射，通过自动编码器式网络预测均匀噪声的复杂分布，填补离散值间的间隙使流模型适应整数数据。耦合层升级为样条流或神经网络自回归变换，增强表达能力以捕获复杂相关性。
Python
"""
Script: glow_flowpp.py
Content: Glow with ActNorm, invertible 1x1 convolution (LU decomposition),
multi-scale architecture, and Flow++ variational dequantization.
Usage: python glow_flowpp.py --model_type glow --dataset mnist --epochs 50
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import numpy as np
import matplotlib.pyplot as plt
from typing import Tuple, List, Dict
import argparse
import scipy.linalg
class ActNorm(InvertibleLayer):
"""
Activation Normalization (ActNorm) layer.
Data-dependent initialization to center and scale activations.
"""
def __init__(self, num_features: int):
super().__init__()
self.num_features = num_features
# Parameters (initialized during first forward pass)
self.register_buffer('initialized', torch.tensor(0, dtype=torch.uint8))
self.register_parameter('log_scale', nn.Parameter(torch.zeros(num_features)))
self.register_parameter('bias', nn.Parameter(torch.zeros(num_features)))
def initialize(self, x: torch.Tensor):
"""Data-dependent initialization."""
with torch.no_grad():
# Compute mean and std per channel/feature
mean = x.mean(dim=0)
std = x.std(dim=0) + 1e-6
# Initialize bias to negative mean (center data)
self.bias.data.copy_(-mean)
# Initialize log_scale to -log(std) (normalize to unit variance)
self.log_scale.data.copy_(-torch.log(std))
self.initialized.fill_(1)
def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
if not self.initialized:
self.initialize(x)
# Apply: y = (x + bias) * exp(log_scale)
y = (x + self.bias) * torch.exp(self.log_scale)
# Log determinant: sum of log_scale for each dimension
log_det = self.log_scale.sum() * x.size(0)
return y, log_det
def inverse(self, y: torch.Tensor) -> torch.Tensor:
return y * torch.exp(-self.log_scale) - self.bias
class InvertibleConv1x1(InvertibleLayer):
"""
Invertible 1x1 convolution using LU decomposition.
Glow-style: W = PL(U + diag(s)), where P is permutation, L lower unit triangular, U upper triangular.
"""
def __init__(self, num_channels: int):
super().__init__()
self.num_channels = num_channels
# Initialize weight as random rotation matrix
w_init = np.random.randn(num_channels, num_channels)
w_init = scipy.linalg.qr(w_init)[0].astype(np.float32) # Orthogonal initialization
# LU decomposition: W = PLU
self.register_buffer('P', torch.from_numpy(scipy.linalg.lu_factor(w_init)[0]))
self.register_parameter('L', nn.Parameter(torch.eye(num_channels)))
self.register_parameter('U', nn.Parameter(torch.from_numpy(scipy.linalg.lu_factor(w_init)[1])))
self.register_parameter('log_s', nn.Parameter(torch.zeros(num_channels)))
# Mask for lower and upper triangular
self.register_buffer('lower_mask', torch.tril(torch.ones(num_channels, num_channels), -1))
self.register_buffer('upper_mask', torch.triu(torch.ones(num_channels, num_channels), 1))
def get_weight(self) -> torch.Tensor:
"""Reconstruct weight matrix from LU factors."""
# L is lower triangular with ones on diagonal + learnable lower part
L = self.L * self.lower_mask + torch.eye(self.num_channels, device=self.L.device)
# U is upper triangular with learnable upper part + diagonal
U = self.U * self.upper_mask + torch.diag(torch.exp(self.log_s))
# W = P @ L @ U
W = self.P @ L @ U
return W
def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
# x: (batch, channels, height, width)
batch_size, c, h, w = x.size()
W = self.get_weight()
# Apply convolution (equivalent to matrix multiplication per spatial location)
W_reshape = W.view(c, c, 1, 1)
y = F.conv2d(x, W_reshape)
# Log determinant: h * w * sum(log_s)
log_det = h * w * self.log_s.sum()
log_det = log_det.expand(batch_size)
return y, log_det
def inverse(self, y: torch.Tensor) -> torch.Tensor:
W = self.get_weight()
W_inv = torch.inverse(W)
W_inv_reshape = W_inv.view(self.num_channels, self.num_channels, 1, 1)
# Padding to maintain size (assuming proper spatial dims)
x = F.conv2d(y, W_inv_reshape, padding=0)
return x
class GlowStep(nn.Module):
"""
Single step of Glow: ActNorm -> Invertible Conv -> Coupling.
"""
def __init__(self, num_channels: int, hidden_dim: int = 256):
super().__init__()
self.actnorm = ActNorm(num_channels)
self.conv1x1 = InvertibleConv1x1(num_channels)
# Coupling layer (simplified affine coupling for images)
# Split channels in half
self.coupling = nn.ModuleDict({
'scale_net': nn.Sequential(
nn.Conv2d(num_channels // 2, hidden_dim, 3, padding=1),
nn.ReLU(),
nn.Conv2d(hidden_dim, hidden_dim, 3, padding=1),
nn.ReLU(),
nn.Conv2d(hidden_dim, num_channels // 2, 3, padding=1),
nn.Tanh()
),
'translate_net': nn.Sequential(
nn.Conv2d(num_channels // 2, hidden_dim, 3, padding=1),
nn.ReLU(),
nn.Conv2d(hidden_dim, hidden_dim, 3, padding=1),
nn.ReLU(),
nn.Conv2d(hidden_dim, num_channels // 2, 3, padding=1)
)
})
def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
log_det_sum = 0
# ActNorm
x, log_det = self.actnorm(x)
log_det_sum += log_det
# 1x1 Conv
x, log_det = self.conv1x1(x)
log_det_sum += log_det
# Coupling (split channel-wise)
x1, x2 = torch.chunk(x, 2, dim=1)
s = self.coupling['scale_net'](x1)
t = self.coupling['translate_net'](x1)
x2 = x2 * torch.exp(s) + t
y = torch.cat([x1, x2], dim=1)
log_det_sum += s.sum(dim=[1, 2, 3])
return y, log_det_sum
def inverse(self, y: torch.Tensor) -> torch.Tensor:
# Reverse coupling
y1, y2 = torch.chunk(y, 2, dim=1)
s = self.coupling['scale_net'](y1)
t = self.coupling['translate_net'](y1)
y2 = (y2 - t) * torch.exp(-s)
x = torch.cat([y1, y2], dim=1)
# Reverse 1x1 Conv
x = self.conv1x1.inverse(x)
# Reverse ActNorm
x = self.actnorm.inverse(x)
return x
class Glow(nn.Module):
"""
Glow model with multi-scale architecture.
Splits off half the dimensions at each scale.
"""
def __init__(self, num_channels: int = 1, num_levels: int = 3,
num_steps_per_level: int = 4, image_size: int = 28):
super().__init__()
self.num_levels = num_levels
self.num_steps_per_level = num_steps_per_level
self.image_size = image_size
self.prior = torch.distributions.Normal(0, 1)
# Build multi-scale architecture
self.flows = nn.ModuleList()
for level in range(num_levels):
# At each level, we process and then split
steps = nn.ModuleList()
for step in range(num_steps_per_level):
steps.append(GlowStep(num_channels))
self.flows.append(steps)
# After each level (except last), we split half the channels
if level < num_levels - 1:
num_channels *= 2 # Squeeze increases channels
image_size //= 2
def squeeze(self, x: torch.Tensor) -> torch.Tensor:
"""
Squeeze operation: (B, C, H, W) -> (B, 4*C, H/2, W/2)
Reshapes spatial dimensions into channels.
"""
batch_size, c, h, w = x.size()
x = x.view(batch_size, c, h // 2, 2, w // 2, 2)
x = x.permute(0, 1, 3, 5, 2, 4).contiguous()
x = x.view(batch_size, c * 4, h // 2, w // 2)
return x
def unsqueeze(self, x: torch.Tensor) -> torch.Tensor:
"""Reverse of squeeze."""
batch_size, c, h, w = x.size()
x = x.view(batch_size, c // 4, 2, 2, h, w)
x = x.permute(0, 1, 4, 2, 5, 3).contiguous()
x = x.view(batch_size, c // 4, h * 2, w * 2)
return x
def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, List[torch.Tensor], torch.Tensor]:
"""
Forward pass returning latent variables at each scale and total log determinant.
"""
z_list = []
log_det_sum = 0
current = x
for level, steps in enumerate(self.flows):
# Squeeze at beginning of level (except first)
if level > 0:
current = self.squeeze(current)
# Apply flow steps
for step in steps:
current, log_det = step(current)
log_det_sum += log_det
# Split (except last level)
if level < self.num_levels - 1:
current, z = torch.chunk(current, 2, dim=1)
z_list.append(z)
# Last level output
z_list.append(current)
return torch.cat([z.flatten(1) for z in z_list], dim=1), z_list, log_det_sum
def log_prob(self, x: torch.Tensor) -> torch.Tensor:
z, _, log_det = self.forward(x)
log_pz = self.prior.log_prob(z).sum(dim=1)
return log_pz + log_det
def sample(self, num_samples: int, device: str = 'cpu') -> torch.Tensor:
# Generate random z
# Need to account for multi-scale dimensions
total_dims = 1 * 28 * 28 # Simplified for single-scale MNIST
z = self.prior.sample((num_samples, total_dims)).to(device)
# Reverse through flows (simplified - assumes single scale for demo)
# Full implementation would track splits and process in reverse
z_spatial = z.view(num_samples, 1, 28, 28)
# Reverse flow
x = z_spatial
for level in reversed(self.flows):
for step in reversed(level):
x = step.inverse(x)
return x
class VariationalDequantization(nn.Module):
"""
Flow++ style variational dequantization.
Learns to add noise to discrete data for continuous flow modeling.
"""
def __init__(self, num_channels: int = 1, num_bits: int = 8):
super().__init__()
self.num_bits = num_bits
self.num_bins = 2 ** num_bits
# Encoder for dequantization noise u
self.encoder = nn.Sequential(
nn.Conv2d(num_channels, 64, 3, padding=1),
nn.ReLU(),
nn.Conv2d(64, 64, 3, padding=1),
nn.ReLU()
)
# Predict distribution parameters for u
self.mu = nn.Conv2d(64, num_channels, 3, padding=1)
self.log_sigma = nn.Conv2d(64, num_channels, 3, padding=1)
def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Dequantize: y = x + u, where u is learned stochastic noise.
"""
# x is discrete in [0, num_bins-1], normalized to [0, 1]
# We want to map to (x, x+1) uniformly in continuous space
# Predict dequantization noise distribution
h = self.encoder(x)
mu = self.mu(h)
log_sigma = self.log_sigma(h)
sigma = torch.exp(log_sigma)
# Sample noise u ~ logistic(mu, sigma) or gaussian
u = mu + sigma * torch.randn_like(x)
# Dequantized value
y = x + u # Now in continuous space
# Log probability of dequantization (for variational bound)
log_qu = -0.5 * ((u - mu) / sigma).pow(2) - log_sigma - 0.5 * np.log(2 * np.pi)
log_qu = log_qu.sum(dim=[1, 2, 3])
# Change of variables: add log(1) = 0 since dy/du = 1
# But we need to account for the discrete-to-continuous mapping
return y, log_qu
def inverse(self, y: torch.Tensor) -> torch.Tensor:
# Round to nearest integer for discrete data
return torch.clamp(torch.round(y), 0, self.num_bins - 1) / self.num_bins
class FlowPlusPlus(nn.Module):
"""
Flow++ combining variational dequantization with improved coupling layers.
"""
def __init__(self, num_channels: int = 1, num_bits: int = 8):
super().__init__()
self.dequantizer = VariationalDequantization(num_channels, num_bits)
# Could combine with Glow steps or RealNVP
self.flow = RealNVP(num_features=784, num_coupling_layers=8) # Reuse RealNVP
def log_prob(self, x: torch.Tensor) -> torch.Tensor:
# Dequantize discrete data
y, log_qu = self.dequantizer(x)
# Apply flow on continuous data
y_flat = y.view(y.size(0), -1) * 2 - 1 # Scale to [-1, 1]
log_py = self.flow.log_prob(y_flat)
# Variational bound: log p(x) >= log p(y) + log q(u|x)
# (Simplified - actual bound includes more terms)
return log_py + log_qu
def sample(self, num_samples: int, device: str = 'cpu') -> torch.Tensor:
z = self.flow.sample(num_samples, device)
# Scale back to [0, 1] then add dequantization
y = (z + 1) / 2
x = self.dequantizer.inverse(y)
return x
def train_glow(dataset: str = 'mnist', epochs: int = 50):
"""
Training pipeline for Glow with multi-scale architecture.
"""
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
transform = transforms.Compose([
transforms.ToTensor(),
# No flattening - keep spatial structure for convolutions
])
train_data = datasets.MNIST('./data', train=True, download=True, transform=transform)
train_loader = DataLoader(train_data, batch_size=64, shuffle=True)
model = Glow(num_channels=1, num_levels=2, num_steps_per_level=3).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4, weight_decay=1e-5)
history = {'loss': [], 'log_prob': []}
for epoch in range(epochs):
epoch_loss = 0
for batch_idx, (data, _) in enumerate(train_loader):
data = data.to(device)
optimizer.zero_grad()
log_px = model.log_prob(data)
loss = -log_px.mean()
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 10.0)
optimizer.step()
epoch_loss += loss.item()
avg_loss = epoch_loss / len(train_loader)
history['loss'].append(avg_loss)
if epoch % 5 == 0:
print(f"Epoch {epoch}/{epochs} | Loss: {avg_loss:.4f}")
visualize_glow(model, train_loader, device, history)
return model, history
def visualize_glow(model: Glow, data_loader: DataLoader, device: torch.device, history: Dict):
"""
Visualization for Glow including multi-scale latent analysis.
"""
fig, axes = plt.subplots(2, 3, figsize=(18, 12))
# 1. Training loss
ax = axes[0, 0]
ax.plot(history['loss'], color='blue', linewidth=2)
ax.set_xlabel('Epoch')
ax.set_ylabel('Negative Log Likelihood')
ax.set_title('Glow Training Loss')
ax.grid(True, alpha=0.3)
# 2. Latent space at different scales
ax = axes[0, 1]
model.eval()
with torch.no_grad():
data, _ = next(iter(data_loader))
data = data.to(device)
z_combined, z_list, _ = model.forward(data)
# Visualize distribution at each scale
colors = plt.cm.viridis(np.linspace(0, 1, len(z_list)))
for i, (z, color) in enumerate(zip(z_list, colors)):
z_flat = z.flatten(1).cpu().numpy().flatten()[:1000] # Sample
ax.hist(z_flat, bins=30, alpha=0.5, color=color, label=f'Scale {i+1}')
ax.set_title('Multi-Scale Latent Distributions')
ax.legend()
ax.grid(True, alpha=0.3)
# 3. ActNorm statistics
ax = axes[0, 2]
# Collect mean/std from first ActNorm layer
if hasattr(model.flows[0][0], 'actnorm'):
actnorm = model.flows[0][0].actnorm
stats = {
'bias': actnorm.bias.detach().cpu().numpy(),
'log_scale': actnorm.log_scale.detach().cpu().numpy()
}
x_pos = np.arange(len(stats['bias']))
ax.bar(x_pos, stats['bias'], alpha=0.7, label='Bias')
ax.set_xlabel('Channel Index')
ax.set_ylabel('Bias Value')
ax.set_title('ActNorm Data-Dependent Initialization')
ax.grid(True, alpha=0.3, axis='y')
# 4. Invertible Conv singular values
ax = axes[1, 0]
if hasattr(model.flows[0][0], 'conv1x1'):
conv = model.flows[0][0].conv1x1
s_values = torch.exp(conv.log_s).detach().cpu().numpy()
ax.bar(range(len(s_values)), s_values, color='purple', alpha=0.7)
ax.set_xlabel('Channel Index')
ax.set_ylabel('Singular Value')
ax.set_title('1x1 Invertible Conv Singular Values')
ax.grid(True, alpha=0.3, axis='y')
# 5. Generated samples
ax = axes[1, 1]
with torch.no_grad():
samples = model.sample(64, device).cpu()
grid = torch.zeros(8, 8, 28, 28)
for i in range(8):
for j in range(8):
idx = i * 8 + j
grid[i, j] = samples[idx].squeeze()
grid = grid.view(8*28, 8*28).numpy()
ax.imshow(grid, cmap='gray')
ax.set_title('Glow Generated Samples')
ax.axis('off')
# 6. Dequantization visualization (if using Flow++)
ax = axes[1, 2]
# Simulate dequantization effect
discrete_vals = torch.linspace(0, 1, 256)
continuous_approx = discrete_vals.unsqueeze(1) + torch.randn(256, 100) * 0.01
for i in range(0, 256, 50):
ax.scatter([i/255]*100, continuous_approx[i].numpy(), alpha=0.3, s=10)
ax.set_xlabel('Discrete Value')
ax.set_ylabel('Continuous Dequantized')
ax.set_title('Variational Dequantization Concept')
ax.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig('glow_analysis.png', dpi=300)
plt.show()
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--model_type', type=str, default='glow',
choices=['glow', 'flowpp'])
parser.add_argument('--epochs', type=int, default=50)
args = parser.parse_args()
train_glow(epochs=args.epochs)13.4 扩散模型与流匹配
13.4.1 DDPM的前向加噪过程（q(xt∣x0) ）与反向去噪（pθ(xt−1∣xt) ）的方差调度（Variance Schedule），去噪分数匹配（Denoising Score Matching）与证据下界（ELBO）
去噪扩散概率模型通过马尔可夫链渐进破坏数据结构，前向过程在有限时间步内向数据注入高斯噪声，方差调度控制扰动速率，确保终止状态逼近易处理先验分布。反向去噪过程学习恢复数据结构，神经网络预测噪声或重构目标实现转移概率参数化。去噪分数匹配目标等价于证据下界优化，去噪器隐式估计数据分布分数函数，通过加权均方误差训练实现似然最大化与样本生成能力的统一。
Python
"""
Script: ddpm_denoising_diffusion.py
Content: DDPM implementation with forward noising, reverse denoising,
variance scheduling (linear/cosine), denoising score matching, and ELBO.
Usage: python ddpm_denoising_diffusion.py --schedule cosine --timesteps 1000 --epochs 50
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import numpy as np
import matplotlib.pyplot as plt
from typing import Tuple, Dict, Optional
import argparse
import math
class VarianceScheduler:
"""
Variance schedule for diffusion process.
Supports linear and cosine scheduling.
"""
def __init__(self, num_timesteps: int = 1000, schedule_type: str = 'linear',
beta_start: float = 1e-4, beta_end: float = 0.02):
self.num_timesteps = num_timesteps
self.schedule_type = schedule_type
if schedule_type == 'linear':
# Linear schedule from beta_start to beta_end
self.betas = torch.linspace(beta_start, beta_end, num_timesteps)
elif schedule_type == 'cosine':
# Cosine schedule (improved)
self.betas = self._cosine_beta_schedule(num_timesteps)
else:
raise ValueError(f"Unknown schedule: {schedule_type}")
self.alphas = 1.0 - self.betas
self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
self.alphas_cumprod_prev = F.pad(self.alphas_cumprod[:-1], (1, 0), value=1.0)
# Calculations for diffusion q(x_t | x_0)
self.sqrt_alphas_cumprod = torch.sqrt(self.alphas_cumprod)
self.sqrt_one_minus_alphas_cumprod = torch.sqrt(1.0 - self.alphas_cumprod)
# Calculations for posterior q(x_{t-1} | x_t, x_0)
self.posterior_variance = (
self.betas * (1.0 - self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
)
self.sqrt_recip_alphas_cumprod = torch.sqrt(1.0 / self.alphas_cumprod)
self.sqrt_recipm1_alphas_cumprod = torch.sqrt(1.0 / self.alphas_cumprod - 1)
# Sampling parameters
self.posterior_mean_coef1 = (
self.betas * torch.sqrt(self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
)
self.posterior_mean_coef2 = (
(1.0 - self.alphas_cumprod_prev) * torch.sqrt(self.alphas) / (1.0 - self.alphas_cumprod)
)
def _cosine_beta_schedule(self, timesteps: int, s: float = 0.008) -> torch.Tensor:
"""
Cosine schedule as proposed in Improved DDPM.
"""
steps = timesteps + 1
x = torch.linspace(0, timesteps, steps)
alphas_cumprod = torch.cos(((x / timesteps) + s) / (1 + s) * math.pi * 0.5) ** 2
alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
return torch.clip(betas, 0.0001, 0.9999)
def q_sample(self, x_0: torch.Tensor, t: torch.Tensor, noise: Optional[torch.Tensor] = None) -> torch.Tensor:
"""
Forward diffusion process: q(x_t | x_0).
"""
if noise is None:
noise = torch.randn_like(x_0)
sqrt_alphas_cumprod_t = self.sqrt_alphas_cumprod[t].view(-1, 1, 1, 1)
sqrt_one_minus_alphas_cumprod_t = self.sqrt_one_minus_alphas_cumprod[t].view(-1, 1, 1, 1)
return sqrt_alphas_cumprod_t * x_0 + sqrt_one_minus_alphas_cumprod_t * noise
def predict_start_from_noise(self, x_t: torch.Tensor, t: torch.Tensor, noise: torch.Tensor) -> torch.Tensor:
"""
Predict x_0 from x_t and predicted noise.
"""
sqrt_recip_alphas_cumprod_t = self.sqrt_recip_alphas_cumprod[t].view(-1, 1, 1, 1)
sqrt_recipm1_alphas_cumprod_t = self.sqrt_recipm1_alphas_cumprod[t].view(-1, 1, 1, 1)
return sqrt_recip_alphas_cumprod_t * x_t - sqrt_recipm1_alphas_cumprod_t * noise
def q_posterior_mean_variance(self, x_0: torch.Tensor, x_t: torch.Tensor, t: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Compute posterior mean and variance: q(x_{t-1} | x_t, x_0).
"""
posterior_mean = (
self.posterior_mean_coef1[t].view(-1, 1, 1, 1) * x_0 +
self.posterior_mean_coef2[t].view(-1, 1, 1, 1) * x_t
)
posterior_variance = self.posterior_variance[t].view(-1, 1, 1, 1)
return posterior_mean, posterior_variance
def p_sample(self, model: nn.Module, x_t: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
"""
Reverse process sampling: p(x_{t-1} | x_t).
"""
# Predict noise
predicted_noise = model(x_t, t)
# Predict x_0
pred_x_0 = self.predict_start_from_noise(x_t, t, predicted_noise)
# Compute posterior mean and variance
model_mean, model_variance = self.q_posterior_mean_variance(pred_x_0, x_t, t)
noise = torch.randn_like(x_t) if t[0] > 0 else torch.zeros_like(x_t)
return model_mean + torch.sqrt(model_variance) * noise
class UNet(nn.Module):
"""
Simple U-Net for noise prediction in diffusion models.
Time embedding via sinusoidal positional encoding.
"""
def __init__(self, in_channels: int = 1, time_emb_dim: int = 256):
super().__init__()
# Time embedding MLP
self.time_mlp = nn.Sequential(
nn.Linear(time_emb_dim, time_emb_dim),
nn.SiLU(),
nn.Linear(time_emb_dim, time_emb_dim)
)
# Encoder
self.enc1 = self._make_block(in_channels, 64)
self.enc2 = self._make_block(64, 128)
self.enc3 = self._make_block(128, 256)
# Bottleneck
self.bottleneck = self._make_block(256, 256)
# Decoder
self.dec3 = self._make_block(256 + 256, 128)
self.dec2 = self._make_block(128 + 128, 64)
self.dec1 = self._make_block(64 + 64, 64)
self.out = nn.Conv2d(64, in_channels, 1)
# Pooling and upsampling
self.pool = nn.MaxPool2d(2)
self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
self.time_emb_dim = time_emb_dim
def _make_block(self, in_ch: int, out_ch: int) -> nn.Module:
return nn.Sequential(
nn.Conv2d(in_ch, out_ch, 3, padding=1),
nn.GroupNorm(8, out_ch),
nn.SiLU(),
nn.Conv2d(out_ch, out_ch, 3, padding=1),
nn.GroupNorm(8, out_ch),
nn.SiLU()
)
def get_time_embedding(self, t: torch.Tensor, channels: int) -> torch.Tensor:
"""
Sinusoidal positional embeddings for time steps.
"""
inv_freq = 1.0 / (10000 ** (torch.arange(0, channels, 2, device=t.device).float() / channels))
pos_enc_a = torch.sin(t.repeat(1, channels // 2) * inv_freq)
pos_enc_b = torch.cos(t.repeat(1, channels // 2) * inv_freq)
pos_enc = torch.cat([pos_enc_a, pos_enc_b], dim=-1)
return pos_enc
def forward(self, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
# Time embedding
t_emb = self.get_time_embedding(t.float().view(-1, 1) / 1000.0, self.time_emb_dim)
t_emb = self.time_mlp(t_emb)
# Encoder with skip connections
e1 = self.enc1(x)
e2 = self.enc2(self.pool(e1))
e3 = self.enc3(self.pool(e2))
# Bottleneck (inject time embedding here for simplicity)
b = self.bottleneck(self.pool(e3))
# Decoder with skip connections
d3 = self.dec3(torch.cat([self.upsample(b), e3], dim=1))
d2 = self.dec2(torch.cat([self.upsample(d3), e2], dim=1))
d1 = self.dec1(torch.cat([self.upsample(d2), e1], dim=1))
return self.out(d1)
class DDPM(nn.Module):
"""
Denoising Diffusion Probabilistic Model (Ho et al., 2020).
"""
def __init__(self, num_timesteps: int = 1000, schedule_type: str = 'linear'):
super().__init__()
self.num_timesteps = num_timesteps
self.scheduler = VarianceScheduler(num_timesteps, schedule_type)
self.model = UNet()
# Register buffers to device
for key, value in self.scheduler.__dict__.items():
if isinstance(value, torch.Tensor):
self.register_buffer(f'scheduler_{key}', value)
def forward(self, x_0: torch.Tensor, t: torch.Tensor, noise: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Forward pass for training: predict noise.
"""
if noise is None:
noise = torch.randn_like(x_0)
x_t = self.scheduler.q_sample(x_0, t, noise)
predicted_noise = self.model(x_t, t)
# Denoising score matching loss (simplified ELBO)
loss = F.mse_loss(predicted_noise, noise)
return loss, predicted_noise
def sample(self, batch_size: int = 64, channels: int = 1, size: int = 28,
device: str = 'cpu') -> torch.Tensor:
"""
Generate samples via reverse process.
"""
self.model.eval()
shape = (batch_size, channels, size, size)
# Start from pure noise
x_t = torch.randn(shape, device=device)
# Reverse diffusion
for t in reversed(range(self.num_timesteps)):
t_batch = torch.full((batch_size,), t, device=device, dtype=torch.long)
x_t = self.scheduler.p_sample(self.model, x_t, t_batch)
return x_t
def train_ddpm(schedule_type: str = 'cosine', num_timesteps: int = 1000,
epochs: int = 50):
"""
Training pipeline for DDPM with denoising score matching.
"""
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,)) # Scale to [-1, 1]
])
train_data = datasets.MNIST('./data', train=True, download=True, transform=transform)
train_loader = DataLoader(train_data, batch_size=128, shuffle=True)
model = DDPM(num_timesteps=num_timesteps, schedule_type=schedule_type).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=2e-4)
history = {'loss': [], 'noise_magnitude': []}
for epoch in range(epochs):
epoch_loss = 0
for batch_idx, (data, _) in enumerate(train_loader):
data = data.to(device)
# Sample random timesteps
t = torch.randint(0, num_timesteps, (data.size(0),), device=device)
optimizer.zero_grad()
loss, pred_noise = model(data, t)
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
optimizer.step()
epoch_loss += loss.item()
history['noise_magnitude'].append(pred_noise.abs().mean().item())
avg_loss = epoch_loss / len(train_loader)
history['loss'].append(avg_loss)
if epoch % 5 == 0:
print(f"Epoch {epoch}/{epochs} | Loss: {avg_loss:.4f}")
visualize_ddpm(model, train_loader, device, history)
return model, history
def visualize_ddpm(model: DDPM, data_loader: DataLoader, device: torch.device,
history: Dict):
"""
Visualization for DDPM including forward/reverse process and noise schedule.
"""
fig, axes = plt.subplots(2, 3, figsize=(18, 12))
# 1. Training loss
ax = axes[0, 0]
ax.plot(history['loss'], color='blue', linewidth=2)
ax.set_xlabel('Epoch')
ax.set_ylabel('MSE Loss')
ax.set_title('DDPM Training Loss (Denoising Score Matching)')
ax.grid(True, alpha=0.3)
# 2. Noise schedule visualization
ax = axes[0, 1]
timesteps = range(model.num_timesteps)
betas = model.scheduler.betas.cpu().numpy()
alphas_cumprod = model.scheduler.alphas_cumprod.cpu().numpy()
ax.plot(timesteps, betas, label='Beta (Variance)', color='red', alpha=0.7)
ax.plot(timesteps, alphas_cumprod, label='Alpha Bar (Signal)', color='blue', alpha=0.7)
ax.set_xlabel('Timestep')
ax.set_ylabel('Value')
ax.set_title(f'Variance Schedule ({model.scheduler.schedule_type})')
ax.legend()
ax.grid(True, alpha=0.3)
# 3. Forward diffusion process (noising)
ax = axes[0, 2]
data, _ = next(iter(data_loader))
x_0 = data[0:1].to(device)
num_show_steps = 10
show_steps = np.linspace(0, model.num_timesteps - 1, num_show_steps, dtype=int)
noisy_images = []
for t in show_steps:
t_tensor = torch.tensor([t], device=device)
x_t = model.scheduler.q_sample(x_0, t_tensor)
noisy_images.append(x_t.cpu().squeeze().numpy())
# Plot
for i, img in enumerate(noisy_images):
plt.subplot(2, num_show_steps, i + 1)
plt.imshow(img, cmap='gray')
plt.axis('off')
plt.title(f't={show_steps[i]}')
# 4. Reverse diffusion process (denoising)
ax = axes[1, 0]
with torch.no_grad():
samples = model.sample(8, device=device).cpu()
grid = torch.zeros(2, 4, 28, 28)
for i in range(8):
grid[i // 4, i % 4] = samples[i].squeeze()
grid = grid.view(2*28, 4*28).numpy()
ax.imshow(grid, cmap='gray')
ax.set_title('Generated Samples (Reverse Process)')
ax.axis('off')
# 5. Noise prediction accuracy across timesteps
ax = axes[1, 1]
with torch.no_grad():
data, _ = next(iter(data_loader))
data = data[:100].to(device)
mse_by_t = []
t_range = range(0, model.num_timesteps, 100)
for t_val in t_range:
t = torch.full((100,), t_val, device=device)
loss, _ = model(data, t)
mse_by_t.append(loss.item())
ax.plot(t_range, mse_by_t, marker='o', color='green')
ax.set_xlabel('Timestep')
ax.set_ylabel('MSE')
ax.set_title('Denoising Difficulty Across Timesteps')
ax.grid(True, alpha=0.3)
# 6. ELBO estimation
ax = axes[1, 2]
# Simplified ELBO visualization: reconstruction quality over training
if len(history['loss']) > 10:
# Use loss as proxy for negative ELBO
elbo_est = [-l for l in history['loss']]
ax.plot(elbo_est, color='purple', linewidth=2)
ax.set_xlabel('Epoch')
ax.set_ylabel('Estimated ELBO')
ax.set_title('Evidence Lower Bound (Proxy)')
ax.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig('ddpm_analysis.png', dpi=300)
plt.show()
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--schedule', type=str, default='cosine',
choices=['linear', 'cosine'])
parser.add_argument('--timesteps', type=int, default=1000)
parser.add_argument('--epochs', type=int, default=50)
args = parser.parse_args()
train_ddpm(args.schedule, args.timesteps, args.epochs)13.4.2 连续时间扩散：随机微分方程（SDE）视角（VP/VE/Sub-VP）与概率流ODE，概率密度变换的Fokker-Planck方程
连续时间极限下扩散过程转化为随机微分方程，漂移系数与扩散系数分别编码确定性趋势与随机扰动强度。方差保持、方差爆炸与亚方差保持设计对应不同噪声调度策略，控制扰动-去噪动态的行为特性。概率流常微分方程去除扩散项仅保留漂移，通过确定性轨迹连接先验与数据分布，实现精确的似然计算与加速采样。Fokker-Planck方程描述概率密度随时间演化，与反向时间SDE共同刻画生成过程的统计动力学，确保概率质量守恒与流平衡。
Python
"""
Script: continuous_time_diffusion_sde.py
Content: Continuous-time diffusion with SDE framework (VP/VE/Sub-VP),
Probability Flow ODE, and Fokker-Planck equation analysis.
Usage: python continuous_time_diffusion_sde.py --sde_type ve --num_steps 50
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import numpy as np
import matplotlib.pyplot as plt
from typing import Tuple, Callable, Dict
import argparse
from scipy.integrate import solve_ivp
class ContinuousSDE:
"""
Continuous-time Stochastic Differential Equation framework for diffusion models.
Implements VP (Variance Preserving), VE (Variance Exploding), and Sub-VP.
"""
def __init__(self, sde_type: str = 've', beta_min: float = 0.1,
beta_max: float = 20.0):
self.sde_type = sde_type
self.beta_min = beta_min
self.beta_max = beta_max
if sde_type == 'vp':
# Variance Preserving (Song et al. 2021)
self.drift = lambda x, t: -0.5 * self.beta(t) * x
self.diffusion = lambda x, t: torch.sqrt(torch.tensor(self.beta(t)))
self.sigma = lambda t: torch.sqrt(1 - torch.exp(-self.beta_integral(t)))
elif sde_type == 've':
# Variance Exploding
self.sigma = lambda t: self.sigma_min * (self.sigma_max / self.sigma_min) ** t
self.drift = lambda x, t: torch.zeros_like(x)
self.diffusion = lambda x, t: torch.sqrt(
2 * (self.sigma(t) ** 2) * torch.log(torch.tensor(self.sigma_max / self.sigma_min))
)
self.sigma_min = 0.01
self.sigma_max = 50.0
elif sde_type == 'subvp':
# Sub-Variance Preserving
self.drift = lambda x, t: -0.5 * self.beta(t) * x
self.diffusion = lambda x, t: torch.sqrt(
self.beta(t) * (1 - torch.exp(-2 * self.beta_integral(t)))
)
self.sigma = lambda t: 1 - torch.exp(-self.beta_integral(t))
else:
raise ValueError(f"Unknown SDE type: {sde_type}")
def beta(self, t: float) -> float:
"""Linear beta schedule."""
return self.beta_min + t * (self.beta_max - self.beta_min)
def beta_integral(self, t: float) -> float:
"""Integral of beta from 0 to t."""
return self.beta_min * t + 0.5 * (self.beta_max - self.beta_min) * t ** 2
def marginal_prob(self, x0: torch.Tensor, t: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Mean and std of marginal probability p(x_t | x_0).
"""
if self.sde_type == 'vp':
mean = torch.exp(-0.5 * self.beta_integral(t)) * x0
std = torch.sqrt(1 - torch.exp(-self.beta_integral(t)))
elif self.sde_type == 've':
mean = x0
std = self.sigma(t)
elif self.sde_type == 'subvp':
mean = torch.exp(-0.5 * self.beta_integral(t)) * x0
std = 1 - torch.exp(-self.beta_integral(t))
return mean, std
def reverse_sde(self, score_fn: Callable, x: torch.Tensor, t: float, dt: float) -> torch.Tensor:
"""
Reverse-time SDE sampling: dx = [f(x,t) - g(t)^2 * score(x,t)] dt + g(t) dw
"""
drift = self.drift(x, t)
diffusion = self.diffusion(x, t)
score = score_fn(x, t)
# Reverse drift
rev_drift = drift - (diffusion ** 2) * score
# Euler-Maruyama step
noise = torch.randn_like(x) * np.sqrt(abs(dt))
x_next = x + rev_drift * dt + diffusion * noise
return x_next
def probability_flow_ode(self, score_fn: Callable, x: torch.Tensor,
t: float, dt: float) -> torch.Tensor:
"""
Probability flow ODE (deterministic): dx = [f(x,t) - 0.5 * g(t)^2 * score(x,t)] dt
"""
drift = self.drift(x, t)
diffusion = self.diffusion(x, t)
score = score_fn(x, t)
# ODE drift (no diffusion term)
ode_drift = drift - 0.5 * (diffusion ** 2) * score
return x + ode_drift * dt
class ScoreNetwork(nn.Module):
"""
Neural network for score function estimation ∇_x log p(x).
"""
def __init__(self, time_embed_dim: int = 128):
super().__init__()
self.time_embed_dim = time_embed_dim
# Time embedding
self.time_mlp = nn.Sequential(
nn.Linear(time_embed_dim, time_embed_dim),
nn.SiLU(),
nn.Linear(time_embed_dim, time_embed_dim)
)
# Simple CNN for score
self.net = nn.Sequential(
nn.Conv2d(1, 64, 3, padding=1),
nn.SiLU(),
nn.Conv2d(64, 128, 3, padding=1),
nn.SiLU(),
nn.Conv2d(128, 128, 3, padding=1),
nn.SiLU(),
nn.Conv2d(128, 64, 3, padding=1),
nn.SiLU(),
nn.Conv2d(64, 1, 3, padding=1)
)
def get_timestep_embedding(self, timesteps: torch.Tensor, embedding_dim: int) -> torch.Tensor:
"""
Sinusoidal timestep embeddings.
"""
half_dim = embedding_dim // 2
emb = math.log(10000) / (half_dim - 1)
emb = torch.exp(torch.arange(half_dim, device=timesteps.device) * -emb)
emb = timesteps[:, None] * emb[None, :]
emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
return emb
def forward(self, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
# Time embedding could be injected here (simplified for demo)
return self.net(x)
class ContinuousDiffusionTrainer:
"""
Trainer for continuous-time diffusion models via score matching.
"""
def __init__(self, sde: ContinuousSDE, score_net: ScoreNetwork, device: str = 'cuda'):
self.sde = sde
self.score_net = score_net.to(device)
self.device = device
# Weighting function for loss (denoising score matching)
if sde.sde_type == 've':
self.lambda_weight = lambda t: (sde.sigma(t) ** 2)
else:
self.lambda_weight = lambda t: 1.0
def score_matching_loss(self, x0: torch.Tensor) -> torch.Tensor:
"""
Denoising score matching loss for continuous time.
"""
# Sample random time
t = torch.rand(x0.size(0), device=self.device) * (0.999 - 0.001) + 0.001
# Sample from marginal distribution
mean, std = self.sde.marginal_prob(x0, t)
noise = torch.randn_like(x0)
xt = mean + std.view(-1, 1, 1, 1) * noise
# Predict score
predicted_score = self.score_net(xt, t)
# Target score is -noise / std (for Gaussian perturbation kernel)
target_score = -noise / std.view(-1, 1, 1, 1)
# Weighted loss
weight = self.lambda_weight(t).view(-1, 1, 1, 1)
loss = torch.mean(weight * ((predicted_score - target_score) ** 2))
return loss
def sample_reverse_sde(self, shape: Tuple, num_steps: int = 1000) -> torch.Tensor:
"""
Sample using reverse-time SDE.
"""
self.score_net.eval()
x = torch.randn(shape, device=self.device)
timesteps = torch.linspace(1, 0, num_steps + 1)
dt = -1.0 / num_steps
with torch.no_grad():
for i in range(num_steps):
t = timesteps[i]
t_batch = torch.full((shape[0],), t, device=self.device)
def score_fn(x, t):
t_tensor = torch.full((x.size(0),), t, device=self.device)
return self.score_net(x, t_tensor)
x = self.sde.reverse_sde(score_fn, x, t.item(), dt)
return x
def sample_probability_flow_ode(self, shape: Tuple, num_steps: int = 100) -> torch.Tensor:
"""
Sample using deterministic probability flow ODE (fewer steps possible).
"""
self.score_net.eval()
x = torch.randn(shape, device=self.device)
timesteps = torch.linspace(1, 0, num_steps + 1)
dt = -1.0 / num_steps
with torch.no_grad():
for i in range(num_steps):
t = timesteps[i]
def score_fn(x, t_val):
t_tensor = torch.full((x.size(0),), t_val, device=self.device)
return self.score_net(x, t_tensor)
x = self.sde.probability_flow_ode(score_fn, x, t.item(), dt)
return x
def train_continuous_diffusion(sde_type: str = 've', epochs: int = 50):
"""
Training pipeline for continuous-time diffusion.
"""
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,))
])
train_data = datasets.MNIST('./data', train=True, download=True, transform=transform)
train_loader = DataLoader(train_data, batch_size=128, shuffle=True)
sde = ContinuousSDE(sde_type=sde_type)
score_net = ScoreNetwork()
trainer = ContinuousDiffusionTrainer(sde, score_net, device)
optimizer = torch.optim.Adam(score_net.parameters(), lr=1e-4)
history = {'loss': []}
for epoch in range(epochs):
epoch_loss = 0
for batch_idx, (data, _) in enumerate(train_loader):
data = data.to(device)
optimizer.zero_grad()
loss = trainer.score_matching_loss(data)
loss.backward()
optimizer.step()
epoch_loss += loss.item()
avg_loss = epoch_loss / len(train_loader)
history['loss'].append(avg_loss)
if epoch % 5 == 0:
print(f"Epoch {epoch}/{epochs} | Score Matching Loss: {avg_loss:.4f}")
# Sampling comparison
print("\nGenerating samples with reverse SDE...")
samples_sde = trainer.sample_reverse_sde((16, 1, 28, 28), num_steps=500)
print("Generating samples with Probability Flow ODE...")
samples_ode = trainer.sample_probability_flow_ode((16, 1, 28, 28), num_steps=50)
visualize_continuous_diffusion(sde, score_net, history, samples_sde, samples_ode, device)
return trainer, history
def visualize_continuous_diffusion(sde: ContinuousSDE, score_net: ScoreNetwork,
history: Dict, samples_sde: torch.Tensor,
samples_ode: torch.Tensor, device: str):
"""
Visualization for continuous diffusion including SDE/ODE comparison and Fokker-Planck analysis.
"""
fig, axes = plt.subplots(2, 3, figsize=(18, 12))
# 1. Training loss
ax = axes[0, 0]
ax.plot(history['loss'], color='blue', linewidth=2)
ax.set_xlabel('Epoch')
ax.set_ylabel('Score Matching Loss')
ax.set_title(f'Continuous Diffusion Training ({sde.sde_type.upper()})')
ax.grid(True, alpha=0.3)
# 2. SDE trajectories (forward process)
ax = axes[0, 1]
# Simulate particle trajectories under SDE
num_particles = 5
timesteps = np.linspace(0, 1, 100)
trajectories = []
x0 = torch.randn(num_particles, 1, 28, 28, device=device) * 0.5
with torch.no_grad():
for t in timesteps:
mean, std = sde.marginal_prob(x0, torch.tensor(t))
# Just track mean trajectory for visualization
trajectories.append(mean[0, 0, 14, 14].cpu().item())
ax.plot(timesteps, trajectories, alpha=0.7)
ax.set_xlabel('Time')
ax.set_ylabel('State Value (center pixel)')
ax.set_title('SDE Forward Trajectories')
ax.grid(True, alpha=0.3)
# 3. Sigma/Noise schedule comparison
ax = axes[0, 2]
t_range = np.linspace(0, 1, 100)
if sde.sde_type == 've':
sigma_vals = [sde.sigma(t).item() for t in t_range]
ax.semilogy(t_range, sigma_vals, label='VE: σ(t)', linewidth=2)
else:
beta_vals = [sde.beta(t) for t in t_range]
ax.plot(t_range, beta_vals, label=f'{sde.sde_type.upper()}: β(t)', linewidth=2)
ax.set_xlabel('Time')
ax.set_ylabel('Noise Level')
ax.set_title('SDE Noise Schedule')
ax.legend()
ax.grid(True, alpha=0.3)
# 4. Generated samples: SDE vs ODE
ax = axes[1, 0]
grid_sde = samples_sde[:8].cpu().view(-1, 28, 28)
grid_sde = torch.cat([grid_sde[i] for i in range(8)], dim=1).numpy()
ax.imshow(grid_sde, cmap='gray')
ax.set_title('Reverse SDE Sampling (Stochastic)')
ax.axis('off')
ax = axes[1, 1]
grid_ode = samples_ode[:8].cpu().view(-1, 28, 28)
grid_ode = torch.cat([grid_ode[i] for i in range(8)], dim=1).numpy()
ax.imshow(grid_ode, cmap='gray')
ax.set_title('Probability Flow ODE (Deterministic)')
ax.axis('off')
# 5. Fokker-Planck density evolution (conceptual)
ax = axes[1, 2]
# Simulate probability density evolution
x_range = np.linspace(-3, 3, 100)
times = [0.0, 0.25, 0.5, 0.75, 1.0]
colors = plt.cm.viridis(np.linspace(0, 1, len(times)))
for t, color in zip(times, colors):
# Analytical solution for Ornstein-Uhlenbeck (approximation)
if sde.sde_type == 'vp':
mean_decay = np.exp(-0.5 * sde.beta_integral(t))
var = 1 - np.exp(-sde.beta_integral(t))
density = (1 / np.sqrt(2 * np.pi * var)) * np.exp(-0.5 * (x_range / np.sqrt(var)) ** 2)
ax.plot(x_range, density, color=color, label=f't={t:.2f}')
ax.set_xlabel('State')
ax.set_ylabel('Probability Density')
ax.set_title('Fokker-Planck Evolution\n(Forward Process)')
ax.legend()
ax.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig('continuous_diffusion_analysis.png', dpi=300)
plt.show()
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--sde_type', type=str, default='ve',
choices=['vp', 've', 'subvp'])
parser.add_argument('--epochs', type=int, default=50)
args = parser.parse_args()
train_continuous_diffusion(args.sde_type, args.epochs)13.4.3 流匹配（Flow Matching）的条件流（Conditional Flow）与边缘流一致性，Rectified Flow的直线路径学习与蒸馏（1步/4步生成），一致性模型（Consistency Models）的PF-ODE跳跃
流匹配框架直接学习传输数据的概率流场，通过回归向量场将先验分布传输至数据分布，条件流与边缘流一致性保证学得的流场等价于概率路径的梯度流。校正流算法通过 straightening 概率流轨迹，学习将弯曲ODE路径转化为直线路径，减少曲率积累的截断误差，经蒸馏后实现单步或少步高质量生成。一致性模型在概率流ODE轨迹上强制自一致性，任意点映射至轨迹终点，通过一致性损失训练直接预测最终样本，实现单步生成同时保持多步精细化能力。
Python
"""
Script: flow_matching_rectified_consistency.py
Content: Flow Matching with conditional flows, Rectified Flow for straight paths,
and Consistency Models for ODE jumping.
Usage: python flow_matching_rectified_consistency.py --method rectified_flow --epochs 100
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import numpy as np
import matplotlib.pyplot as plt
from typing import Tuple, Callable, Dict
import argparse
class FlowMatching:
"""
Flow Matching (Lipman et al., 2023; Liu et al., 2023).
Learns vector field u_t(x) that generates probability path.
"""
def __init__(self, sigma_min: float = 1e-4):
self.sigma_min = sigma_min
def sample_time(self, batch_size: int, device: str) -> torch.Tensor:
"""Sample time uniformly."""
return torch.rand(batch_size, device=device)
def sample_conditional_path(self, x0: torch.Tensor, x1: torch.Tensor,
t: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Conditional flow: x_t = (1-t)*x_0 + t*x_1.
Returns x_t and target velocity (dx_t/dt).
"""
# Interpolate
t_expanded = t.view(-1, 1, 1, 1)
x_t = (1 - t_expanded) * x0 + t_expanded * x1
# Target velocity: x1 - x0 (constant velocity along straight line)
target_velocity = x1 - x0
return x_t, target_velocity
def loss(self, v_theta: Callable, x1: torch.Tensor) -> torch.Tensor:
"""
Flow matching loss: MSE between predicted and target velocity.
"""
batch_size = x1.size(0)
device = x1.device
# Sample noise (prior)
x0 = torch.randn_like(x1)
# Sample time
t = self.sample_time(batch_size, device)
# Sample from conditional path
x_t, target_v = self.sample_conditional_path(x0, x1, t)
# Predict velocity
pred_v = v_theta(x_t, t)
# MSE loss
return F.mse_loss(pred_v, target_v)
class RectifiedFlow:
"""
Rectified Flow (Liu et al., 2023).
Straightens flow trajectories through reflow.
"""
def __init__(self, base_model: nn.Module):
self.model = base_model
self.reflow_iteration = 0
def sample_ode(self, x0: torch.Tensor, num_steps: int = 100) -> torch.Tensor:
"""
Sample by solving ODE: dx/dt = v(x, t).
"""
self.model.eval()
x = x0
dt = 1.0 / num_steps
trajectory = [x.cpu()]
with torch.no_grad():
for i in range(num_steps):
t = torch.ones(x.size(0), device=x.device) * (i * dt)
v = self.model(x, t)
x = x + v * dt
trajectory.append(x.cpu())
return x, trajectory
def reflow(self, x0_batch: torch.Tensor, num_steps: int = 100) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Generate paired (x_0, x_1) by sampling base model, then train new model on straight path.
"""
with torch.no_grad():
x1_generated, _ = self.sample_ode(x0_batch, num_steps)
# These pairs define a straighter path than the original curved ODE
return x0_batch, x1_generated
def compute_straightness(self, trajectory: list) -> float:
"""
Measure how straight the trajectory is (lower is straighter).
Straightness = E[||x1 - x0 - ∫v dt||^2].
"""
x0 = trajectory[0]
x1 = trajectory[-1]
straight_line = x1 - x0
# Actual path
path_integral = torch.zeros_like(x0)
for i in range(len(trajectory) - 1):
path_integral += trajectory[i+1] - trajectory[i]
deviation = torch.mean((straight_line - path_integral) ** 2).item()
return deviation
class ConsistencyModel(nn.Module):
"""
Consistency Models (Song et al., 2023).
Maps any point on ODE trajectory to the endpoint (t=1).
"""
def __init__(self, base_dim: int = 1, time_embed_dim: int = 128):
super().__init__()
self.time_embed_dim = time_embed_dim
# Time embedding
self.time_mlp = nn.Sequential(
nn.Linear(time_embed_dim, time_embed_dim),
nn.SiLU(),
nn.Linear(time_embed_dim, time_embed_dim)
)
# Consistency function f: (x, t) -> x_T (final sample)
# Parameterized as skip connection: f(x, t) = c_skip(t)*x + c_out(t)*F(x, t)
self.F = nn.Sequential(
nn.Conv2d(base_dim, 64, 3, padding=1),
nn.SiLU(),
nn.Conv2d(64, 128, 3, padding=1),
nn.SiLU(),
nn.Conv2d(128, 64, 3, padding=1),
nn.SiLU(),
nn.Conv2d(64, base_dim, 3, padding=1)
)
def get_time_embedding(self, t: torch.Tensor) -> torch.Tensor:
"""Sinusoidal embeddings."""
half_dim = self.time_embed_dim // 2
emb = torch.exp(torch.arange(half_dim, device=t.device) *
-(np.log(10000) / (half_dim - 1)))
emb = t[:, None] * emb[None, :]
emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
return emb
def c_skip(self, t: torch.Tensor) -> torch.Tensor:
"""Skip connection coefficient."""
return torch.ones_like(t[:, None, None, None]) # Start with identity
def c_out(self, t: torch.Tensor) -> torch.Tensor:
"""Output scaling."""
return t[:, None, None, None] # Scale with time
def forward(self, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
"""
Consistency function f(x, t).
"""
# Time embedding could condition F (simplified here)
t_emb = self.get_time_embedding(t)
t_emb = self.time_mlp(t_emb)
# Skip connection
output = self.c_skip(t) * x + self.c_out(t) * self.F(x)
return output
def consistency_loss(self, x1: torch.Tensor, v_theta: Callable,
num_distill_steps: int = 4) -> torch.Tensor:
"""
Consistency loss: enforce f(x_t, t) = f(x_{t'}, t') for any t, t'.
"""
batch_size = x1.size(0)
device = x1.device
# Sample noise
x0 = torch.randn_like(x1)
# Sample two different times
t1 = torch.rand(batch_size, device=device) * 0.9 + 0.05 # Avoid boundaries
t2 = torch.rand(batch_size, device=device) * 0.9 + 0.05
# Get x_t1 and x_t2 via ODE solver using v_theta
with torch.no_grad():
# Euler step from x0 to t1
x_t1 = x0 + t1.view(-1, 1, 1, 1) * v_theta(x0, torch.zeros_like(t1))
# Euler step from x0 to t2
x_t2 = x0 + t2.view(-1, 1, 1, 1) * v_theta(x0, torch.zeros_like(t2))
# Consistency: f(x_t1, t1) should equal f(x_t2, t2) (both should map to x1)
pred1 = self.forward(x_t1, t1)
pred2 = self.forward(x_t2, t2)
return F.mse_loss(pred1, pred2)
class VectorFieldNet(nn.Module):
"""
Neural network approximating the velocity field v_theta(x, t).
"""
def __init__(self, time_embed_dim: int = 128):
super().__init__()
self.time_embed_dim = time_embed_dim
self.time_mlp = nn.Sequential(
nn.Linear(time_embed_dim, time_embed_dim),
nn.SiLU(),
nn.Linear(time_embed_dim, time_embed_dim)
)
# U-Net style architecture for velocity field
self.encoder = nn.Sequential(
nn.Conv2d(1, 64, 3, padding=1),
nn.SiLU(),
nn.Conv2d(64, 128, 3, padding=1),
nn.SiLU()
)
self.bottleneck = nn.Sequential(
nn.Conv2d(128 + time_embed_dim, 128, 3, padding=1),
nn.SiLU()
)
self.decoder = nn.Sequential(
nn.Conv2d(128, 64, 3, padding=1),
nn.SiLU(),
nn.Conv2d(64, 1, 3, padding=1)
)
def get_time_embedding(self, t: torch.Tensor) -> torch.Tensor:
half_dim = self.time_embed_dim // 2
emb = torch.exp(torch.arange(half_dim, device=t.device) *
-(np.log(10000) / (half_dim - 1)))
emb = t[:, None] * emb[None, :]
emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
return emb
def forward(self, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
# Time embedding
t_emb = self.get_time_embedding(t)
t_emb = self.time_mlp(t_emb)
t_emb = t_emb[:, :, None, None].expand(-1, -1, x.size(2), x.size(3))
# Encode
h = self.encoder(x)
# Inject time
h = torch.cat([h, t_emb], dim=1)
h = self.bottleneck(h)
# Decode
v = self.decoder(h)
return v
def train_flow_model(method: str = 'flow_matching', epochs: int = 100):
"""
Training pipeline for Flow Matching, Rectified Flow, or Consistency Models.
"""
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,))
])
train_data = datasets.MNIST('./data', train=True, download=True, transform=transform)
train_loader = DataLoader(train_data, batch_size=128, shuffle=True)
v_theta = VectorFieldNet().to(device)
optimizer = torch.optim.Adam(v_theta.parameters(), lr=2e-4)
if method == 'consistency':
consistency_model = ConsistencyModel().to(device)
consistency_opt = torch.optim.Adam(consistency_model.parameters(), lr=2e-4)
flow_matching = FlowMatching()
else:
flow_matching = FlowMatching()
history = {'loss': [], 'straightness': []}
for epoch in range(epochs):
epoch_loss = 0
for batch_idx, (data, _) in enumerate(train_loader):
data = data.to(device)
if method == 'flow_matching':
optimizer.zero_grad()
loss = flow_matching.loss(v_theta, data)
loss.backward()
optimizer.step()
epoch_loss += loss.item()
elif method == 'rectified_flow':
# First train base flow
optimizer.zero_grad()
loss = flow_matching.loss(v_theta, data)
loss.backward()
optimizer.step()
epoch_loss += loss.item()
elif method == 'consistency':
# Train both flow and consistency model
optimizer.zero_grad()
flow_loss = flow_matching.loss(v_theta, data)
flow_loss.backward()
optimizer.step()
consistency_opt.zero_grad()
cons_loss = consistency_model.consistency_loss(data, v_theta)
cons_loss.backward()
consistency_opt.step()
epoch_loss += (flow_loss.item() + cons_loss.item()) / 2
avg_loss = epoch_loss / len(train_loader)
history['loss'].append(avg_loss)
# Measure trajectory straightness periodically
if method == 'rectified_flow' and epoch % 10 == 0:
rf = RectifiedFlow(v_theta)
x0 = torch.randn(10, 1, 28, 28, device=device)
_, traj = rf.sample_ode(x0, num_steps=50)
straightness = rf.compute_straightness(traj)
history['straightness'].append(straightness)
if epoch % 10 == 0:
print(f"Epoch {epoch}/{epochs} | Loss: {avg_loss:.4f}")
# Distillation / Sampling
if method == 'rectified_flow':
# Reflow: generate straight path data
print("\nPerforming reflow for straightening...")
rf = RectifiedFlow(v_theta)
x0_test = torch.randn(100, 1, 28, 28, device=device)
x0_new, x1_new = rf.reflow(x0_test, num_steps=100)
# Train new model on straight paths (distillation)
v_theta_straight = VectorFieldNet().to(device)
opt_straight = torch.optim.Adam(v_theta_straight.parameters(), lr=2e-4)
for epoch in range(20): # Short reflow training
for i in range(0, len(x0_new), 10):
batch_x0 = x0_new[i:i+10]
batch_x1 = x1_new[i:i+10]
opt_straight.zero_grad()
t = torch.rand(10, device=device)
xt, target = flow_matching.sample_conditional_path(batch_x0, batch_x1, t)
pred = v_theta_straight(xt, t)
loss = F.mse_loss(pred, target)
loss.backward()
opt_straight.step()
# Test 1-step vs 4-step generation
print("\nTesting 1-step generation...")
with torch.no_grad():
x1_1step = x0_test + v_theta_straight(x0_test, torch.ones(100, device=device))
print("Testing 4-step generation...")
with torch.no_grad():
x1_4step = x0_test
for t in [0.25, 0.5, 0.75, 1.0]:
t_batch = torch.ones(100, device=device) * t
x1_4step = x1_4step + 0.25 * v_theta_straight(x1_4step, t_batch)
visualize_flow_methods(v_theta, v_theta_straight if method == 'rectified_flow' else None,
consistency_model if method == 'consistency' else None,
history, method)
else:
visualize_flow_methods(v_theta, None,
consistency_model if method == 'consistency' else None,
history, method)
return v_theta, history
def visualize_flow_methods(v_theta: nn.Module, v_theta_straight: Optional[nn.Module],
consistency_model: Optional[nn.Module],
history: Dict, method: str):
"""
Visualization for Flow Matching and related methods.
"""
fig, axes = plt.subplots(2, 3, figsize=(18, 12))
# 1. Training loss
ax = axes[0, 0]
ax.plot(history['loss'], color='blue', linewidth=2)
ax.set_xlabel('Epoch')
ax.set_ylabel('Flow Matching Loss')
ax.set_title(f'{method.upper()} Training')
ax.grid(True, alpha=0.3)
# 2. Trajectory comparison (curved vs straight)
if method == 'rectified_flow' and v_theta_straight is not None:
ax = axes[0, 1]
device = next(v_theta.parameters()).device
# Generate trajectory with base model
x0 = torch.randn(1, 1, 28, 28, device=device)
rf_base = RectifiedFlow(v_theta)
_, traj_base = rf_base.sample_ode(x0, num_steps=20)
# Generate with straightened model
rf_straight = RectifiedFlow(v_theta_straight)
_, traj_straight = rf_straight.sample_ode(x0, num_steps=20)
# Plot trajectories (project to 2D via center pixel)
base_path = [t[0, 0, 14, 14].cpu().item() for t in traj_base]
straight_path = [t[0, 0, 14, 14].cpu().item() for t in traj_straight]
ax.plot(base_path, label='Base Flow (Curved)', marker='o', alpha=0.7)
ax.plot(straight_path, label='Reflowed (Straight)', marker='s', alpha=0.7)
ax.set_xlabel('Step')
ax.set_ylabel('State (Center Pixel)')
ax.set_title('Trajectory Straightening')
ax.legend()
ax.grid(True, alpha=0.3)
# 3. Straightness evolution
if method == 'rectified_flow' and history['straightness']:
ax = axes[0, 2]
ax.plot(range(0, len(history['loss']), 10), history['straightness'],
marker='o', color='green', linewidth=2)
ax.set_xlabel('Epoch')
ax.set_ylabel('Trajectory Straightness')
ax.set_title('Reflow Progress')
ax.grid(True, alpha=0.3)
# 4. Generated samples (Flow Matching)
ax = axes[1, 0]
device = next(v_theta.parameters()).device
with torch.no_grad():
x0 = torch.randn(8, 1, 28, 28, device=device)
t_vals = torch.linspace(0, 1, 10)
# Show evolution
samples = []
for t in [0.0, 0.3, 0.6, 1.0]:
t_batch = torch.ones(8, device=device) * t
if t == 0:
xt = x0
elif t == 1.0:
xt = x0 + v_theta(x0, torch.zeros_like(t_batch))
# Multi-step refinement
for _ in range(9):
xt = xt + v_theta(xt, torch.ones_like(t_batch))
else:
xt = (1-t) * x0 + t * (x0 + v_theta(x0, torch.zeros_like(t_batch)))
samples.append(xt.cpu())
# Plot final samples
final_samples = samples[-1].view(-1, 28, 28)
grid = torch.cat([final_samples[i] for i in range(8)], dim=1).numpy()
ax.imshow(grid, cmap='gray')
ax.set_title('Flow Matching Samples')
ax.axis('off')
# 5. Consistency model samples (single step)
if method == 'consistency' and consistency_model is not None:
ax = axes[1, 1]
with torch.no_grad():
x0 = torch.randn(8, 1, 28, 28, device=device)
t_batch = torch.ones(8, device=device) * 0.5 # Any intermediate time
# Consistency model jumps directly to t=1
x1_direct = consistency_model(x0, t_batch)
samples = x1_direct.cpu().view(-1, 28, 28)
grid = torch.cat([samples[i] for i in range(8)], dim=1).numpy()
ax.imshow(grid, cmap='gray')
ax.set_title('Consistency Model (1-Step)')
ax.axis('off')
# 6. Step ablation (1-step vs multi-step)
if method == 'rectified_flow':
ax = axes[1, 2]
device = next(v_theta.parameters()).device
with torch.no_grad():
x0 = torch.randn(50, 1, 28, 28, device=device)
# 1-step
x1_1step = x0 + v_theta_straight(x0, torch.ones(50, device=device))
# 4-step
x = x0
for t in [0.25, 0.5, 0.75, 1.0]:
t_batch = torch.ones(50, device=device) * t
x = x + 0.25 * v_theta_straight(x, t_batch)
x1_4step = x
# Compute sample quality proxy (variance across batch)
quality_1 = torch.var(x1_1step).item()
quality_4 = torch.var(x1_4step).item()
ax.bar(['1-Step', '4-Step'], [quality_1, quality_4], color=['red', 'blue'], alpha=0.7)
ax.set_ylabel('Sample Variance (Diversity)')
ax.set_title('Step Count vs Sample Quality')
plt.tight_layout()
plt.savefig(f'flow_matching_{method}.png', dpi=300)
plt.show()
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--method', type=str, default='flow_matching',
choices=['flow_matching', 'rectified_flow', 'consistency'])
parser.add_argument('--epochs', type=int, default=100)
args = parser.parse_args()
train_flow_model(args.method, args.epochs)13.4.4 引导生成：Classifier Guidance的梯度缩放与无分类器引导（CFG）的隐式分类器，PNDM/DDIM的确定性采样与步数加速
引导生成技术利用条件信息控制扩散模型的采样过程，分类器引导借助显式分类器的梯度信息调整扩散轨迹，通过梯度缩放系数平衡生成样本的逼真度与类别一致性，缩放系数增强引导强度使生成样本更忠实于条件约束，过高则牺牲样本多样性。无分类器引导摒弃显式分类器，在训练阶段随机丢弃条件信息学习无条件与条件生成联合分布，采样时结合条件与无条件分数估计实现隐式分类，避免额外分类器训练与对抗攻击脆弱性。
概率流采样器将随机微分方程转化为确定性常微分方程，利用离散化求解实现大步长跳跃，数据预测参数化与噪声预测参数化支持解析积分减少迭代次数。PNDM与DDIM等方法通过非马尔可夫扩散过程设计，允许任意采样步数配置，利用历史梯度信息或解析转移公式，在十步至五十步内实现高质量生成，显著降低采样延迟满足实时应用需求。
Python
"""
Script: guided_sampling_cfg_ddim.py
Content: Classifier Guidance, Classifier-Free Guidance (CFG),
DDIM/PNDM deterministic sampling with step acceleration.
Usage: python guided_sampling_cfg_ddim.py --guidance_type cfg --num_steps 50
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import numpy as np
import matplotlib.pyplot as plt
from typing import Tuple, Optional, Callable
import argparse
class Classifier(nn.Module):
"""
Explicit classifier for classifier guidance.
"""
def __init__(self, num_classes: int = 10):
super().__init__()
self.net = nn.Sequential(
nn.Linear(784, 256),
nn.ReLU(),
nn.Linear(256, 256),
nn.ReLU(),
nn.Linear(256, num_classes)
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.net(x.view(-1, 784))
def gradient(self, x: torch.Tensor, target_class: int) -> torch.Tensor:
"""Compute gradient of log p(y|x) w.r.t. x."""
x.requires_grad_(True)
logits = self.forward(x)
log_prob = F.log_softmax(logits, dim=-1)[:, target_class]
grad = torch.autograd.grad(log_prob.sum(), x)[0]
return grad
class ConditionalUNet(nn.Module):
"""
UNet supporting both conditional and unconditional generation for CFG.
"""
def __init__(self, num_classes: int = 10, time_embed_dim: int = 128,
dropout_prob: float = 0.1):
super().__init__()
self.num_classes = num_classes
self.dropout_prob = dropout_prob
# Class embedding
self.class_embed = nn.Embedding(num_classes, time_embed_dim)
# Time embedding (shared)
self.time_mlp = nn.Sequential(
nn.Linear(time_embed_dim, time_embed_dim),
nn.SiLU(),
nn.Linear(time_embed_dim, time_embed_dim)
)
# Network (simplified)
self.net = nn.Sequential(
nn.Linear(784 + time_embed_dim, 512),
nn.ReLU(),
nn.Linear(512, 512),
nn.ReLU(),
nn.Linear(512, 784)
)
# Null token for unconditional generation
self.null_token = nn.Parameter(torch.zeros(time_embed_dim))
def get_time_embedding(self, t: torch.Tensor, dim: int) -> torch.Tensor:
half_dim = dim // 2
emb = torch.exp(torch.arange(half_dim, device=t.device) *
-(np.log(10000) / (half_dim - 1)))
emb = t[:, None] * emb[None, :]
return torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
def forward(self, x: torch.Tensor, t: torch.Tensor,
c: Optional[torch.Tensor] = None) -> torch.Tensor:
batch_size = x.size(0)
x_flat = x.view(batch_size, -1)
# Time embedding
t_emb = self.get_time_embedding(t, 128)
t_emb = self.time_mlp(t_emb)
# Class embedding with dropout for CFG training
if c is not None and self.training:
# Random dropout for CFG
mask = torch.rand(batch_size, device=x.device) > self.dropout_prob
c_emb = self.class_embed(c) * mask.unsqueeze(1).float()
c_emb = c_emb + self.null_token.unsqueeze(0) * (~mask).unsqueeze(1).float()
elif c is not None:
c_emb = self.class_embed(c)
else:
c_emb = self.null_token.unsqueeze(0).expand(batch_size, -1)
# Combine
h = torch.cat([x_flat, t_emb + c_emb], dim=1)
return self.net(h).view_as(x)
class GuidedDiffusion:
"""
Diffusion model with Classifier Guidance or Classifier-Free Guidance.
"""
def __init__(self, model: ConditionalUNet, classifier: Optional[Classifier] = None,
num_timesteps: int = 1000):
self.model = model
self.classifier = classifier
self.num_timesteps = num_timesteps
# Linear beta schedule
self.betas = torch.linspace(1e-4, 0.02, num_timesteps)
self.alphas = 1.0 - self.betas
self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
self.sqrt_alphas_cumprod = torch.sqrt(self.alphas_cumprod)
self.sqrt_one_minus_alphas_cumprod = torch.sqrt(1.0 - self.alphas_cumprod)
self.alphas_cumprod_prev = F.pad(self.alphas_cumprod[:-1], (1, 0), value=1.0)
self.posterior_variance = self.betas * (1.0 - self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
def q_sample(self, x_0: torch.Tensor, t: torch.Tensor, noise: Optional[torch.Tensor] = None):
if noise is None:
noise = torch.randn_like(x_0)
sqrt_acp = self.sqrt_alphas_cumprod[t].view(-1, 1, 1, 1)
sqrt_omacp = self.sqrt_one_minus_alphas_cumprod[t].view(-1, 1, 1, 1)
return sqrt_acp * x_0 + sqrt_omacp * noise
def predict_noise(self, x_t: torch.Tensor, t: torch.Tensor, c: Optional[torch.Tensor] = None):
return self.model(x_t, t, c)
def classifier_guidance(self, x_t: torch.Tensor, t: torch.Tensor,
target_class: int, guidance_scale: float = 2.0) -> torch.Tensor:
"""
Explicit classifier guidance.
"""
with torch.enable_grad():
x_t_grad = x_t.clone().requires_grad_(True)
class_grad = self.classifier.gradient(x_t_grad, target_class)
# Modify score with classifier gradient
noise_pred = self.predict_noise(x_t, t, None)
guided_noise = noise_pred - guidance_scale * self.sqrt_one_minus_alphas_cumprod[t].view(-1, 1, 1, 1) * class_grad
return guided_noise
def classifier_free_guidance(self, x_t: torch.Tensor, t: torch.Tensor,
c: torch.Tensor, guidance_scale: float = 7.5) -> torch.Tensor:
"""
Classifier-Free Guidance (CFG).
"""
# Conditional prediction
noise_cond = self.predict_noise(x_t, t, c)
# Unconditional prediction
noise_uncond = self.predict_noise(x_t, t, None)
# Weighted combination
noise_pred = noise_uncond + guidance_scale * (noise_cond - noise_uncond)
return noise_pred
def ddim_sample(self, x_t: torch.Tensor, t: torch.Tensor, c: Optional[torch.Tensor] = None,
guidance_scale: float = 7.5, use_cfg: bool = True) -> torch.Tensor:
"""
DDIM deterministic sampling step.
"""
batch_size = x_t.size(0)
# Predict noise with guidance
if use_cfg and c is not None:
noise_pred = self.classifier_free_guidance(x_t, t, c, guidance_scale)
else:
noise_pred = self.predict_noise(x_t, t, c)
# Predict x_0
alpha_t = self.alphas_cumprod[t][0]
pred_x0 = (x_t - torch.sqrt(1 - alpha_t) * noise_pred) / torch.sqrt(alpha_t)
# DDIM deterministic update
if t[0] > 0:
alpha_prev = self.alphas_cumprod[t - 1][0]
sigma_t = 0 # Deterministic (eta = 0)
x_t_prev = (
torch.sqrt(alpha_prev) * pred_x0 +
torch.sqrt(1 - alpha_prev - sigma_t**2) * noise_pred
)
else:
x_t_prev = pred_x0
return x_t_prev
def pndm_sample(self, x_t: torch.Tensor, t: torch.Tensor,
old_noise_preds: list, c: Optional[torch.Tensor] = None) -> torch.Tensor:
"""
PNDM (Pseudo Numerical Methods for Diffusion Models) sampling.
Uses Runge-Kutta-like gradient accumulation.
"""
noise_pred = self.predict_noise(x_t, t, c)
if len(old_noise_preds) < 3:
# Use DDIM for first few steps
return self.ddim_sample(x_t, t, c, use_cfg=False)
# PNDM update using history
# Simplified version: use combination of current and previous gradients
noise_combined = (55 * noise_pred - 59 * old_noise_preds[-1] +
37 * old_noise_preds[-2] - 9 * old_noise_preds[-3]) / 24
# Update x
alpha_t = self.alphas_cumprod[t][0]
if t[0] > 0:
alpha_prev = self.alphas_cumprod[t - 1][0]
x_t_prev = (x_t - torch.sqrt(1 - alpha_t) * noise_combined) * torch.sqrt(alpha_prev / alpha_t) + \
torch.sqrt(1 - alpha_prev) * noise_combined
else:
x_t_prev = (x_t - torch.sqrt(1 - alpha_t) * noise_combined) / torch.sqrt(alpha_t)
return x_t_prev
def train_guided_diffusion(guidance_type: str = 'cfg', epochs: int = 50):
"""
Training pipeline for guided diffusion models.
"""
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,))
])
train_data = datasets.MNIST('./data', train=True, download=True, transform=transform)
train_loader = DataLoader(train_data, batch_size=128, shuffle=True)
# Model
model = ConditionalUNet(dropout_prob=0.1).to(device) if guidance_type == 'cfg' else \
ConditionalUNet(dropout_prob=0.0).to(device)
# Classifier for classifier guidance
classifier = Classifier().to(device) if guidance_type == 'classifier' else None
if classifier is not None:
# Pre-train classifier
print("Pre-training classifier...")
clf_opt = torch.optim.Adam(classifier.parameters(), lr=1e-3)
for epoch in range(5):
for data, labels in train_loader:
data, labels = data.to(device), labels.to(device)
clf_opt.zero_grad()
logits = classifier(data)
loss = F.cross_entropy(logits, labels)
loss.backward()
clf_opt.step()
# Diffusion
diffusion = GuidedDiffusion(model, classifier)
optimizer = torch.optim.Adam(model.parameters(), lr=2e-4)
history = {'loss': []}
for epoch in range(epochs):
epoch_loss = 0
for batch_idx, (data, labels) in enumerate(train_loader):
data, labels = data.to(device), labels.to(device)
t = torch.randint(0, diffusion.num_timesteps, (data.size(0),), device=device)
noise = torch.randn_like(data)
x_t = diffusion.q_sample(data, t, noise)
optimizer.zero_grad()
# Predict noise
if guidance_type == 'cfg':
# Randomly drop labels for unconditional training
pred_noise = model(x_t, t, labels)
else:
pred_noise = model(x_t, t, None)
loss = F.mse_loss(pred_noise, noise)
loss.backward()
optimizer.step()
epoch_loss += loss.item()
avg_loss = epoch_loss / len(train_loader)
history['loss'].append(avg_loss)
if epoch % 5 == 0:
print(f"Epoch {epoch}/{epochs} | Loss: {avg_loss:.4f}")
# Sampling with different methods
print("\nGenerating samples with DDIM (50 steps)...")
samples_ddim_50 = sample_with_method(diffusion, device, num_steps=50, method='ddim',
guidance_type=guidance_type, target_class=7)
print("Generating samples with DDIM (10 steps)...")
samples_ddim_10 = sample_with_method(diffusion, device, num_steps=10, method='ddim',
guidance_type=guidance_type, target_class=7)
if guidance_type == 'cfg':
print("Generating samples with different guidance scales...")
scales = [1.0, 3.0, 7.5, 15.0]
samples_scales = []
for scale in scales:
samples = sample_with_method(diffusion, device, num_steps=50, method='ddim',
guidance_type=guidance_type, target_class=7,
guidance_scale=scale)
samples_scales.append((scale, samples))
else:
samples_scales = None
visualize_guided_sampling(diffusion, history, samples_ddim_50, samples_ddim_10,
samples_scales, guidance_type)
return diffusion, history
def sample_with_method(diffusion: GuidedDiffusion, device: str, num_steps: int = 50,
method: str = 'ddim', guidance_type: str = 'cfg',
target_class: int = 7, guidance_scale: float = 7.5,
batch_size: int = 8) -> torch.Tensor:
"""
Sample using specified method and guidance.
"""
model = diffusion.model
model.eval()
shape = (batch_size, 1, 28, 28)
x = torch.randn(shape, device=device)
# Select timesteps
timesteps = torch.linspace(diffusion.num_timesteps - 1, 0, num_steps, dtype=torch.long)
with torch.no_grad():
old_noise_preds = []
for i, t_val in enumerate(timesteps):
t = torch.full((batch_size,), t_val, device=device, dtype=torch.long)
c = torch.full((batch_size,), target_class, device=device, dtype=torch.long)
if method == 'ddim':
x = diffusion.ddim_sample(x, t, c, guidance_scale, use_cfg=(guidance_type == 'cfg'))
elif method == 'pndm':
x = diffusion.pndm_sample(x, t, old_noise_preds, c)
old_noise_preds.append(diffusion.predict_noise(x, t, c))
if len(old_noise_preds) > 4:
old_noise_preds.pop(0)
return x.cpu()
def visualize_guided_sampling(diffusion: GuidedDiffusion, history: Dict,
samples_ddim_50: torch.Tensor, samples_ddim_10: torch.Tensor,
samples_scales: Optional[list], guidance_type: str):
"""
Visualization for guided sampling comparison.
"""
fig, axes = plt.subplots(2, 3, figsize=(18, 12))
# 1. Training loss
ax = axes[0, 0]
ax.plot(history['loss'], color='blue', linewidth=2)
ax.set_xlabel('Epoch')
ax.set_ylabel('MSE Loss')
ax.set_title(f'{guidance_type.upper()} Training Loss')
ax.grid(True, alpha=0.3)
# 2. Sample quality vs steps (DDIM)
ax = axes[0, 1]
# Show 50-step vs 10-step
grid_50 = torch.cat([samples_ddim_50[i].squeeze() for i in range(8)], dim=1).numpy()
ax.imshow(grid_50, cmap='gray')
ax.set_title('DDIM 50 Steps')
ax.axis('off')
ax = axes[0, 2]
grid_10 = torch.cat([samples_ddim_10[i].squeeze() for i in range(8)], dim=1).numpy()
ax.imshow(grid_10, cmap='gray')
ax.set_title('DDIM 10 Steps')
ax.axis('off')
# 3. Guidance scale comparison (CFG)
if guidance_type == 'cfg' and samples_scales is not None:
ax = axes[1, 0]
for idx, (scale, samples) in enumerate(samples_scales[:2]):
plt.subplot(1, 2, idx + 1)
grid = torch.cat([samples[i].squeeze() for i in range(8)], dim=1).numpy()
plt.imshow(grid, cmap='gray')
plt.title(f'Scale = {scale}')
plt.axis('off')
# 4. Classifier-Free vs Classifier Guidance comparison
ax = axes[1, 1]
methods = ['No Guidance', 'Classifier Guidance', 'CFG (Scale=7.5)']
qualities = [0.7, 0.85, 0.9] # Proxy quality scores
colors = ['gray', 'blue', 'green']
ax.bar(methods, qualities, color=colors, alpha=0.7)
ax.set_ylabel('Sample Quality (Proxy)')
ax.set_title('Guidance Method Comparison')
ax.set_ylim(0, 1)
# 5. Step ablation curve
ax = axes[1, 2]
steps = [5, 10, 20, 50, 100]
fids = [45, 25, 15, 8, 5] # Simulated FID scores (lower is better)
ax.plot(steps, fids, marker='o', color='red', linewidth=2)
ax.set_xlabel('Number of Sampling Steps')
ax.set_ylabel('FID Score (Lower Better)')
ax.set_title('DDIM Step Efficiency')
ax.set_xscale('log')
ax.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig(f'guided_sampling_{guidance_type}.png', dpi=300)
plt.show()
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--guidance_type', type=str, default='cfg',
choices=['none', 'classifier', 'cfg'])
parser.add_argument('--num_steps', type=int, default=50)
parser.add_argument('--epochs', type=int, default=50)
args = parser.parse_args()
train_guided_diffusion(args.guidance_type, args.epochs)
58
扫一扫
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
