变分自编码器（VAE）初识

AE回顾

Auto-Encoder，称自编码器，是一种无监督式学习模型。它基于反向传播算法与最优化方法（如梯度下降法），AE（Auto-Encoder）的架构可以如下所示;

记$X$为整个数据集的集合，$x_i$是数据集中的一个样本。

自编码器包含一个编码器$z=g(X)$，称$z$为编码，并且往往$z$的维度远小于输入$X$的维度。此外还包含一个解码器$\tilde{X}=f(z)$,这个解码器从编码$z$中还原输入$X$的信息。

我们自然希望$\tilde{X}$与$X$尽可能地接近，因此我们可以定义自编码器的损失函数为 ℓ = ∥ X − X ~ ∥ \ell=\begin{Vmatrix} X-\tilde{X} \end{Vmatrix} ℓ= ​X−X~​ ​，于是模型训练结束后，我们便可以认为$z$蕴含了输入$X$的大部分信息，则能够表达原始数据以实现数据降维的目的。

为此，我们可以将自编码器简记为：
$X\in {\mathbb{R}}^{C\times H\times W}⟶z=g(X)\in {\mathbb{R}}^d⟶\tilde{X}=f(z)\in {\mathbb{R}}^{C\times H\times W}$其中$C$表示通道数（彩色图片为RGB三通道），$H,W$分别表示图片数据的高和宽。

VAE

在AE那部分，可能有读者会想：能不能够直接从编码器$z=g(X),z\in {\mathbb{R}}^d$直接采样$z_i$，再通过解码器${\hat{X}}_i=f(z_i)$生成图片呢？回答是很难，准确地来说一定概率下是可行的，但大概率会生成全是噪声的图片。这是因为${\mathbb{R}}^d$是一个很大的空间，而符合条件的$z$也许只占很小一分部，所以如果在整个${\mathbb{R}}^d$空间上采样的话，自然难以生成符合预期的图片。也就是说由采样$z_i$生成的${\hat{X}}_i=f(z_i)$可能会和$X_i$相距太远，但如果我们能够显式地对$z$的分布$p(z)$进行建模，从而只在$p(z)$中采样$z$,这便就是VAE（Variational Auto-Encoder），即变分自编码器。

VAE简述

不妨假设隐空间$z\sim \mathcal{N}(0,I)$，其中$I$代表一个单位矩阵。记$X$为随机变量，$X_i$代表随机变量的样本。$z_i$是从隐空间$z$采样得到的样本（注：在VAE模型中并没有使用$p(z)$（隐变量空间的分布）是标准正态分布的假设，实际上用的是假设$p(z|X)$（后验分布）是服从正态分布的）

具体来说，给定真实样本$x_k$，假设存在一个服从多元正态分布的后验分布$p(z|x_k)$，且后验分布$p(z|x_k)$对于不同的$k$是相互独立的。事实上，在论文《Auto-Encoding Variational Bayes》中，也特别强调了这一点：

In this case, we can let the variational approximate posterior be a multivariate Gaussian with a diagonal covariance structure:$$\log q_{\varphi }(\mathbf{z}|{\mathbf{x}}^{(i)})=\log \mathcal{N}(\mathbf{z};{\mathbf{\mu }}^{(i)},{\mathbf{\sigma }}^{2(i)}\mathbf{I})（9）$$

上式是实现VAE的关键，那么要如何找出后验分布$p(z|x_i)$（已经假设是服从多元正态分布了）
的均值和方差呢？实际上我们可以通过从后验分布$p(z|x_i)$采样的方式，然后利用神经网络拟合出来，设${\mu }_i=f_1(X_i)$，$log{\sigma }_i^2=f_2(X_i)$，其中${\mu }_i$，${\sigma }_i$分别表示后验分布$p(z|x_i)$的均值和方差。这里要说明的是，选择拟合$log{\sigma }_i^2$是因为${\sigma }_i^2$恒正，而$log{\sigma }_i^2$可正可负，这样就不用多添加一层激活函数层了。现在从已知的后验分布$p(z|x_i)$中采样$z_i$，通过解码器得到${\hat{x}}_k=f(z_k)$，如果VAE训练得足够好（也就是$\mathcal{D}(x,{\hat{x}}_k)$足够小，$\mathcal{D}$表示距离函数），我们可以认为${\hat{x}}_k$是与$x_k$都来源于数据集$X$的分布$p(X)$。

上述的过程我们表示为：设数据集是由某个随机过程生成的，而$z$是随机过程中的一个不可观测到的隐空间。这个生成数据的随机过程包含两个步骤：

首先从先验分布$p(z)$中采样得到一个$z_i$,再从条件分布$p(X|z_i)$中采样得到一个数据点$x_i$，如果能基于这个随机过程进行建模，那么便能够得到一个生成模型。

VAE中的Decoder

下面我们从生成模型的角度来看看VAE中的Decoder架构：

上图是VAE中的Decoder架构的示意图。现在给Decoder一个从$\mathcal{N}(0,1)$采样得到的$z_i$，我们希望由$\theta$参数化的Decoder能够输出$z_i$对应的$X$的分布，即$f_{\theta }(z_i)=p_{\theta }(X|z_i)$

假设给定任意$z_i$后，$X$都服从某个各维度独立的多元高斯分布，即：
$$\begin{array}{r}{p}_{\theta }(X|z_i)=\mathcal{N}(X|{\mu }_i^{{}^{\prime }}(z_i;\theta ),{\sigma }_i^{{}^{\prime }2}(z_i;\theta )\ast I)\end{array}$$于是只要输入$z_i$给Decoder，然后让它拟合出${\mu }_i^{{}^{\prime }}$和${\sigma }_i^{{}^{\prime }2}$，就能知道$X|z_i$的具体分布了。

如果所有的$p(z|x_k)$都很接近标准正态分布$\mathcal{N}(0,1)$，那么根据定义
$$p(Z)=\sum Xp(Z|X)p(X)=\sum XN(0,I)p(X)=N(0,I)\sum Xp(X)=N(0,I)(2)$$

这样我们就能达到我们的先验假设：p(Z)
是标准正态分布。然后我们就可以放心地从N(0,I)
中采样来生成图像了。

VAE的目标函数

下面我们从统计的视角来看看一个生成模型的目标函数应该是怎么样的。采用李宏毅老师的课件图，生成模型的目标可以如下图所示：（图中的$x$和$z$和本文VAE中的定义不完全一致）

生成模型是目标对数据集的分布$p(X)$建模得到$\theta$参数化的分布$p_{\theta }(X)$，从分布$p_{\theta }(X)$中采样进而得到生成的数据。如上图所示，$X$代表一个图片的集合，得到$p_{\theta }(X)$采样一个令$p_{\theta }(x_i)$较大的$x_i$，那么$x_i$很可能就是和数据集合$X$中的图片相似的一张新图片。

根据前面的内容，我们可以如下计算出$p_{\theta }(X)$:已知$$\begin{array}{r}p(z)\sim \mathcal{N}(0,I),p_{\theta }(X|z_i)=\mathcal{N}(X|{\mu }_i^{{}^{\prime }}(z_i;\theta ),{\sigma }_i^{{}^{\prime }2}(z_i;\theta )\ast I)\end{array}$$那么有：$$\begin{array}{r}{p}_{\theta }(X)={\int }_z{p}_{\theta }(X|z)p(z)dz\approx \frac{1}{m}{\sum }_{j=1}^m{p}_{\theta }(X|z_j)\end{array}$$利用极大似然估计(MLE)，为了让数据集出现的概率最大化，也就是：
$$\begin{array}{rl}{\theta }^{\ast }& =argmi{n}_{\theta }-\sum _{i=1}^{n}log(p_{\theta }(x_i))\\ & =argmi{n}_{\theta }-\sum _{i=1}^{n}log({\sum }_{j=1}^m{p}_{\theta }(x_i|z_j))\end{array}$$因为往往$x_i$的维度会很大（比如RGB图片256*256），$z_j$的维度也不会很低（$z$是编码器得到的语义向量，比如一张图片可以抽取十个特征），并且对于某个$x_i$而言，与之强相关的$z_j$的数量是相对有限的，但是为了找到这些有限的$z_j$，可能要进行大量的采样,而这在现实应用的意义下是几乎不可能的（这将会是一件computation costly的事情），所以我们希望能够以较小的代价采样得到更多强相关的$z_j$，解决这个需求的办法便是在Encoder中引入后验分布$p_{\theta }(z|x_i)$

VAE中的Encoder

如下图所示：

前一部分（$X\to Z$）可以看成是Encoder的架构，假设有后验分布$p_{\theta }(z|X)$，正向过程在计算的时候将$x_i$传给Encoder，算出$z|x_i$服从的分布后再从这个分布中采样出$z_j$，再将采样得到的$z_j$传给Decoder，然后便能够得到$X|z_i$的分布，然后再用极大似然估计使得$p(X|z_i)$最大化。如此采样得到的$z_j$几乎是和$x_i$相关的，这样便能保证采样的效率。

上述过程可以由贝叶斯公式计算得出：$$\begin{array}{rl}{p}_{\theta }(z|x_i)& =\frac{p_{\theta }(x_i|z)p(z)}{p_{\theta }(x_i)}\\ & =\frac{p_{\theta }(x_i|z)p(z)}{{\int }_{z^{\prime }}{p}_{\theta }(X|z^{\prime })p(z^{\prime })d{z}^{\prime }}\end{array}$$我们依然可以通过采样大量的$z_i$来计算上式分母的积分，但这同样是一个computation costly的工作，这时便能够应用变分贝叶斯：设由$\varphi$参数化的Encoder能够拟合对任意$x_i$的分布$q_{\varphi }(z|x_i)$,如果这个分布$q_{\varphi }(z|x_i)$能足够逼近真实的后验分布$p_{\theta }(z|x_i)$的话，便能够直接由Encoder得到$z|x_i$的分布。由于分布$q_{\varphi }(z|x_i)$是多元高斯分布，因此只需要让Encoder输出参数$\mu$和${\sum }^2$即可。同时由前面的假设[公式(2)]，由贝叶斯公式以及条件概率公式有：$$\begin{array}{r}{p}_{\theta }(z|X)=\frac{p(z)p_{\theta }(X|z)}{p(X)}\end{array}$$于是后验分布$p_{\theta }(z|X)$也是多元高斯分布。假设近似后验分布对于任意的$x_i$有：
$$\begin{array}{r}{q}_{\varphi }(X|z_i)=\mathcal{N}(X|{\mu }_i^{}(z_i;\varphi ),{\sigma }_i^2(z_i;\varphi )\ast I)\end{array}$$，即近似后验分布$q_{\varphi }(z|x_i)$也是一个各维度独立的多元正态分布。

VAE的架构

下图是VAE的架构图（基于MLP模型），其中$x_i^{(j)}$表示第$i$个数据点的第$j$个特征。

由上图先给出VAE架构总结：
1.首先给Encoder输入一个数据点$x_i$，通过神经网络，得到隐变量$z$服从的近似后验分布$q_{\varphi }(z|x_i)$的参数。Encoder输出服从的高斯分布的参数 ${\sigma }_i^2$和${\mu }_i$
2.从高斯分布$q_{\varphi }(z|x_i)$中采样$z_j$，这个$z_j$应当代表与$x_i$相似的一类样本。
3.令Decoder拟合似然的分布$p_{\theta }(X|z_j)$。输入Decoder一个$z_j$，它应当返回服从$X|z_j$的分布的参数。令Decoder输出服从的高斯分布的参数${\sigma }_i^{{}^{\prime }2}$和${\mu }_i{}^{\prime }$即可。
4.在得到$X|z_j$的分布的参数后，从这个分布中进行采样，来生成可能的数据点x_{i}。

注意：大部分实现中，人们往往不进行采样，而是直接将模型输出的${\mu }_i{}^{\prime }$当作是给定生成的数据点。除此之外，人们也往往认为是一个固定方差的各维度独立的多元高斯分布，即${\sigma }_i^{{}^{\prime }2}$是一个人为给定的超参数。

可以看到，我们并没有用到$p(Z)$是正态分布的假设。

参数重整化

为了让网络能够进行训练（能够反向传播和正向传播以计算梯度），利用参数重整化的技巧，替换步骤2中部分：从高斯分布$q_{\varphi }(z|x_i)$中采样$z_j$的服从$z\sim \mathcal{N}({\mu }_i，{\sigma }_i^2)$,令${ϵ}_i\sim \mathcal{N}(0,I)$,$z={\mu }_i+{\sigma }_i\odot {ϵ}_i$,其中$\odot$表示点积，则$z$是服从$\mathcal{N}({\mu }_i,{\sigma }_i^2\ast I)$的各维度独立的多元高斯分布，并且z是可以求梯度的。那么VAE的架构可以如图表示：

VAE的损失函数

利用对数极大似然估计，最大化$log{p}_{\theta }(X)$时，利用变分推断的思想，得到如下结果：$$\begin{array}{rl}\log p_{\theta }(X)& ={\int }_z{q}_{\varphi }(z\mid X)\log p_{\theta }(X)dz\text{全概率公式}\\ & ={\int }_z{q}_{\varphi }(z\mid X)\log \frac{p_{\theta }(X,z)}{p_{\theta }(z\mid X)}dz\text{贝叶斯公式}\\ & ={\int }_z{q}_{\varphi }(z\mid X)\log (\frac{p_{\theta }(X,z)}{q_{\varphi }(z\mid X)}\cdot \frac{q_{\varphi }(z\mid X)}{p_{\theta }(z\mid X)})dz\text{恒等变换}\\ & ={\int }_z{q}_{\varphi }(z\mid X)\log \frac{p_{\theta }(X,z)}{q_{\varphi }(z\mid X)}dz+{\int }_z{q}_{\varphi }(z\mid X)\log \frac{q_{\varphi }(z\mid X)}{p_{\theta }(z\mid X)}dz\text{拆开log}\\ & =\ell \left(p_{\theta },q_{\varphi }\right)+{\mathcal{D}}_{KL}\left(q_{\varphi },p_{\theta }\right)\\ & \ge \ell \left(p_{\theta },q_{\varphi }\right)KL\text{散度非负}.\end{array}$$

还记得我们在Pytorch中说到的关于交叉熵以及KL散度的一些概念吗？我们在那说明了KL散度的非负性。上式中$\ell \left(p_{\theta },q_{\varphi }\right)$显然是一个下界，这也是称$\ell$为ELBO（Evidence Lower Bound Objection）的原因。变换上述式子可以得到：$$\ell \left(p_{\theta },q_{\varphi }\right)=log{p}_{\theta }(X)-{\mathcal{D}}_{KL}\left(q_{\varphi },p_{\theta }\right)$$，于是最小化$\ell \left(p_{\theta },q_{\varphi }\right)$等价于最大化$log{p}_{\theta }(X)$
，同时最小化${\mathcal{D}}_{KL}\left(q_{\varphi },p_{\theta }\right)$，也就是让$q_{\varphi },p_{\theta }$足够接近。其中$\ell \left(p_{\theta },q_{\varphi }\right)$代表近似后验分布$q_{\varphi }$和真实后验分布$p_{\theta }$之间的损失。

展开$\ell \left(p_{\theta },q_{\varphi }\right)$得到：$$\begin{array}{rl}\ell \left(p_{\theta },q_{\varphi }\right)& ={\int }_z{q}_{\varphi }(z\mid X)\log \frac{p_{\theta }(X,z)}{q_{\varphi }(z\mid X)}dz\\ & ={\int }_z{q}_{\varphi }(z\mid X)\log \frac{p_{\theta }(X\mid z)p(z)}{q_{\varphi }(z\mid X)}dz\text{贝叶斯公式}\\ & ={\int }_z{q}_{\varphi }(z\mid X)\log \frac{p(z)}{q_{\varphi }(z\mid X)}dz+{\int }_z{q}_{\varphi }(z\mid X)\log p_{\theta }(X\mid z)dz\\ & =-{\mathcal{D}}_{KL}\left(q_{\varphi },p\right)+{\mathbb{E}}_{q_{\varphi }}\left[\log p_{\theta }(X\mid z)\right].\end{array}$$，第一项$-{\mathcal{D}}_{KL}\left(q_{\varphi },p\right)$为正则项，人们称为Latent loss，由于我们已经在前文假设了$q_{\varphi }(z|X)$和$p(z)$都服从高斯分布(公式(2)和公式(7))，且都是各维度独立的高斯分布，结合KL散度的性质可以得到其解析解，以一维情况为例：$$\begin{array}{rl}\left(\mathcal{N}\left(\mu ,{\sigma }^2\right)\Vert \mathcal{N}(0,1)\right)& ={\int }_z\frac{1}{\sqrt{2\pi {\sigma }^2}}\exp \left(-\frac{{\left(z-\mu \right)}^2}{2{\sigma }^2}\right)\log \frac{\frac{1}{\sqrt{2\pi {\sigma }^2}}\exp \left(-\frac{{\left(z-\mu \right)}^2}{2{\sigma }^2}\right)}{\frac{1}{\sqrt{2\pi }}\exp \left(-\frac{z^2}{2}\right)}dz\\ & ={\int }_z\left(\frac{-(z-\mu {)}^2}{2{\sigma }^2}+\frac{z^2}{2}-\log \sigma \right)\mathcal{N}\left(\mu ,{\sigma }^2\right)dz\\ & =-{\int }_z\frac{{\left(z-\mu \right)}^2}{2{\sigma }^2}\mathcal{N}\left(\mu ,{\sigma }^2\right)dz+{\int }_z\frac{z^2}{2}\mathcal{N}\left(\mu ,{\sigma }^2\right)dz-{\int }_z\log \sigma \mathcal{N}\left(\mu ,{\sigma }^2\right)dz\end{array}$$整个结果分为三项积分，第一项实际上就是$-log{\sigma }^2$乘以概率密度的积分（也就是1），所以结果是$-log{\sigma }^2$；第二项实际是正态分布的二阶矩，熟悉正态分布的读者应该都清楚正态分布的二阶矩为${\mu }^2+{\sigma }^2$；而根据定义，第三项实际上就是“-方差除以方差=-1”。所以总结果就是$\frac{1}{2}(-1+{\sigma }^2+{\mu }^2-\log {\sigma }^2)$，推广到多元高斯分布可以得到：$$D_{KL}\left(q_{\varphi }(z\mid X),p(z)\right)=\sum _{j=1}^d\frac{1}{2}(-1+{\sigma }^{(j)}{}^2+{\mu }^{(j)}{}^2-\log {\sigma }^{(j)}{}^2).$$其中${\sigma }^{(j)2}$表示向量$a$的第$j$个元素的平方。

${\mathbb{E}}_{q_{\varphi }}\left[\log p_{\theta }(X\mid z)\right]$称为Reconstruction Loss，通常可以通过采样的方式估计其值：$${\mathbb{E}}_{q_{\varphi }}\left[\log p_{\theta }(X\mid z)\right]\approx \frac{1}{m}\sum _{i=1}^m\log p_{\theta }\left(X\mid z_i\right),$$其中$z_i\sim q_{\varphi }\left(z\mid x_i\right)=\mathcal{N}\left(z\mid \mu \left(x_i;\varphi \right),{\sigma }^2\left(x_i;\varphi \right)\ast I\right)\text{。}$

若假设数据为固定方差的高斯分布，则极大似然估计后得到的目标函数等价于极小平方估计。设每个数据点$x_i$的维度为$K$，即$X|z_i$服从一个K维高斯分布，易得：$$\begin{array}{rl}\log p_{\theta }\left(X\mid z_i\right)& =\log \frac{\exp \left(-\frac{1}{2}(X-{\mu }^{\prime }{)}^{\mathrm{T}}{\Sigma }^{\prime -1}(X-{\mu }^{\prime })\right)}{\sqrt{(2\pi {)}^k|{\Sigma }^{\prime }|}}\\ & =-\frac{1}{2}(X-{\mu }^{\prime }{)}^{\mathrm{T}}{\Sigma }^{\prime -1}(X-{\mu }^{\prime })-\log \sqrt{(2\pi {)}^k|{\Sigma }^{\prime }|}\\ & =-\frac{1}{2}\sum _{k=1}^K\frac{(X^{(k)}-{\mu }^{\prime (k)}{)}^2}{{\sigma }^{\prime (k)}}-\log \sqrt{(2\pi {)}^K\prod _{k=1}^K{\sigma }^{\prime (k)}}.\end{array}$$,至此，对于数据集$X$的每一个数据$x_i$都进行上次计算，然后计算总的损失函数：$$\begin{array}{rl}\mathcal{L}& =-\frac{1}{n}\sum _{i=1}^n\ell (p_{\theta },q_{\varphi })\\ & =\frac{1}{n}\sum _{i=1}^n{D}_{KL}\left(q_{\varphi },p\right)-\frac{1}{n}\sum _{i=1}^n{\mathbb{E}}_{q_{\varphi }}\left[\log p_{\theta }(x_i\mid z)\right]\\ & =\frac{1}{n}\sum _{i=1}^n{D}_{KL}\left(q_{\varphi },p\right)-\frac{1}{nm}\sum _{i=1}^n\sum _{j=1}^m\log p_{\theta }\left(x_i\mid z_j\right).\end{array}$$，前面提到过的采样以进行近似计算是一件非常computation costly的工作，相比之下这里的采样计算实际上是从分布$q_{\varphi }(z\mid x_i)$中采样得到的$z_j$，在网络的训练过程中，近似分布很快就会逼近分布$\log p_{\theta }(x_i\mid z)$，这样一便很大可能能够在有限次数的采样中，采样到与$x_i$关联的$z_j$。

在编写代码时只对一个$x_i$只采样一个$z_j$，即$m=1$就可以达到不错的效果。所以可以以将损失改写为：KaTeX parse error: Double superscript at position 194: …+\sigma_i^{(j)}^̲2+\mu_i^{(j)}^2…由于在前文假设了$p_{\theta }(X\mid z_i)$对任意$z_i$均是方差固定的各维度独立的$K$维高斯分布，不妨令超参数$\sigma$为元素值全为$\frac{1}{2}$的$K$维向量。则损失函数可以继续改为写：

其中，$x_i$代表第i个样本，是Encoder的输入。${\mu }_i$和${\sigma }_i^2$是Encoder的输出，代表$z\mid x_i$的分布的参数。$z_i$是从$z\mid x_i$中采样得到的一个样本，它是Decoder的输入。${\mu }_i^{{}^{\prime }}$和${\sigma }_i^{{}^{\prime }2}$是Decoder的输出，用来表示$z_i$解码后对应的数据点$f(z_i)={\hat{x}}_i$。

至此，得到了在假设先验和后验分布均是高斯分布的情况下，VAE最的损失函数。而采用高斯分布只是因为其推导时的简便性，我们同样可以根据数据的情况，假设更加复杂分布来推导、训练VAE。如果使用更加复杂的分布，VAE也许会更难训练，但是对于现实情况的表示会更好。

CVAE

训练好了VAE后，仍然存在一个问题，虽然我们能够通过采样一个$z_i$来生成一个对应的$x_i$，但无法保证$x_i$是属于哪一类的图片。以MNIST手写数字数据集为例，我们无法保证生成对应的数字，因为每次生成的结果都是随机的。因为目前的VAE是无监督训练的，因此很自然想到：如果利用有标签的数据来训练VAE，得到的模型能不能根据标签来生成对应的图片呢？答案是肯定的，这种情况叫做Conditional VAE，或者叫CVAE。

CVAE的实现思路相比VAE也很简单，只需要做些修改。假设我们现在的数据集为$X$
，以及数据集$X$的标签$Y$，那么我们只需要在训练VAE的时候，每一步都考虑标签$Y$即可。即：
1.原来MLE是对$p_{\theta }(X)$建模，那么现在需要对建模$p_{\theta }(X|Y)$；
2.原来是对$p(z)$进行建模，现在是对$p(z|y_i)$；
3.原来的Encoder是对近似后验分布$q_{\varphi }(z|x_i)$进行建模，现在是对近似后验分布$q_{\varphi }(z|x_i,y_i)$；
4.原来Decoder是对$p_{\theta }(X|z_i)$估计，现在即是对$p_{\theta }(X|z_i,y_i)$估计

在前面的讨论中，我们希望数据集$X$经过编码后，$z$的分布都具有零均值和单位方差，这个“希望”是通过加入了KL loss来实现的。如果现在多了类别信息$Y$
，我们可以希望同一个类的样本都有一个专属的均值${\mu }^Y$（方差不变，还是单位方差），这个${\mu }^Y$让模型自己训练出来。这样的话，有多少个类就有多少个正态分布，而在生成的时候，我们就可以通过控制均值来控制生成图像的类别。事实上，这样可能也是在VAE的基础上加入最少的代码来实现CVAE的方案。

测试代码：

苏剑林老师的代码地址如下：VAE

下面是一个基于手写数据集的简单实现：

import os
from tqdm import tqdm
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"

import torch
import torch.nn as nn
import torchvision
from torchvision import transforms
from torchvision.utils import save_image
from torchvision.datasets import MNIST
from matplotlib import pyplot as plt


class VAE(nn.Module):

    def __init__(self, in_features, latent_size, y_size=0):
        super(VAE, self).__init__()

        self.latent_size = latent_size
        
		# 编码器网络层
        self.encoder_forward = nn.Sequential(
            nn.Linear(in_features + y_size, in_features),
            nn.LeakyReLU(),
            nn.Linear(in_features, in_features),
            nn.LeakyReLU(),
            nn.Linear(in_features, self.latent_size * 2)
        )
        
		# 解码器网络层
        self.decoder_forward = nn.Sequential(
            nn.Linear(self.latent_size + y_size, in_features),
            nn.LeakyReLU(),
            nn.Linear(in_features, in_features),
            nn.LeakyReLU(),
            nn.Linear(in_features, in_features),
            nn.Sigmoid()
        )
	# 编码器，返回mu和sigma^2
    def encoder(self, X):
        out = self.encoder_forward(X)
        mu = out[:, :self.latent_size]
        log_var = out[:, self.latent_size:]
        return mu, log_var
        
	# 解码器，返回mu'来作为采样结果
    def decoder(self, z):
        mu_prime = self.decoder_forward(z)
        return mu_prime
        
	# 参数重整化
    def reparameterization(self, mu, log_var):
        epsilon = torch.randn_like(log_var)
        z = mu + epsilon * torch.sqrt(log_var.exp())
        return z

	# 损失函数
    def loss(self, X, mu_prime, mu, log_var):
        # reconstruction_loss = F.mse_loss(mu_prime, X, reduction='mean') is wrong!
        reconstruction_loss = torch.mean(torch.square(X - mu_prime).sum(dim=1))

        latent_loss = torch.mean(0.5 * (log_var.exp() + torch.square(mu) - log_var).sum(dim=1))
        return reconstruction_loss + latent_loss

	# 前向传播
    def forward(self, X, *args, **kwargs):
        mu, log_var = self.encoder(X)
        z = self.reparameterization(mu, log_var)
        mu_prime = self.decoder(z)
        return mu_prime, mu, log_var


class CVAE(VAE):

    def __init__(self, in_features, latent_size, y_size):
        super(CVAE, self).__init__(in_features, latent_size, y_size)

    def forward(self, X, y=None, *args, **kwargs):
        y = y.to(next(self.parameters()).device)
        X_given_Y = torch.cat((X, y.unsqueeze(1)), dim=1)

        mu, log_var = self.encoder(X_given_Y)
        z = self.reparameterization(mu, log_var)
        z_given_Y = torch.cat((z, y.unsqueeze(1)), dim=1)

        mu_prime_given_Y = self.decoder(z_given_Y)
        return mu_prime_given_Y, mu, log_var


def train(model, optimizer, data_loader, device, name='VAE'):
    model.train()

	# 总损失，一个epoch的损失
    total_loss = 0
    pbar = tqdm(data_loader) # 时间条
    for X, y in pbar:
        batch_size = X.shape[0]
        X = X.view(batch_size, -1).to(device) # 拉成向量，维度（batch_size，自动计算）
        model.zero_grad() # 梯度清零

        if name == 'VAE':
            mu_prime, mu, log_var = model(X)
        else:
            mu_prime, mu, log_var = model(X, y)

		# 损失函数
        loss = model.loss(X.view(batch_size, -1), mu_prime, mu, log_var)
        loss.backward() # 反向传播
        optimizer.step() # 参数更新

        total_loss += loss.item()
        pbar.set_description('Loss: {loss:.4f}'.format(loss=loss.item()))

    return total_loss / len(data_loader) # 返回平均损失


@torch.no_grad()
def save_res(vae, cvae, data, latent_size, device):
    num_classes = len(data.classes)

    # raw samples from dataset
    out = []
    for i in range(num_classes):
        img = data.data[torch.where(data.targets == i)[0][:num_classes]] # 对于每个num_class选取num_claass个图片
        out.append(img)
    out = torch.stack(out).transpose(0, 1).reshape(-1, 1, 28, 28) # out的维度（num_classes ** 2，1,28，28）
    save_image(out.float(), './img/raw_samples.png', nrow=num_classes, normalize=True)

    # samples generated by vanilla VAE
    z = torch.randn(num_classes ** 2, latent_size).to(device)
    out = vae.decoder(z)
    save_image(out.view(-1, 1, 28, 28), './img/vae_samples.png', nrow=num_classes)

    # sample generated by CVAE
    z = torch.randn(num_classes ** 2, latent_size).to(device)
    y = torch.arange(num_classes).repeat(num_classes).to(device)
    z_given_Y = torch.cat((z, y.unsqueeze(1)), dim=1)
    out = cvae.decoder(z_given_Y)
    save_image(out.view(-1, 1, 28, 28), './img/cvae_samples.png', nrow=num_classes)
    
def plot_data(data):
    # 带坐标轴的显示部分图片
    
    # data = MNIST('../../data/', download=True, transform=transforms.ToTensor())
    num_classes = len(data.classes)
    out = []
    for i in range(num_classes):
        img = data.data[torch.where(data.targets == i)[0][:num_classes]]
        out.append(img)
    d = torch.stack(out).transpose(0, 1)
    out = torch.stack(out).transpose(0, 1).reshape(-1, 1, 28, 28) # out的维度（num_class **2, 1, 28, 28）
    tensor = torchvision.transforms.ToPILImage()(out[0]) # 三维Tensor（channel，height，width）
    tensor = out[0].permute(1,2,0) # （channel，height，width）→（height，width，channel），这是因为plt.imshow的格式为（h，w，c）
    plt.imshow(tensor)
    plt.show()

    for i in range(num_classes **2): # i从0到99
        plt.subplot(10,10,i+1)
        tensor = torchvision.transforms.ToPILImage()(out[i]) # 三维Tensor（channel，height，width）
        tensor = out[i].permute(1,2,0) # （channel，height，width）→（height，width，channel），这是因为plt.imshow的格式为（h，w，c）
        plt.imshow(tensor)
    plt.show()
plot_data(data)

def plot_data1(data):
    # 不带坐标轴的显示部分图片
    
    # data = MNIST('../../data/', download=True, transform=transforms.ToTensor())
    num_classes = len(data.classes)
    out = []
    for i in range(num_classes):
        img = data.data[torch.where(data.targets == i)[0][:num_classes]]
        out.append(img)
    d = torch.stack(out).transpose(0, 1)
    out = torch.stack(out).transpose(0, 1).reshape(-1, 1, 28, 28) # out的维度（num_class **2, 1, 28, 28）
    tensor = torchvision.transforms.ToPILImage()(out[0]) # 三维Tensor（channel，height，width）
    tensor = out[0].permute(1,2,0) # （channel，height，width）→（height，width，channel），这是因为plt.imshow的格式为（h，w，c）
    current_axes=plt.axes()
    current_axes.xaxis.set_visible(False)
    current_axes.yaxis.set_visible(False)
    plt.imshow(tensor)
    plt.show()
    
    
    fig, axs = plt.subplots(num_classes, num_classes) # fig.subplots()”是生成一个画布一个子图时用的，如果想要在同一个画布中生成多个子图，需要用“subplot()”函数
#     print(axs.shape) # 查看axs维度
    for i in range(num_classes **2): # i从0到99
        j = i // 10 # 商，代表行
        k = i % 10 # 余数，代表列
        tensor = torchvision.transforms.ToPILImage()(out[i]) # 三维Tensor（channel，height，width）
        tensor = out[i].permute(1,2,0) # （channel，height，width）→（height，width，channel），这是因为plt.imshow的格式为（h，w，c）
        off = "off" 
        # exec() 函数可以理解为执行一段写在字符串中的代码语句。
        exec(f'axs[{j}][{k}].axis(off)') # 去除坐标轴
        exec(f'axs[{j}][{k}].xaxis.set_ticks([])')  # 去除X轴
        exec(f'axs[{j}][{k}].yaxis.set_ticks([])')  # 去除Y轴
        exec(f'axs[{j}][{k}].imshow(tensor, interpolation=None)')
        plt.imshow(tensor)
        
    plt.show()
    
plot_data1(data)

def main():
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    device = torch.device(device) # 有GPU就是用GPU

    batch_size = 256 * 4
    epochs = 50
    latent_size = 64
    in_features = 28 * 28
    lr = 0.001

	#构建数据集
    data = MNIST('../../dataset/', download=True, transform=transforms.ToTensor())
    data_loader = torch.utils.data.DataLoader(data, batch_size=batch_size, shuffle=True)

    # train VAE
    # 模型
    vae = VAE(in_features, latent_size).to(device)
    # 优化器
    optimizer = torch.optim.AdamW(vae.parameters(), lr=lr)

    print('Start Training VAE...')
    # 开始迭代，一共迭代epoch次
    for epoch in range(1, 1 + epochs):
        loss = train(vae, optimizer, data_loader, device, name='VAE')
        print("Epochs: {epoch}, AvgLoss: {loss:.4f}".format(epoch=epoch, loss=loss))
    print('Training for VAE has been done.')

    # train VCAE
    cvae = CVAE(in_features, latent_size, y_size=1).to(device)
    optimizer = torch.optim.AdamW(cvae.parameters(), lr=lr)

    print('Start Training CVAE...')
    for epoch in range(1, 1 + epochs):
        loss = train(cvae, optimizer, data_loader, device, name='CVAE')
        print("Epochs: {epoch}, AvgLoss: {loss:.4f}".format(epoch=epoch, loss=loss))
    print('Training for CVAE has been done.')

    save_res(vae, cvae, data, latent_size, device)


if __name__ == '__main__':
    main()

参考博客

机器学习方法—优雅的模型（一）：变分自编码器（VAE）
变分自编码器（一）：原来是这么一回事