AIGC 中的变分自编码器（VAE）代码与实现

3.1 导入必要的库

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from torchvision.utils import save_image
import os

3.2 定义 VAE 的结构

编码器与解码器的实现：

# 定义 VAE 模型
class VAE(nn.Module):
    def __init__(self, input_dim=784, hidden_dim=400, latent_dim=20):
        super(VAE, self).__init__()
        # 编码器
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.fc_mu = nn.Linear(hidden_dim, latent_dim)
        self.fc_logvar = nn.Linear(hidden_dim, latent_dim)
        # 解码器
        self.fc2 = nn.Linear(latent_dim, hidden_dim)
        self.fc3 = nn.Linear(hidden_dim, input_dim)
        self.sigmoid = nn.Sigmoid()

    def encode(self, x):
        h1 = torch.relu(self.fc1(x))
        mu = self.fc_mu(h1)
        logvar = self.fc_logvar(h1)
        return mu, logvar

    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mu + eps * std

    def decode(self, z):
        h2 = torch.relu(self.fc2(z))
        return self.sigmoid(self.fc3(h2))

    def forward(self, x):
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar)
        return self.decode(z), mu, logvar

3.3 定义损失函数

# 损失函数包含重建误差和 KL 散度
def loss_function(recon_x, x, mu, logvar):
    BCE = nn.functional.binary_cross_entropy(recon_x, x, reduction='sum')
    KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
    return BCE + KLD

3.4 加载数据集

# 加载 MNIST 数据集
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])
dataset = datasets.MNIST(root='./data', train=True, transform=transform, download=True)
dataloader = DataLoader(dataset, batch_size=128, shuffle=True)

3.5 训练模型

# 训练 VAE 模型
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
vae = VAE().to(device)
optimizer = optim.Adam(vae.parameters(), lr=1e-3)
epochs = 10
for epoch in range(epochs):
    vae.train()
    train_loss = 0
    for batch_idx, (data, _) in enumerate(dataloader):
        data = data.view(-1, 784).to(device)
        optimizer.zero_grad()
        recon_batch, mu, logvar = vae(data)
        loss = loss_function(recon_batch, data, mu, logvar)
        loss.backward()
        train_loss += loss.item()
        optimizer.step()
    print(f'Epoch [{epoch+1}/{epochs}], Loss: {train_loss/len(dataloader.dataset):.4f}')
    # 保存生成的样本
    with torch.no_grad():
        z = torch.randn(64, 20).to(device)
        sample = vae.decode(z).cpu()
        save_image(sample.view(64, 1, 28, 28), f'./results/sample_{epoch+1}.png')

4. VAE 的应用

4.1 图像生成

• 利用训练好的 VAE 模型，可以生成与训练数据分布相似的图像。
• 通过对隐变量 ( z ) 进行插值，可以生成不同风格的图像。

示例：生成图像

# 从隐空间采样并生成图像
vae.eval()
with torch.no_grad():
    z = torch.randn(16, 20).to(device)
    # 生成随机潜在向量
    sample = vae.decode(z).cpu()
    save_image(sample.view(16, 1, 28, 28), 'generated_images.png')

4.2 数据压缩

• VAE 的编码器能够将高维数据压缩到低维隐变量空间，实现数据降维和压缩。

4.3 数据补全

• VAE 可用于缺失数据补全，通过生成模型预测缺失部分。

4.4 多模态生成

• 通过扩展，VAE 可用于生成跨模态内容（如从文本生成图像）。

5. VAE 与其他生成模型的对比

6. 总结

变分自编码器（VAE）作为一种生成式模型，凭借其概率建模能力和隐空间结构化表示，在图像生成、数据降维、数据补全等领域展现了强大的能力。尽管 VAE 生成的样本质量可能不如 GAN，但其稳定性和解释性使其成为许多应用场景的首选模型。