变分自编码器(Variational AutoEncoder)

TL;DR

最近，AIGC是极火热的讨论话题，而文生图可以说是AIGC的代表性工作。目前，效果最好的文生图模型是基于扩散模型的，当进一步深入扩散模型时，又对他的损失函数产生了很大的疑问。通过查找各方资料，才发现扩散模型与变分自编码器在损失定义上同出一门，理解了变分自编码器的损失自然也能理解扩散模型的损失。

另外，变分自编码器已经作为基础模型，集成到许多后续工作中，例如：

1. Stable Diffusion用变分自编码器获取图片的潜在表征(latents)进行前向扩散，避免直接在像素空间中前向扩散，极大地提升了计算效率；
2. 作为变分自编码器的拓展性工作，向量化离散变分自编码器(Vector Quantised-Variational AutoEncoder, VQ-VAE)已经被广泛用作图像分词器，如BEIT、DALL·E等。

可以说，变分自编码器是过不去的一个坎，极有必要对变分自编码器做细致的了解。

但是，查阅已有资料发现，有关变分自编码器的教程总是伴随复杂的公式推导，而实现的代码又难以与公式严格对应。另外,理论部分还涉及变分推断、ELBO、重参数等等多种技巧，让人摸不着头脑。本文将从基本原理入手，逐步介绍变分自编码器的概念、损失函数、推断过程等关键内容，旨在对变分自编码器理论的来龙去脉进行详细的解释，并将推导过程与具体实现相结合，帮助更好地理解变分自编码器。

理论部分

什么是自编码器？：自编码器(AutoEncoder, AE)是一种无监督方式训练的神经网络，主要思想是将高维的输入数据进行编码、压缩，得到低维的特征表示，然后将该特征解码回原始数据，从而学习数据的特征表示。可以用于数据压缩、降维、异常检测、图像去噪等。

如图所示，自编码器包含两个部分：

1. 编码器(Encoder)：将原始高维数据映射到低维隐空间中，以得到低维特征表示；
2. 解码器(Decoder)：低维隐空间中的特征表示作为输入，将其重新映射到原始数据空间，以得到重建数据。

记原始输入数据点为$x$，编码器为$g_{\phi}$，编码后的特征为$z$，解码器为$f_{\theta}$，解码重建后的数据为$x'$，那么就有

$$\begin{aligned} z &= g_{\phi}(x) \\ x' &= f_{\theta}(z) \end{aligned} \tag{1}$$

其中$\phi$和$\theta$分别为编码器$g(\cdot)$和解码器$f(\cdot)$的参数。最终的目标是学习一个恒等映射，即

$$x' \approx f_{\theta}(g_{\phi}(x)) \tag{2}$$

损失可以用$x'$与$x$间的距离度量定义，如熵、MSE等，下面用MSE定义损失

$$L_{AE} (\theta, \phi) = \frac{1}{n} \sum_{i=1}^n (x^{(i)} - f_{\theta}(g_{\phi}(x^{(i)})))^2 \tag{3}$$

自编码器与内容生成：那么训练结束后，获得了编码器、解码器两个网络，除了对原始数据的压缩、降维，是否还可以用来生成数据？比如在隐空间随机取一个特征，用解码器对这个特征进行重构，从而得到新的数据。

这听起来是合理的，但事实上这样做的结果却不尽如人意，原因是：

1. 自编码器的训练目标是重构输入数据，模型规模较大、数据量较小的情况下，能做到一对一的映射，但也引入了过拟合问题；
2. 训练过程中没有对隐空间作任何限制，也就是说隐空间是以任意方式组织的，导致是不连续的，呈现不规则的、无界的分布。

也就是说，隐空间中随机选取特征可能不具有任何实际含义，导致解码后的结果无意义。

变分自编码器如何解决这个问题？：变分自编码器(Variational AutoEncoder)是一种改进的自编码器，目的是使自编码器能应用于内容生成。其思想是：将原始数据编码为隐空间中的概率分布，而不是特定的单个特征，使隐空间具有可采样的特性。

进一步地，为了使隐空间具有可采样的特性，可以令隐变量$z$服从某简单分布（如正态分布），那么可以通过下面步骤采样得到隐层表征，并重构生成数据：

1. 从先验概率$p_{\theta}(z)$中采样，得到特征$z^{(i)}$；
2. 用似然函数$p_{\theta}(x|z=z^{(i)})$重构数据，得到$x'$。

那么，接下来的问题就是如何估计变分自编码器的参数$\theta$。在解决这个问题前，先从贝叶斯模型角度讲解“变分推断”是怎么回事。

从贝叶斯模型谈起：假设输入变量为$x$，隐变量是$z$（在分类问题中即标签$y$，回归问题中就是预测值），那么贝叶斯模型中有

• 先验概率$p(z)$
• 似然函数$p(x|z)$
• 后验概率$p(z|x)$

它们之间的联系可以用贝叶斯公式描述：

$$p(z|x) = \frac{p(x|z) p(z)}{p(x)} \tag{4.1}$$

其中

$$p(x) = \int p(x, z) dz = \int p(x|z) p(z) dz \tag{4.2}$$

其中，$p(z)$和$p(x|z)$可以从数据集估计得到，那么目的就是为了求解后验概率分布$p(z|x)$。将已知项代入上式就能得到结果，但可以看到，$p(z|x) = \frac{p(x|z) p(z)}{\int p(x|z) p(z) dz}$涉及积分计算，这就很难求解了，需要通过近似推断的方法求解，这就引入了变分推断。

“变分”是什么意思？：“变分”来自变分推断(Variational Inference, VI)，是通过引入一个已知分布（如高斯分布）$q(z|x)$来逼近复杂分布$p(z|x)$，设已知分布参数为$\phi$、复杂分布参数为$\theta$，将两个分布记作$q_{\phi}(z|x)$和$p_{\theta}(z|x)$。那么希望两个分布越接近越好，可以用KL散度来度量。

但注意到，KL散度是非对称的：

• $\text{KL}(P||Q) = \mathbb{E}_{z \sim P(z)} \log \frac{P(z)}{Q(z)}$，是指用分布$Q$近似分布$P$，需要保证任意$P(z) > 0$的地方都有$Q(z) > 0$，结果是$Q$的分布会覆盖整个$P$的分布；
• $\text{KL}(Q||P) = \mathbb{E}_{z \sim Q(z)} \log \frac{Q(z)}{P(z)}$，是指用分布$P$近似分布$Q$，当$P(z) \rightarrow 0$时一定有$Q(z) \rightarrow 0$，结果是使$Q$逼近$P$的其中一个峰。

在变分推断中，一般用反向KL散度，即

$$\begin{aligned} \phi^* &= \arg \min_{\phi} \text{KL}(q_{\phi}(z|x) || p_{\theta}(z|x)) \\ &= \arg \min_{\phi} \mathbb{E}_{z \sim q_{\phi}(z|x)} \log \frac{q_{\phi}(z|x)}{p_{\theta}(z|x)} \end{aligned} \tag{5}$$

其中$p_{\theta}(z|x)$未知，需要经过一系列变换才能进行优化。

变分推断与ELBO：对上式进行变换，由贝叶斯公式有$p_{\theta}(z|x) = \frac{p_{\theta}(x|z) p_{\theta}(z)}{p_{\theta}(x)}$，代入可以得到

$$\begin{aligned} \text{KL}(q_{\phi}(z|x) || p_{\theta}(z|x)) &= \mathbb{E}_{z \sim q_{\phi}(z|x)} \log \frac{q_{\phi}(z|x) p_{\theta}(x)}{p_{\theta}(x|z) p_{\theta}(z)} \\ &= \mathbb{E}_{z \sim q_{\phi}(z|x)} \log \frac{q_{\phi}(z|x)}{p_{\theta}(x|z) p_{\theta}(z)} + \log p_{\theta}(x) & \scriptstyle{\mathbb{E}_{z \sim q_{\phi}(z|x)} \log p_{\theta}(x) = \log p_{\theta}(x)}\\ &= \mathbb{E}_{z \sim q_{\phi}(z|x)} \left( \log \frac{q_{\phi}(z|x)}{p_{\theta}(z)} - \log p_{\theta}(x|z) \right) + \log p_{\theta}(x) \\ &= \text{KL}(q_{\phi}(z|x)||p_{\theta}(z)) - \mathbb{E}_{z \sim q_{\phi}(z|x)}\log p_{\theta}(x|z) + \log p_{\theta}(x) \\ \end{aligned} \tag{6}$$

多项式移项整理后，可以得到

$$\log p_{\theta}(x) = \text{KL}(q_{\phi}(z|x) || p_{\theta}(z|x)) - \text{KL}(q_{\phi}(z|x)||p_{\theta}(z)) + \mathbb{E}_{z \sim q_{\phi}(z|x)}\log p_{\theta}(x|z) \tag{7}$$

由于KL散度非负，即$\text{KL}(q_{\phi}(z|x) || p_{\theta}(z|x)) \geq 0$，因此

$$\log p_{\theta}(x) \geq - \text{KL}(q_{\phi}(z|x)||p_{\theta}(z)) + \mathbb{E}_{z \sim q_{\phi}(z|x)}\log p_{\theta}(x|z) \tag{8}$$

右边多项式可以视作$\log p_{\theta}(x)$的下界，或称证据变量$x$的下界，定义为证据下界(Evidence Lower Bound, ELBO)，即

$$-L_{\text{VI}} = - \text{KL}(q_{\phi}(z|x)||p_{\theta}(z)) + \mathbb{E}_{z \sim q_{\phi}(z|x)}\log p_{\theta}(x|z) \tag{9}$$

那么优化目标就可以进行转换，即

$$\phi^* = \arg \min_{\phi} \text{KL}(q_{\phi}(z|x) || p_{\theta}(z|x)) = \arg \min_{\phi} L_{\text{VI}} \tag{10}$$

回到变分自编码器：VAE的训练目标定义为最大化真实数据的概率分布，也即

$$\begin{aligned} \theta^* &= \arg \max_{\theta} \prod_{i=1}^n p_{\theta} (x^{(i)}) \\ &= \arg \max_{\theta} \sum_{i=1}^n \log p_{\theta} (x^{(i)}) \\ \end{aligned} \tag{11}$$

上面提到，用贝叶斯公式直接展开上式，会引入积分项导致难以求解。而由式$(8)$又可知，$(-L_{VI})$是$\log p_{\theta} (x)$的一个下界，那么通过最大化下界，可以间接地最大化$\log p_{\theta} (x)$，也就是

$$\theta^*, \phi^* = \arg \max_{\theta, \phi} \sum_{i=1}^n - \text{KL}(q_{\phi}(z^{(i)}|x^{(i)})||p_{\theta}(z^{(i)})) + \mathbb{E}_{z \sim q_{\phi}(z|x^{(i)})}\log p_{\theta}(x^{(i)}|z) \tag{12}$$

通常最小化损失，因此记变分自编码器的损失为

$$L_{\text{VAE}} = \frac{1}{n} \sum_{i=1}^n - \mathbb{E}_{z \sim q_{\phi}(z|x^{(i)})}\log p_{\theta}(x^{(i)}|z) + \text{KL}(q_{\phi}(z^{(i)}|x^{(i)})||p_{\theta}(z^{(i)})) \tag{13}$$

其中，$q_{\phi}(z|x)$是编码器部分，$p_{\theta}(x|z)$是解码器部分，$p_{\theta}(z)$是期望的令$z$服从的已知简单分布（如正态分布、均匀分布等）。

损失的具体形式：写到这里，已经完成了形式化的损失函数定义，许多教程在这里就结束了。但阅读一些具体实现的代码，发现损失如式$(14)$所示，很难将其联系到式$(13)$上：

$$L_{\text{VAE}} = \frac{1}{n} \sum_{i=1}^n ||x^{(i)} - x'^{(i)}||^2 + \frac{1}{2} ||\mu^{(i)2} + \sigma^{(i)2} - \log \sigma^{(i)2} - 1||^2 \tag{14}$$

其中$x^{(i)}$是样本点，$x'^{(i)}$是重构后的样本点。上面引入近似分布（也即编码器）$q_{\phi}(z|x)$是高斯分布，即$q_{\phi}(z^{(i)}|x^{(i)}) \sim \mathcal{N}(\mu^{(i)}, \sigma^{(i)2}I)$，$\mu^{(i)}$和$\sigma^{(i)2}$表示$x^{(i)}$输入对应的均值、方差。

接下来说明，如何从式$(13)$得到$(14)$。

形式化损失与具体损失的联系：回到式$(13)$，我们可以将其拆分为重构损失、正则项损失两部分：

$$\begin{cases} L_{\text{recon}} &= \frac{1}{n} \sum_{i=1}^n - \mathbb{E}_{z \sim q_{\phi}(z|x^{(i)})}\log p_{\theta}(x^{(i)}|z) \\ L_{\text{regu}} &= \frac{1}{n} \sum_{i=1}^n \text{KL}(q_{\phi}(z^{(i)}|x^{(i)})||p_{\theta}(z^{(i)})) \end{cases} \tag{15}$$

其中：

• $z \sim q_{\phi}(z|x^{(i)})$表示采样过程，涉及到重参数技巧；
• $L_{\text{recon}}$是重构损失，与自编码器一致，$L_{\text{regu}}$是正则项损失，目的是更好地组织隐空间，使其具有可采样的特性，并防止过拟合；
• 注意到这两项是相互对抗的，因为最小化$L_{\text{regu}}$使$\text{KL}(q_{\phi}(z^{(i)}|x^{(i)})||p_{\theta}(z^{(i)})) = 0$时，$z$就没有了任何差异，这样重建准确率就很低，导致$L_{\text{recon}}$很高，因此最终目的是达到两项的平衡状态。

再看式$(15)$中各项概率分布：

• $p_{\theta}(z)$：为了方便采样，一般令$z \sim \mathcal{N}(0, I)$，这是人为指定的；
• $q_{\phi}(z|x)$：编码器部分，前面变分推断部分已经提到，用高斯分布拟合，得到$\mathcal{N}(\mu, \sigma^2 I)$；
• $p_{\theta}(x|z)$：解码器部分，还没定，也可以选择一个简单分布拟合，如伯努利分布或者高斯分布。

当$p_{\theta}(x|z)$采用伯努利分布，即多元二项分布，有

$$p_{\theta}(x|z) = \prod_{k=1}^{d} p_{\theta}(z_k)^{x_{k}} (1 - p_{\theta}(z_k))^{1 - x_{k}} \tag{16.1}$$

其中$d$表示随机变量$x$的维度，此时$x_k \in \{ 0, 1 \}, k = 1, \cdots, d$，那么

$$\begin{aligned} L_{\text{recon}} &= \frac{1}{n} \sum_{i=1}^n - \mathbb{E}_{z \sim q_{\phi}(z|x^{(i)})}\log p_{\theta}(x^{(i)}|z) \\ &= \frac{1}{n} \sum_{i=1}^n \log \left( - \prod_{k=1}^{d} p_{\theta}(z^{(i)}_k)^{x^{(i)}_k} (1 - p_{\theta}(z^{(i)}_k))^{1 - x^{(i)}_k} \right) \\ &= \frac{1}{n} \sum_{i=1}^n \sum_{k=1}^{d} \left( - x^{(i)}_k \log p_{\theta}(z^{(i)}_k) - (1 - x^{(i)}_k) \log (1 - p_{\theta}(z^{(i)}_k)) \right) \end{aligned} \tag{16.2}$$

此时用二元交叉熵作为损失函数。

当$p_{\theta}(x|z)$采用高斯分布，回顾多维高斯分布：若随机变量$x \sim \mathcal{N}(\mu, \Sigma)$，有

$$p(x) = \frac{1}{(2\pi)^{d/2} |\Sigma|^{1/2}} \exp \left[ - \frac{1}{2} (x - \mu)^T \Sigma^{-1} (x - \mu) \right] \tag{17.1}$$

很容易得到$p_{\theta}(x^{(i)}|z)$的表达式，进一步地，简化假设各分量独立（即$\Sigma$为对角阵$\sigma^2 I$），$\mu$为关于$z$的函数，那么

$$\begin{aligned} L_{\text{recon}} &= \frac{1}{n} \sum_{i=1}^n - \mathbb{E}_{z \sim q_{\phi}(z|x^{(i)})}\log p_{\theta}(x^{(i)}|z) \\ &= \frac{1}{n} \sum_{i=1}^n \log \left( - \frac{1}{\prod_{k=1}^d \sqrt{(2 \pi)^d \sigma_k^2(z^{(i)})}} \exp \left( - \frac{1}{2} ||\frac{x^{(i)} - \mu(z^{(i)})}{\sigma(z^{(i)})}||^2 \right) \right) \\ &= \frac{1}{n} \sum_{i=1}^n \left( \frac{1}{2} ||\frac{x^{(i)} - \mu(z^{(i)})}{\sigma(z^{(i)})}||^2 + \frac{1}{2} \sum_{k=1}^d \log (2 \pi)^d \sigma_k^2(z^{(i)}) \right) \\ &= \frac{1}{n} \sum_{i=1}^n \left( \frac{1}{2} ||\frac{x^{(i)} - \mu(z^{(i)})}{\sigma(z^{(i)})}||^2 + \frac{d}{2} \sum_{k=1}^d \log 2 \pi + \frac{1}{2} \sum_{k=1}^d \sigma_k^2(z^{(i)}) \right) \end{aligned} \tag{17.2}$$

为简化计算，令方差项$\sigma(z)$为常数$c$，损失可以简化为MSE损失：

$$L_{\text{recon}} = \frac{1}{n} \sum_{i=1}^n \frac{1}{2c} ||x^{(i)} - \mu_{\theta}(z^{(i)})||^2 \cancel{+ C} \tag{17.3}$$

注意到，$\mu_{\theta}(z^{(i)})$即重构的数据$x'^{(i)}$。

再看正则项损失，有

$$\begin{cases} q_{\phi}(z^{(i)}|x^{(i)}) &= \frac{1}{ \prod_{k=1}^h \sqrt{(2 \pi)^h \sigma_k^2(x^{(i)})} } \exp \left( - \frac{1}{2} ||\frac{z^{(i)} - \mu(x^{(i)})}{\sigma(x^{(i)})}||^2 \right) \\ p_{\theta}(z^{(i)}) &= \frac{1}{ \prod_{k=1}^h \sqrt{(2 \pi)^h} } \exp \left( - \frac{1}{2} ||z^{(i)}||^2 \right) \\ \end{cases} \tag{18.1}$$

$$\begin{aligned} L_{\text{regu}} &= \frac{1}{n} \sum_{i=1}^n \text{KL}(q_{\phi}(z^{(i)}|x^{(i)})||p_{\theta}(z^{(i)})) \\ &= \frac{1}{n} \sum_{i=1}^n \int q_{\phi}(z^{(i)}|x^{(i)}) \log \frac{ q_{\phi}(z^{(i)}|x^{(i)}) }{ p_{\theta}(z^{(i)}) } d z^{(i)} \\ &= \cdots & \scriptstyle{20.1式代入计算，略} \\ &= \frac{1}{n} \sum_{i=1}^n \frac{1}{2} ||\mu^2(x^{(i)}) + \sigma^2(x^{(i)}) - \log \sigma^2(x^{(i)}) - 1||^2 \end{aligned} \tag{18.2}$$

也即

$$L_{\text{regu}} = \frac{1}{n} \sum_{i=1}^n \frac{1}{2} ||\mu^{(i)2} + \sigma^{(i)2} - \log \sigma^{(i)2} - 1||^2 \tag{18.3}$$

实现细节

编码器与解码器网络：变分推断中提到用高斯分布来逼近$p_{\theta}(z|x)$，也就是说希望编码器$q_{\phi}(z|x)$输出高斯概率分布。直接令神经网络$g_{\phi}(x)$拟合分布参数$\mu$和$\sigma^2$（考虑到$\sigma^2$非负，一般用$\log \sigma^2$），那么有

$$\mu, \log \sigma^2 = g_{\phi}(x) \tag{19.1}$$

解码器部分就比较简单了，只要将采样得到的$z$重建，同样用神经网络$f_{\theta}(z)$表示，也就是

$$x' = f_{\theta}(z) \tag{19.2}$$

隐层特征$z$的采样：目前，已经令编码器得到分布$\mathcal{N}(\mu^{(i)}, \sigma^{(i)2} I)$了，那么如何得到隐层特征$z^{(i)}$呢？能够直接从分布中采样得到呢？答案是不可以，因为采样操作是不可导的，导致最终误差无法通过网络反传到编码器实现参数更新。

解决方法是采用重参数技巧(Reparameterization Trick)，希望从正态分布$\mathcal{N}(\mu, \sigma^2 I)$中采样，可以先从标准正态分布$\mathcal{N}(0, I)$中采样$\epsilon$，然后用以下变换得到$z$（由正态分布性质可证）：

$$z = \mu \epsilon + \sigma \tag{20}$$

这样做，就可以把不可导的采样操作移除到梯度计算图之外，实现误差反传。

具体实现：下面是在MNIST数据集上进实现的的变分自编码器

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

# 定义变分自编码器模型
class VAE(nn.Module):
    def __init__(self, input_size, hidden_size, latent_size):
        super(VAE, self).__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.latent_size = latent_size
        
        self.encoder = nn.Sequential(
            nn.Linear(self.input_size, self.hidden_size),
            nn.ReLU(),
            nn.Linear(self.hidden_size, self.hidden_size),
            nn.ReLU()
        )
        
        self.mean = nn.Linear(self.hidden_size, self.latent_size)
        self.logvar = nn.Linear(self.hidden_size, self.latent_size)
        
        self.decoder = nn.Sequential(
            nn.Linear(self.latent_size, self.hidden_size),
            nn.ReLU(),
            nn.Linear(self.hidden_size, self.hidden_size),
            nn.ReLU(),
            nn.Linear(self.hidden_size, self.input_size),
            nn.Sigmoid()
        )
        
    def encode(self, x):
        h = self.encoder(x)
        mean = self.mean(h)
        logvar = self.logvar(h)
        return mean, logvar
    
    def reparameterize(self, mean, logvar):
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        z = mean + eps * std
        return z
    
    def decode(self, z):
        x_hat = self.decoder(z)
        return x_hat
    
    def forward(self, x):
        mean, logvar = self.encode(x)
        z = self.reparameterize(mean, logvar)
        x_hat = self.decode(z)
        return x_hat, mean, logvar

# 定义训练函数
def train(model, dataloader, optimizer, criterion, device):
    model.train()
    train_loss = 0
    for batch_idx, (data, _) in enumerate(dataloader):
        data = data.view(data.size(0), -1)
        data = data.to(device)
        optimizer.zero_grad()
        recon_batch, mu, logvar = model(data)
        loss = criterion(recon_batch, data, mu, logvar)
        loss.backward()
        train_loss += loss.item()
        optimizer.step()
    return train_loss / len(dataloader.dataset)

# 定义测试函数
@torch.no_grad()
def test(model, dataloader, criterion, device):
    model.eval()
    test_loss = 0
    for data, _ in dataloader:
        data = data.view(data.size(0), -1)
        data = data.to(device)
        recon_batch, mu, logvar = model(data)
        test_loss += criterion(recon_batch, data, mu, logvar).item()
    return test_loss / len(dataloader.dataset)

# 定义损失函数
def loss_fn(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

if __name__ == "__main__":
    # 加载数据集
    batch_size = 128
    train_dataset = datasets.MNIST(root='./data', train=True, transform=transforms.ToTensor(), download=True)
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    test_dataset = datasets.MNIST(root='./data', train=False, transform=transforms.ToTensor(), download=True)
    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=True)

    # 初始化模型和优化器
    input_size = 784
    hidden_size = 256
    latent_size = 20
    model = VAE(input_size, hidden_size, latent_size).to('cuda')
    optimizer = optim.Adam(model.parameters(), lr=1e-3)

    # 训练模型
    epochs = 10
    for epoch in range(1, epochs+1):
        train_loss = train(model, train_loader, optimizer, loss_fn, 'cuda')
        test_loss = test(model, test_loader, loss_fn, 'cuda')
        print('Epoch {}: Train Loss {:.4f}, Test Loss {:.4f}'.format(epoch, train_loss, test_loss))

    torch.save(model.state_dict(), 'vae.pth')

可以用下面代码进行推断

import torch
from torchvision.utils import save_image
from vae import VAE

# 加载VAE模型
input_size = 784
hidden_size = 256
latent_size = 20

vae = VAE(input_size, hidden_size, latent_size).to('cuda')
vae.load_state_dict(torch.load('vae.pth'))
vae.eval()

# 从标准正态分布中采样潜在向量
z = torch.randn(64, latent_size)

# 生成新的样本
with torch.no_grad():
    z = z.to("cuda")
    x_hat = vae.decode(z)

# 将生成的样本保存到文件中
save_image(x_hat.view(64, 1, 28, 28), 'generated_samples.png')

可以多训练几轮，达到更好的效果

参考资料

• Auto-Encoding Variational Bayes - arXiv.org
• 半小时理解变分自编码器 - 知乎
• 变分自编码器VAE：原来是这么一回事 | 附开源代码 - 知乎
• 变分推断(Variational Inference)初探 - 知乎
• 变分推断（Variational Inference）- 知乎
• 一文解释 VAE+ELBO - 知乎
• 变分自编码器（一）：原来是这么一回事
• 变分自编码器（二）：从贝叶斯观点触发
• From Autoencoder to Beta-VAE - Lil’Log
• A Beginner’s Guide to Variational Methods: Mean-Field Approximation - Eric Jang