Diffusion 模型生成图片太慢了？用 AE(AutoEncoder) 和 VAE 加速一波！

Diffusion 模型生成图片太慢了？用 AE(AutoEncoder) 和 VAE(Variational AutoEncoder) 压缩图片加速一波！

本文分为两个部分：理论学习和代码实践。

先预览一下 AE 和 VAE 模型的效果：

理论学习

本文价值

Diffusion 模型生成图片的效果堪称惊艳，但是推理速度慢的问题被广泛吐槽。本文介绍 Diffusion 模型推理加速的一种常见方式：用 AE(AutoEncoder) 和 VAE(Variational AutoEncoder) 进行图片压缩/反压缩。理论部分学完之后立即用代码进行实践，彻底掌握 AE/VAE。

AE 基础知识

自编码器（AutoEncoder，AE）是一种无监督学习的神经网络模型，主要用于数据压缩和特征学习。它的核心结构包括两个主要部分：编码器和解码器。

编码器负责将输入数据压缩到一个低维的潜在空间，这个过程可以看作是提取输入数据的关键特征。解码器则尝试从这个压缩的表示重构原始输入，目标是使重构的输出尽可能接近原始输入。

AE 通过最小化重构误差来训练，这促使网络学习输入数据的最重要特征。训练完成后，编码器可以用于降维、特征提取或数据压缩，而完整的 AE 可以用于去噪或异常检测等任务。

AE 的优点包括结构简单、训练相对快速，以及可以学习紧凑的特征表示。然而，它也存在一些局限性，如生成能力有限，难以生成新的、有意义的样本。

VAE 基础知识

变分自编码器（VAE）是自编码器的一种概率变体，它结合了变分推断和神经网络，用于生成模型和表示学习。VAE 的核心思想是将输入数据编码为概率分布，而不是固定的向量。

VAE 的结构包括编码器、采样层和解码器。编码器将输入映射到潜在空间的均值和方差，采样层从这个分布中采样，解码器则从采样的潜在向量重构输入。

VAE 的训练目标包括两部分：重构损失和 KL 散度。重构损失确保模型能够准确重建输入，而 KL 散度则作为正则化项，使潜在空间的分布接近标准正态分布。

相比传统的自编码器，VAE 具有更强的生成能力，可以生成新的、合理的样本。它的潜在空间是连续的，便于插值，并且具有一定的正则化效果，有助于减少过拟合。

VAE 广泛应用于图像生成、异常检测、数据增强等领域。然而，它的训练过程可能较为复杂和不稳定，且 KL 散度项可能导致模型忽视部分输入信息。

对比 AE 和 VAE

特性	AutoEncoder (AE)	Variational AutoEncoder (VAE)
模型结构	• 编码器：将输入压缩到潜在空间 • 解码器：从潜在空间重构输入 • 通常是确定性的	• 编码器：将输入映射到潜在空间的均值和方差 • 采样层：从编码器输出的分布中采样 • 解码器：从采样的潜在向量重构输入
潜在空间	离散的点	连续的概率分布
损失函数	重构损失（如 MSE）	重构损失 + KL 散度（正则化项）
训练过程	1. 前向传播 2. 计算重构误差 3. 反向传播 4. 更新参数	1. 前向传播（包括采样步骤） 2. 计算重构误差和 KL 散度 3. 反向传播 4. 更新参数
推理过程	1. 输入通过编码器 2. 潜在向量通过解码器	1. 输入通过编码器得到分布参数 2. 从分布中采样 3. 采样的向量通过解码器
优点	• 结构简单，易于理解和实现 • 训练相对快速 • 可以学习紧凑的特征表示	• 生成能力强 • 潜在空间连续，便于插值 • 可以生成新的、合理的样本 • 具有正则化效果，减少过拟合
缺点	• 生成能力有限 • 潜在空间可能不连续 • 难以生成新的、有意义的样本 • 可能过拟合训练数据	• 结构和训练过程较复杂 • 训练可能更慢且不稳定 • KL 散度项可能导致模型忽视部分输入
应用场景	• 数据压缩 • 降噪 • 特征提取	• 图像生成 • 异常检测 • 数据增强 • 条件生成
潜在空间特性	• 不保证平滑或有意义的插值 • 可能存在"空洞"	• 通常呈现高斯分布 • 允许平滑插值 • 潜在空间更加结构化
生成新样本	困难，通常需要额外的技巧	容易，直接从先验分布采样即可
对异常值的敏感度	较高，可能过度拟合异常值	较低，由于正则化效果而更鲁棒
可解释性	潜在表示可能难以解释	潜在空间通常更有结构，便于解释

代码实践

为了深刻理解 AE 和 VAE 的模型结构和训练/评测流程，我们自己动手实现在 MNIST 数据集上训练和评测 AE 和 VAE 模型。先看一下最终 AE 和 VAE 的重构效果：

我们先实现最简单的 Baseline 版本，这个版本的效果比较一般，主要是因为 latent space 维度为 3；然后我们改进 Baseline 版本，改进点包括：

将 latent space 维度从 3 --> 16；
增加了网络深度和宽度；
使用了 LeakyReLU 激活函数；
添加了批归一化；
使用了 AdamW 优化器和学习率调度器；
增加了训练轮数到 500；
添加了简单的数据增强（随机旋转和噪声）；
对 VAE 的 KL 散度使用了 β 参数（设为 0.5）来平衡重构质量和潜在空间的规则性。

Baseline 版本

先看效果图，再看代码：

完整的代码如下：

import torch,os
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import matplotlib.pyplot as plt

# 设置随机种子以确保结果可复现
torch.manual_seed(42)

# 准备MNIST数据集
transform = transforms.Compose([transforms.ToTensor()])
train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)

# 创建数据加载器
# batch_size: 每批处理的样本数
# shuffle: 是否在每个epoch打乱数据
train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=128, shuffle=False)

class AE(nn.Module):
    def __init__(self):
        super(AE, self).__init__()
        # 编码器：将28x28的输入压缩到3维潜在空间
        self.encoder = nn.Sequential(
            nn.Linear(28*28, 128),
            nn.ReLU(),
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Linear(64, 12),
            nn.ReLU(),
            nn.Linear(12, 3)
        )
        # 解码器：将3维潜在空间重构为28x28的输出
        self.decoder = nn.Sequential(
            nn.Linear(3, 12),
            nn.ReLU(),
            nn.Linear(12, 64),
            nn.ReLU(),
            nn.Linear(64, 128),
            nn.ReLU(),
            nn.Linear(128, 28*28),
            nn.Sigmoid()  # 使用Sigmoid确保输出在[0,1]范围内
        )

    def forward(self, x):
        """
        前向传播函数
        
        参数:
        x (torch.Tensor): 输入图像张量，形状为 (batch_size, 1, 28, 28)
        
        返回:
        torch.Tensor: 重构后的图像张量，形状为 (batch_size, 1, 28, 28)
        """
        x = x.view(-1, 28*28)  # 将输入展平
        encoded = self.encoder(x)
        decoded = self.decoder(encoded)
        return decoded.view(-1, 1, 28, 28)  # 重塑为原始图像形状

# 初始化AE模型
ae_model = AE()
ae_optimizer = optim.Adam(
    ae_model.parameters(),
    lr=0.0001
)
ae_criterion = nn.MSELoss()  # 使用均方误差作为重构损失

class VAE(nn.Module):
    def __init__(self):
        super(VAE, self).__init__()
        # 编码器
        self.encoder = nn.Sequential(
            nn.Linear(28*28, 128),
            nn.ReLU(),
            nn.Linear(128, 64),
            nn.ReLU()
        )
        # 均值和对数方差的全连接层
        self.fc_mu = nn.Linear(64, 3)
        self.fc_logvar = nn.Linear(64, 3)
        # 解码器
        self.decoder = nn.Sequential(
            nn.Linear(3, 64),
            nn.ReLU(),
            nn.Linear(64, 128),
            nn.ReLU(),
            nn.Linear(128, 28*28),
            nn.Sigmoid()
        )

    def encode(self, x):
        """
        编码函数
        
        参数:
        x (torch.Tensor): 输入图像张量，形状为 (batch_size, 784)
        
        返回:
        tuple(torch.Tensor, torch.Tensor): 均值和对数方差，每个的形状为 (batch_size, 3)
        """
        h = self.encoder(x)
        return self.fc_mu(h), self.fc_logvar(h)

    def reparameterize(self, mu, logvar):
        """
        重参数化技巧
        
        参数:
        mu (torch.Tensor): 均值，形状为 (batch_size, 3)
        logvar (torch.Tensor): 对数方差，形状为 (batch_size, 3)
        
        返回:
        torch.Tensor: 采样得到的潜在变量，形状为 (batch_size, 3)
        """
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mu + eps * std

    def decode(self, z):
        """
        解码函数
        
        参数:
        z (torch.Tensor): 潜在变量，形状为 (batch_size, 3)
        
        返回:
        torch.Tensor: 重构的图像，形状为 (batch_size, 784)
        """
        return self.decoder(z)

    def forward(self, x):
        """
        前向传播函数
        
        参数:
        x (torch.Tensor): 输入图像张量，形状为 (batch_size, 1, 28, 28)
        
        返回:
        tuple: (重构图像, 均值, 对数方差)
            - 重构图像 (torch.Tensor): 形状为 (batch_size, 1, 28, 28)
            - 均值 (torch.Tensor): 形状为 (batch_size, 3)
            - 对数方差 (torch.Tensor): 形状为 (batch_size, 3)
        """
        x = x.view(-1, 28*28)
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar)
        return self.decode(z).view(-1, 1, 28, 28), mu, logvar

# 初始化VAE模型
vae_model = VAE()
vae_optimizer = optim.Adam(
    vae_model.parameters(),
    lr=0.0001
)

def vae_loss(recon_x, x, mu, logvar):
    """
    VAE损失函数：重构损失 + KL散度
    
    参数:
    recon_x (torch.Tensor): 重构的图像，形状为 (batch_size, 784)
    x (torch.Tensor): 原始图像，形状为 (batch_size, 784)
    mu (torch.Tensor): 均值，形状为 (batch_size, 3)
    logvar (torch.Tensor): 对数方差，形状为 (batch_size, 3)
    
    返回:
    torch.Tensor: 标量，表示总损失
    """
    BCE = nn.functional.binary_cross_entropy(recon_x.view(-1, 28*28), x.view(-1, 28*28), reduction='sum')
    KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
    return BCE + KLD

def train(epoch, model, optimizer, criterion, is_vae=False):
    """
    训练函数
    
    参数:
    epoch (int): 当前训练的轮数
    model (nn.Module): 要训练的模型（AE或VAE）
    optimizer (torch.optim.Optimizer): 优化器
    criterion (callable): 损失函数（仅用于AE）
    is_vae (bool): 是否为VAE模型
    
    返回:
    None
    """
    model.train()
    train_loss = 0
    for batch_idx, (data, _) in enumerate(train_loader):
        optimizer.zero_grad()
        if is_vae:
            recon_batch, mu, logvar = model(data)
            loss = vae_loss(recon_batch, data, mu, logvar)
        else:
            recon_batch = model(data)
            loss = criterion(recon_batch, data)
        loss.backward()
        train_loss += loss.item()
        optimizer.step()
    
    print(f'====> Epoch: {epoch} Average loss: {train_loss / len(train_loader.dataset):.4f}')

# 训练AE和VAE模型
num_epochs = 200

def compare_reconstructions(ae_model, vae_model):
    """
    比较AE和VAE模型的重构效果
    
    参数:
    ae_model (nn.Module): 训练好的AE模型
    vae_model (nn.Module): 训练好的VAE模型
    
    返回:
    None (显示图像)
    """
    ae_model.eval()
    vae_model.eval()
    with torch.no_grad():
        data = next(iter(test_loader))[0][:8]  # 获取8个测试样本
        ae_recon = ae_model(data)
        vae_recon, _, _ = vae_model(data)

        # 将原始图像、AE重构和VAE重构拼接在一起
        comparison = torch.cat([data, ae_recon, vae_recon])
        plt.figure(figsize=(12, 4))
        for i in range(24):
            plt.subplot(3, 8, i+1)
            plt.imshow(comparison[i].squeeze().numpy(), cmap='gray')
            plt.axis('off')
        plt.tight_layout()
        plt.show()
        
        # 创建保存路径
        save_dir = "/root/autodl-tmp/projects/diffuser/handwritten_algos/ldm/res_images"
        os.makedirs(save_dir, exist_ok=True)
        
        # 保存图像
        save_path = os.path.join(save_dir, f"{epoch}.png")
        plt.savefig(save_path)
        print(f"=> saved to {save_path}")
        plt.close()  # 关闭图像，防止内存泄漏

for epoch in range(1, num_epochs + 1):
    train(epoch, ae_model, ae_optimizer, ae_criterion)
    train(epoch, vae_model, vae_optimizer, None, is_vae=True)
    # 比较重构效果
    compare_reconstructions(ae_model, vae_model)

改进版本

先看最终效果，再看代码：

可以看出只要训练得当，AE 和 VAE 都能很好的压缩/反压缩图片，重构的效果很棒，肉眼看不出显著的差异！

完整代码如下：

import torch
import os
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import matplotlib.pyplot as plt

# 设置随机种子
torch.manual_seed(42)

# 检查是否有可用的 GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# 数据预处理和加载
transform = transforms.Compose([
    transforms.RandomRotation(10),
    transforms.ToTensor(),
])

train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)

train_loader = DataLoader(train_dataset, batch_size=128*8, num_workers=8, shuffle=True)
test_loader  = DataLoader(test_dataset, batch_size=128*8, num_workers=8, shuffle=False)

class ImprovedAE(nn.Module):
    def __init__(self):
        super(ImprovedAE, self).__init__()
        self.encoder = nn.Sequential(
            nn.Linear(28*28, 256),
            nn.BatchNorm1d(256),
            nn.LeakyReLU(),
            nn.Linear(256, 128),
            nn.BatchNorm1d(128),
            nn.LeakyReLU(),
            nn.Linear(128, 64),
            nn.BatchNorm1d(64),
            nn.LeakyReLU(),
            nn.Linear(64, 16)
        )
        self.decoder = nn.Sequential(
            nn.Linear(16, 64),
            nn.BatchNorm1d(64),
            nn.LeakyReLU(),
            nn.Linear(64, 128),
            nn.BatchNorm1d(128),
            nn.LeakyReLU(),
            nn.Linear(128, 256),
            nn.BatchNorm1d(256),
            nn.LeakyReLU(),
            nn.Linear(256, 28*28),
            nn.Sigmoid()
        )

    def forward(self, x):
        x = x.view(-1, 28*28)
        encoded = self.encoder(x)
        decoded = self.decoder(encoded)
        return decoded.view(-1, 1, 28, 28)

class ImprovedVAE(nn.Module):
    def __init__(self):
        super(ImprovedVAE, self).__init__()
        self.encoder = nn.Sequential(
            nn.Linear(28*28, 256),
            nn.BatchNorm1d(256),
            nn.LeakyReLU(),
            nn.Linear(256, 128),
            nn.BatchNorm1d(128),
            nn.LeakyReLU()
        )
        self.fc_mu = nn.Linear(128, 16)
        self.fc_logvar = nn.Linear(128, 16)
        self.decoder = nn.Sequential(
            nn.Linear(16, 128),
            nn.BatchNorm1d(128),
            nn.LeakyReLU(),
            nn.Linear(128, 256),
            nn.BatchNorm1d(256),
            nn.LeakyReLU(),
            nn.Linear(256, 28*28),
            nn.Sigmoid()
        )

    def encode(self, x):
        h = self.encoder(x)
        return self.fc_mu(h), self.fc_logvar(h)

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

    def decode(self, z):
        return self.decoder(z)

    def forward(self, x):
        x = x.view(-1, 28*28)
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar)
        return self.decode(z).view(-1, 1, 28, 28), mu, logvar

# 初始化改进后的模型并移动到 GPU
ae_model = ImprovedAE().to(device)
vae_model = ImprovedVAE().to(device)

# 使用 AdamW 优化器
ae_optimizer = optim.AdamW(ae_model.parameters(), lr=0.001, weight_decay=1e-5)
vae_optimizer = optim.AdamW(vae_model.parameters(), lr=0.001, weight_decay=1e-5)

# 学习率调度器
ae_scheduler = optim.lr_scheduler.ReduceLROnPlateau(ae_optimizer, patience=10, factor=0.5, verbose=True)
vae_scheduler = optim.lr_scheduler.ReduceLROnPlateau(vae_optimizer, patience=10, factor=0.5, verbose=True)

ae_criterion = nn.MSELoss()

def vae_loss(recon_x, x, mu, logvar, beta=1.0):
    BCE = nn.functional.binary_cross_entropy(recon_x.view(-1, 28*28), x.view(-1, 28*28), reduction='sum')
    KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
    return BCE + beta * KLD

def train(epoch, model, optimizer, scheduler, criterion, is_vae=False):
    model.train()
    train_loss = 0
    for batch_idx, (data, _) in enumerate(train_loader):
        data = data.to(device)  # 将数据移动到 GPU
        optimizer.zero_grad()
        if is_vae:
            recon_batch, mu, logvar = model(data)
            loss = vae_loss(recon_batch, data, mu, logvar, beta=0.5)
        else:
            recon_batch = model(data)
            loss = criterion(recon_batch, data)
        loss.backward()
        train_loss += loss.item()
        optimizer.step()
    
    avg_loss = train_loss / len(train_loader.dataset)
    print(f'====> Epoch: {epoch} Average loss: {avg_loss:.4f}')
    scheduler.step(avg_loss)
    return avg_loss

def compare_reconstructions(ae_model, vae_model, epoch):
    ae_model.eval()
    vae_model.eval()
    with torch.no_grad():
        data = next(iter(test_loader))[0][:8].to(device)  # 将数据移动到 GPU
        ae_recon = ae_model(data)
        vae_recon, _, _ = vae_model(data)

        # 将重构结果移回 CPU 以进行可视化
        data = data.cpu()
        ae_recon = ae_recon.cpu()
        vae_recon = vae_recon.cpu()

        comparison = torch.cat([data, ae_recon, vae_recon])
        plt.figure(figsize=(12, 4))
        for i in range(24):
            plt.subplot(3, 8, i+1)
            plt.imshow(comparison[i].squeeze().numpy(), cmap='gray')
            plt.axis('off')
        plt.tight_layout()
        
        # 创建保存路径
        save_dir = "/root/autodl-tmp/projects/diffuser/handwritten_algos/ldm/res_images_v2"
        os.makedirs(save_dir, exist_ok=True)
        
        # 保存图像
        save_path = os.path.join(save_dir, f"epoch_{epoch}.png")
        plt.savefig(save_path)
        print(f"=> saved to {save_path}")
        plt.close()  # 关闭图像，防止内存泄漏

# 训练循环
num_epochs = 500
for epoch in range(1, num_epochs + 1):
    ae_loss = train(epoch, ae_model, ae_optimizer, ae_scheduler, ae_criterion)
    vae_loss_val = train(epoch, vae_model, vae_optimizer, vae_scheduler, None, is_vae=True)
    
    if epoch % 2 == 0:
        print(f"Epoch {epoch}/{num_epochs}")
        print(f"AE Loss: {ae_loss:.4f}, VAE Loss: {vae_loss_val:.4f}")
        compare_reconstructions(ae_model, vae_model, epoch)
        # save the model
        save_ae_path = '/root/autodl-tmp/projects/diffuser/ckpts/custom/ae_vae/ae.pth'
        save_vae_path = '/root/autodl-tmp/projects/diffuser/ckpts/custom/ae_vae/vae.pth'
        torch.save(ae_model, save_ae_path); print(f"=> saved model to {save_ae_path}")
        torch.save(vae_model, save_vae_path); print(f"=> saved model to {save_vae_path}")

print("Training complete!")