变分自编码器（VAE）原理与 PyTorch 实战实现

定义 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

定义损失函数

损失函数包含两部分：重建误差（BCE）和 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

加载数据集

这里加载 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)

训练模型

准备好设备、优化器和训练循环。注意保存生成的样本以便观察效果。

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')

VAE 的应用

图像生成

利用训练好的 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')

数据压缩

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

数据补全

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

多模态生成

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

VAE 与其他生成模型的对比

总结

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

通过上述原理讲解和代码实现，希望能帮助大家深入理解 VAE 的工作机制及其在 AIGC 中的实际应用。如果感兴趣，不妨尝试在自己的数据集上进行训练与测试！