AIGC 中的变分自编码器（VAE）原理与 PyTorch 实战

3.2 定义 VAE 的结构

我们需要构建编码器和解码器。编码器输出均值 $\mu$ 和对数方差 $\log\sigma^2$，并通过重参数化技巧采样得到 $z$。

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 加载数据集

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 训练模型

在实际训练中，我们需要注意设备选择（CPU/GPU）以及梯度清零的细节。

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 图像生成

利用训练好的模型，可以从隐空间随机采样生成新图像。通过对 $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 数据压缩

编码器能将高维数据映射到低维隐变量空间，实现降维和特征提取。

4.3 数据补全

结合生成能力，VAE 可用于预测缺失的数据部分。

4.4 多模态生成

扩展后的 VAE 架构可支持跨模态内容生成，例如从文本描述生成图像。

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

6. 总结

变分自编码器（VAE）凭借其概率建模能力和隐空间的可解释性，在图像生成、数据降维等领域展现了独特价值。虽然其生成样本的清晰度有时不如 GAN，但其训练过程的稳定性和理论完备性使其成为许多场景下的首选。

建议读者在实际项目中尝试在自己的数据集上复现上述代码，通过调整超参数观察隐空间的变化，这将有助于更深入地理解生成模型的运作机制。