PyTorch计算机视觉（9）------变分自编码器（VAE）详解与实现
0. 前言
本节介绍用于无监督学习的变分自编码器 (Variational Autoencoder, VAE)，以生成模拟的二次曲线和 MNIST 图像。VAE 模型易于训练，但其生成的图像质量不如生成对抗网络 (Generative Adversarial Network, GAN)。我们已经学习了如何使用 GAN 生成伪造图像，然而，由于模型坍塌等问题，DCGAN 和 cDCGAN 的模型训练仍然具有挑战性。调整初始学习率、学习率调度器和其他超参数以实现成功的模型训练需要很长时间，我们学习了使用 WGAN-GP 稳定模型训练。变分自编码器 (Variational Autoencoder, VAE) 易于训练，通过学习数据集底层结构的紧凑表示，VAE 可用于生成与给定数据集中图像相似的新数据样本(例如图像)。
1. 变分自编码器简介
变分自编码器 (Variational Autoencoder, VAE) 系统中也有两个模型。编码器模型将输入数据映射到低维表示(将每个图像压缩成低维潜在向量)。潜在向量可以有一维、二维、三维甚至超过一百维。例如，MNIST 图像可以压缩为仅具有两个元素的二维潜在空间向量。然后，系统中的解码器网络使用该潜在向量生成与原始 MNIST 图像相似的图像。图像压缩率为 392 (= 784/2)。然而，VAE 并不是为了图像压缩而设计的，因为 VAE 算法的表现不如 JPEG 算法。VAE 比 GAN 提出的时间更早，其理论更难以理解，并且 VAE 生成的合成图像通常不如 GAN 生成的图像清晰。那么，我们为什么要花时间学习 VAE 呢？
我们感兴趣的并非伪造图像，而是 VAE 系统生成的那些低维潜在向量。潜在向量中的每个元素都是 VAE 系统学习到的图像特征。通过调整潜在向量中这些元素的值，我们可以编辑图像以生成新图像。例如，我们可以将图像中的人脸表情从愤怒变为微笑，或将图像中的月相从满月变为新月。借助基于 VAE 算法的新技术，如矢量量化变分自编码器 (Vector Quantized Variational Autoencoder, VQ-VAE)，我们可以在无监督学习中取得优异结果。我们可以将 VAE 与 GAN 结合，以利用两种技术的优势进行模型训练。VAE 已应用于许多领域，包括图像生成、异常检测和强化学习。
2. VAE 与 beta-VAE
在 VAE 系统用于图像生成的应用中，数据集中的每个图像都是高维观测空间中的一个点。例如，每个 MNIST 图像是 784 维空间中的一个点。通常，我们的计算机内存不足以处理数据集中的所有图像。我们必须将图像数据集切片成包含多个批次的数据加载器。它们的批大小是一个常数 n n n。批次中的图像是观测数据： x 1 , x 2 , ... , x n x_1, x_2, ..., x_n x1,x2,...,xn。每个图像由具有以下元素的潜在向量生成： z 1 , z 2 , ... , z m z_1, z_2, ..., z_m z1,z2,...,zm，其中 m ≪ n m ≪ n m≪n。下表显示了观测数据点和未知潜在向量的未知联合概率。 p ( x j ) p(x_j) p(xj) 的边际概率也是未知的 ( j = 1 , 2 , ... , n ) (j = 1, 2, ..., n) (j=1,2,...,n)。我们假设先验分布 p ( z i ) p(z_i) p(zi) 遵循标准正态分布 N ( 0 , I ) N(0,I) N(0,I)，其中 I I I 是单位矩阵。根据以下公式：
L = − l n ( P ( X ) ) = l n ∏ i = 1 n p ( x j ) = − ∑ j = 1 n l n ( ∑ i = 1 m = 3 p ( x j ∣ z i ) p ( z i ) ) L=-ln(P(X))=ln∏{i=1}^np(x_j)=-\sum{j=1}^nln(\sum_{i=1}^{m=3}p(x_j|z_i)p(z_i)) L=−ln(P(X))=lni=1∏np(xj)=−j=1∑nln(i=1∑m=3p(xj∣zi)p(zi))
我们有 p ( x j ) = ∑ j = 1 m p ( x j , z i ) = ∑ z p ( x j , z ) p(x_j) = ∑_{j=1}^mp(x_j, z_i) = ∑_z p(x_j, z) p(xj)=∑j=1mp(xj,zi)=∑zp(xj,z)。根据数据集的负对数似然，我们得到以下损失函数：
L = − l n ( P ( X ) ) = l n ∏ j = 1 n p ( x j ) = − ∑ j = 1 n l n ( ∑ z p ( x j , z ) ) L=-ln(P(X))=ln∏{j=1}^np(x_j)=-\sum{j=1}^nln(\sum_zp(x_j,z)) L=−ln(P(X))=lnj=1∏np(xj)=−j=1∑nln(z∑p(xj,z))
由于联合概率未知，我们无法直接通过上述方程对损失函数进行处理。但我们可以构建一个编码器网络，该网络以数据加载器中的每批图像作为输入，为每个图像获取均值向量和方差向量。这些均值向量和方差向量随后用于为每个图像构建多元正态分布 q ( z ∣ x ) q q(z|x)q q(z∣x)q。图像潜在向量的每个元素通过从 q ( z ∣ x ) q(z|x) q(z∣x) 中采样获得。利用 VAE 编码器的先验概率及关于 g ( x ) = − l n ( x ) g(x)=−ln(x) g(x)=−ln(x) 的 Jensen 不等式，我们可以间接找到 VAE 系统损失函数的最小值。在下面的负对数函数内部是 p / q p/q p/q 关于 z z z 的平均值，它应小于或等于 − l n ( p / q ) −ln(p/q) −ln(p/q) 的平均值：
L = − ∑ j = 1 n l n ( ∑ z q ( z ∣ x j ) p ( x j , z ) q ( z ∣ x j ) ) ≤ − ∑ j = 1 n ∑ z q ( z ∣ x j ) l n p ( x j , z ) q ( z ∣ x j ) L ≤ ∑ j = 1 n ∑ z q ( z ∣ x j ) l n q ( z ∣ x j ) p ( z ) − ∑ z q ( z ∣ x j ) l n ( p ( x j ∣ z ) ) L ≤ ∑ j = 1 n ( D K L ( q ( z ∣ x j ) ∣ ∣ p ( z ) ) − E z ∼ q ( z ∣ x j ) l n ( p ( x j ∣ z ) ) ) L=-\sum_{j=1}^nln(\sum_zq(z|x_j)\frac {p(x_j,z)}{q(z|x_j)})\leq -\sum_{j=1}^n\sum_zq(z|x_j)ln\frac {p(x_j,z)}{q(z|x_j)}\\ L\leq \sum_{j=1}^n\\sum_zq(z\|x_j)ln\\frac {q(z\|x_j)}{p(z)}-\\sum_zq(z\|x_j)ln(p(x_j\|z))\\ L\leq \sum_{j=1}^n(D_{KL}(q(z|x_j)||p(z))-E_{z\sim q(z|x_j)}ln(p(x_j|z))) L=−j=1∑nln(z∑q(z∣xj)q(z∣xj)p(xj,z))≤−j=1∑nz∑q(z∣xj)lnq(z∣xj)p(xj,z)L≤j=1∑nz∑q(z∣xj)lnp(z)q(z∣xj)−z∑q(z∣xj)ln(p(xj∣z))L≤j=1∑n(DKL(q(z∣xj)∣∣p(z))−Ez∼q(z∣xj)ln(p(xj∣z)))
我们还可以为 VAE 系统构建解码器网络 P。解码器使用来自编码器模型 Q 的潜在向量生成伪造图像 x ^ \hat x x^。我们将解码器训练为具有单位方差的标准正态分布 p ( x ∣ z ) p(x|z) p(x∣z)：
p ( x j ∣ z ) = 1 ( 2 π ) m ∣ I ∣ e − 1 2 ( x j − μ θ ( z ) ) T I ( x j − μ θ ( z ) ) p(x_j|z ) = \frac{1}{(\sqrt {2π)^{m} |I|}} e^{-\frac{1}{2} (x_j - μ_θ(z))^T I(x_j - μ_θ(z))} p(xj∣z)=(2π)m∣I∣ 1e−21(xj−μθ(z))TI(xj−μθ(z))
其中 θ θ θ 是解码器的一组参数， x ^ j = μ θ ( z ) \hat x_j = μ_θ(z) x^j=μθ(z) 是解码器的输出。根据以上信息，VAE 系统的损失函数可以写成以下方程(忽略常数项)：
L ≤ ∑ j = 1 n ∑ i = 1 m ( u i , j 2 + σ i , j 2 − l n ( σ i , j 2 − 1 ) + 1 2 ∑ j = 1 n ∣ ∣ x j − x ^ j ∣ ∣ 2 L\leq \sum_{j=1}^n\sum_{i=1}^m(u_{i,j}^2+\sigma_{i,j}^2-ln(\sigma_{i,j}^2-1)+\frac 12\sum_{j=1}^n||x_j-\hat x_j||^2 L≤j=1∑ni=1∑m(ui,j2+σi,j2−ln(σi,j2−1)+21j=1∑n∣∣xj−x^j∣∣2
如果我们能够最小化以上方程右侧的值，那么逻辑上我们也能最小化方程左侧的值，因为左侧的值小于或等于右侧的值。方程右侧的第一部分称为"正则化损失"，第二部分称为"重构损失"。在上述表格和方程中，我们将潜在变量 z z z 视为离散变量。事实上，在 VAE 系统中，潜在变量 z z z 可以是连续变量。我们也可以通过以下数学推导得到相同的 VAE 损失函数结果，其中 ∫ z q ( z ∣ x j ) d z = 1 \int_z q(z \mid x_j) dz = 1 ∫zq(z∣xj)dz=1 且 D K L ≥ 0 D_{KL} \geq 0 DKL≥0。
L = − l n ( P ( X ) ) = l n ∏ j = 1 n p ( x j ) = − ∑ j = 1 n l n ( p ( x j ) ) ∫ z q ( z ∣ x j ) d z L = − ∑ j = 1 n ∫ z q ( z ∣ x j ) l n ( p ( x j ) ) d z = − ∑ j = 1 n ∫ z q ( z ∣ x j ) l n ( p ( z , x j ) p ( z ∣ x j ) ) d z L = − ∑ j = 1 n ∫ z q ( z ∣ x j ) l n ( p ( z , x j ) q ( z ∣ x j ) q ( z ∣ x j ) p ( z ∣ x j ) ) d z L = − ∑ j = 1 n ∫ z q ( z ∣ x j ) l n ( p ( z , x j ) q ( z ∣ x j ) ) − ∑ j = 1 n ∫ z q ( z ∣ x j ) l n ( q ( z ∣ x j ) p ( z ∣ x j ) ) d z L = ∑ j = 1 n ∫ z Q = q ( z ∣ x j ) l n ( p ( z , x j q ( z ∣ x j ) ) d z L ≤ ∑ j = 1 n ∫ z q ( z ∣ x j ) l n ( p ( x j ∣ z ) p ( z ) q ( z ∣ x j ) ) L ≤ ∑ j = 1 n ∫ z q ( z ∣ x j ) l n ( q ( z ∣ x j ) p ( z ) ) d z − ∑ j = 1 n ∫ z q ( z ∣ x j ) l n ( p ( x j ∣ z ) ) d z L ≤ ∑ j = 1 n ( D K L ( q ( z ∣ x j ) ∣ ∣ p ( z ) ) − E z ∼ q ( z ∣ x j ) l n ( p ( x j ∣ z ) ) ) L=-ln(P(X))=ln∏{j=1}^np(x_j)=-\sum{j=1}^nln(p(x_j))\\int_zq(z\|x_j)dz\\ L=-\sum_{j=1}^n\int_zq(z|x_j)ln(p(x_j))dz=-\sum_{j=1}^n\int_zq(z|x_j)ln(\frac {p(z,x_j)}{p(z|x_j)})dz\\ L=-\sum_{j=1}^n\int_zq(z|x_j)ln(\frac {p(z,x_j)q(z|x_j)}{q(z|x_j)p(z|x_j)})dz\\ L=-\sum_{j=1}^n\int_zq(z|x_j)ln(\frac {p(z,x_j)}{q(z|x_j)})-\sum_{j=1}^n\int_zq(z|x_j)ln(\frac {q(z|x_j)}{p(z|x_j)})dz\\ L=\sum_{j=1}^n\int_zQ=q(z|x_j)ln(\frac {p(z,x_j}{q(z|x_j)})dz\\ L\leq\sum_{j=1}^n\int_zq(z|x_j)ln(\frac{p(x_j|z)p(z)}{q(z|x_j)})\\ L\leq\sum_{j=1}^n\int_zq(z|x_j)ln(\frac{q(z|x_j)}{p(z)})dz-\sum_{j=1}^n\int_zq(z|x_j)ln(p(x_j|z))dz\\ L\leq\sum_{j=1}^n(D_{KL}(q(z|x_j)||p(z))-E_{z\sim q(z|x_j)}ln(p(x_j|z))) L=−ln(P(X))=lnj=1∏np(xj)=−j=1∑nln(p(xj))∫zq(z∣xj)dzL=−j=1∑n∫zq(z∣xj)ln(p(xj))dz=−j=1∑n∫zq(z∣xj)ln(p(z∣xj)p(z,xj))dzL=−j=1∑n∫zq(z∣xj)ln(q(z∣xj)p(z∣xj)p(z,xj)q(z∣xj))dzL=−j=1∑n∫zq(z∣xj)ln(q(z∣xj)p(z,xj))−j=1∑n∫zq(z∣xj)ln(p(z∣xj)q(z∣xj))dzL=j=1∑n∫zQ=q(z∣xj)ln(q(z∣xj)p(z,xj)dzL≤j=1∑n∫zq(z∣xj)ln(q(z∣xj)p(xj∣z)p(z))L≤j=1∑n∫zq(z∣xj)ln(p(z)q(z∣xj))dz−j=1∑n∫zq(z∣xj)ln(p(xj∣z))dzL≤j=1∑n(DKL(q(z∣xj)∣∣p(z))−Ez∼q(z∣xj)ln(p(xj∣z)))
损失函数涉及到两个部分之间的平衡问题，两部分损失值的量级并不相同。2017 年 Higgins 在论文中提出解决方案，引入超参数β来调节损失函数两部分的贡献权重，并将采用修正损失函数的 VAE 称为 beta-VAE。当 β = 1.0 β=1.0 β=1.0 时，即为标准 VAE 系统；若 β > 1.0 β>1.0 β>1.0，正则化损失 D K L D_{KL} DKL 会快速趋近于零，导致出现生成图像的"模型坍塌"问题；当 β = 0.0 β=0.0 β=0.0 时，VAE 系统退化为自编码器系统。增大 β β β 值虽然会使生成图像变得模糊，但能增强对生成图像在不同阶段进行调整的能力。
L = β ∑ j = 1 n ∑ i = 1 m ( u i , j 2 + σ i , j 2 − l n ( σ i , j 2 − 1 ) + 1 2 ∑ j = 1 n ∣ ∣ x j − x ^ j ∣ ∣ 2 L=\beta \sum_{j=1}^n\sum_{i=1}^m(u_{i,j}^2+\sigma_{i,j}^2-ln(\sigma_{i,j}^2-1)+\frac 12\sum_{j=1}^n||x_j-\hat x_j||^2 L=βj=1∑ni=1∑m(ui,j2+σi,j2−ln(σi,j2−1)+21j=1∑n∣∣xj−x^j∣∣2
3. beta-VAE 生成二次曲线
在二次曲线生成中，beta-VAE 系统的编码器模型 Q 和解码器模型 P 仅使用两个全连接层。为便于数据可视化，将潜在维度超参数 z_dim 设为 2，这意味着每条包含 21 个数据点的曲线均由二维潜在向量生成。系统解码器将从二维高斯分布中采样两个数值，重构出包含 21 个点的曲线。超参数 beta 设置为 0.5，以平衡损失函数两部分的贡献权重。在二次曲线中加入了噪声，如下图所示，由于噪声干扰，64 条真实曲线并不平滑。本项目无需预设数据集，而是在数据加载器需要每批次数据时实时生成训练数据，每个 epoch 训练的批次数设为 5。
带噪声的批次中每条真实曲线经编码器模型 Q 处理后，会得到多元高斯分布的两个向量：均值向量 μ μ μ 和标准差向量 σ σ σ。随后通过公式 z = μ + σ ε z=μ+σε z=μ+σε 为每条真实曲线生成潜在向量，其中 ε ε ε 为标准多元正态分布 N ( 0 , I ) N(0,I) N(0,I) 的随机数。后续项目代码中将展示一个技巧：如何计算标准差向量 σ σ σ------其每个元素必须为正数。解码器模型 P 利用这批潜在向量 Z Z Z 生成一批伪造图像 X ^ \hat X X^。下图显示，在训练初始阶段生成的合成曲线噪声明显，而经过 20000 个 epoch 训练后噪声被有效滤除，如下图所示，生成的曲线变得平滑。
python
import torch
import torch.nn as nn
from tqdm import trange
import torch.optim.lr_scheduler as lr_scheduler
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
n_epochs = 20000
batch_size = 64
z_dim = 2
n_hidden = 256
n_points = 21
lr = 1e-4
beta = 0.5
n_batch = 5
n_samples = n_batch*batch_size
x = np.linspace(-1, 1, n_points)
def training_data():
a = np.random.uniform(1, 2, size=batch_size).reshape(-1,1)
noises = 0.015*np.random.randn(n_points)
y_values = a*(x+noises)**2 + (a-1)
return y_values
y = training_data()
def show_curves(x, y):
fig, ax = plt.subplots(figsize=(4,4))
for i in range(batch_size):
ax.plot(x, y[i,:])
ax.set(xlabel='x', ylabel='y')
ax.set(xticks=np.arange(-1.0, 1.1, 0.5))
ax.grid(color='g', linestyle=':')
plt.show()
show_curves(x,y)编码器类使用一批真实曲线作为输入，这些曲线通过一个全连接层处理。编码器输出三个张量：潜在向量 z z z、均值张量 μ μ μ 和对数方差张量 l n ( σ 2 ) ln(σ^2) ln(σ2)。潜在向量 z z z 通过公式 z = μ + σ ε z=μ+σε z=μ+σε 从正态分布 N ( μ , σ 2 ) N(μ,σ^2) N(μ,σ2) 中采样获得，其中 ε ε ε 是由标准正态分布 N ( 0 , I ) N(0,I) N(0,I) 生成的随机数。这是计算标准差(必须为正数)的技巧：通过对数方差可通过公式 σ = e x p ( 0.5 ⋅ l n ( σ 2 ) ) σ=exp(0.5·ln(σ^2)) σ=exp(0.5⋅ln(σ2)) 求得标准差。解码器 P 包含两个全连接层，以一批潜在向量作为输入，输出一批生成曲线。
python
class encoder(nn.Module):
def __init__(self):
super().__init__()
self.net = nn.Sequential(
nn.Linear(n_points, n_hidden),
nn.LeakyReLU(0.2, inplace=True))
self.mu = nn.Linear(n_hidden, z_dim)
self.logvar = nn.Linear(n_hidden, z_dim)
def forward(self, curves):
outputs = self.net(curves)
mu = self.mu(outputs)
logvar = self.logvar(outputs)
std = torch.exp(0.5*logvar)
epsilon = torch.randn_like(std)
z = mu + std*epsilon
return z, mu, logvar
Q = encoder().cuda()
class decoder(nn.Module):
def __init__(self):
super().__init__()
self.net = nn.Sequential(
nn.Linear(z_dim, n_hidden),
nn.ReLU(True),
nn.Linear(n_hidden, n_points))
def forward(self, z):
fake_curves = self.net(z)
return fake_curves
P = decoder().cuda()根据公式定义损失函数。若损失函数的两部分采用"均值"而非"总和"计算，则需要将 β β β 值调整至 0.01 左右。模型训练采用 Adam 优化器。训练函数使用 pandas 数据框记录训练过程中的损失值。每个训练 epoch 生成五批真实曲线用于训练：每批训练数据经编码器模型 Q 处理后，其输出 z 作为模型 P 的输入以生成模拟曲线，随后计算损失值并通过 Adam 优化器更新两个模型的参数。损失函数曲线如下图所示。
python
def loss_function(fake, real, mu, logvar):
MSE = 0.5*torch.nn.functional.mse_loss(fake, real, reduction='sum')
Dkl = torch.sum(mu.pow(2) + logvar.exp() - logvar - 1)
return MSE, Dkl
optimizer = torch.optim.Adam([{'params': P.parameters()},
{'params': Q.parameters()}], lr=lr)
scheduler = lr_scheduler.OneCycleLR(optimizer, max_lr=lr,
steps_per_epoch=n_batch, epochs=n_epochs, pct_start=0.6)
def fit(epochs):
torch.cuda.empty_cache()
# The DataFrame df is a recorder of the training history
df = pd.DataFrame(np.empty([epochs, 3]),
index = np.arange(epochs), columns=['loss','Dkl', 'MSE'])
for i in trange(epochs): #trange from tqdm for a progress bar
loss_total = 0.0; loss_prior=0; loss_rec=0
for _ in range(n_batch): #training with 5 batches of training data
real_curves = torch.cuda.FloatTensor(training_data())
z, mu, logvar = Q(real_curves)
fake_curves = P(z)
MSE, Dkl = loss_function(fake_curves, real_curves, mu, logvar)
loss = MSE + beta*Dkl
loss_total += loss.item()
loss_prior += Dkl.item()
loss_rec += MSE.item()
optimizer.zero_grad()
loss.backward()
optimizer.step()
scheduler.step()
df.iloc[i, 0] = loss_total/n_samples
df.iloc[i, 1] = loss_prior/n_samples
df.iloc[i, 2] = loss_rec/n_samples
if i==0 or (i+1)%4000==0:
print(
"Epoch={:2}, Loss={:.2f}, Dkl_Loss={:.2f}, Rec_Loss={:.2f}"
.format(i+1, df.iloc[i,0], df.iloc[i,1], df.iloc[i,2]))
y2cpu = fake_curves.detach().cpu()
show_curves(x, y2cpu)
return df
train_history = fit(n_epochs)利用五批训练数据计算训练后编码器 Q 的输出，包括均值 (mu)、标准差 (std) 和潜在向量 (z)。这些潜在向量 (z) 的分布如下图所示，其分布模式符合二维标准正态分布函数。若设置 β = 4.0 β=4.0 β=4.0，mu 数组的每个元素将趋近于 0.0，std 数组的每个元素则趋近于 1.0；若 β = 0.0 β=0.0 β=0.0，下图中的点会排列成直线，且 std 向量的每个元素趋近于零。超参数 β β β 可用于将潜在向量压缩至符合标准正态分布的区域------这正是自编码器网络的功能体现。
python
history = train_history
fig, ax = plt.subplots(figsize=(4,4))
history.plot(ax=ax, color=['r', 'b', 'g'], style=[':', ':', ':'])
ax.grid(which='major', axis='both', color='g', linestyle=':')
ax.set( xlabel='epoch', ylim=[0, 4], ylabel='Loss');
torch.cuda.empty_cache()
Q.eval()
m = 64 #epochs
X= np.empty(m*n_batch) #X for z[0]
Y= np.empty(m*n_batch) #Y for z[1]
Z = np.empty((m*n_batch,2)) #Z for z
for i in range(n_batch):
curves_real = torch.cuda.FloatTensor(training_data())
z, mu, logvar = Q(curves_real)
Z = z.detach().cpu().numpy()
X[i*m: (i+1)*m]= Z[:,0]
Y[i*m: (i+1)*m]= Z[:,1]
fig, ax = plt.subplots(figsize=(4,4))
ax.plot(X, Y, 'g+')
ax.grid(which='major', axis='both', color='g', linestyle=':')
ax.set(xlim=[-3,3], ylim=[-3,3], xlabel ='z[0]', ylabel='z[1]',
title ='Latent Space (z_dim=2)')
plt.show()
fig, ax = plt.subplots(figsize=(4,4))
noise1 = torch.FloatTensor([-2, 0]).cuda()
latent_curve1 = P(noise1).detach().cpu().numpy()
ax.plot(x, latent_curve1, 'k-', label='first')
noise2 = torch.FloatTensor([2, 0]).cuda()
latent_curve2 = P(noise2).detach().cpu().numpy()
ax.plot(x, latent_curve2, 'b-', label='second')
noise3 = torch.FloatTensor([0, -9]).cuda()
latent_curve3 = P(noise3).detach().cpu().numpy()
ax.plot(x, latent_curve3, 'r:+', label='third')
ax.legend()
ax.set(xlabel='x', ylabel='y', xticks=np.arange(-1.0, 1.1, 0.5))
ax.grid(color='g', linestyle=':')在下图中，标记为 ⊗ 和 ⊕ 的两个潜在向量坐标分别为 [-2,0] 和 [2,0]。通过解码器 P 生成的两条平滑模拟曲线如下图所示。潜在向量中 z[0] 的值可调节生成曲线的垂直位置，z[1] 的值则能控制曲线的平滑度。下图中锯齿状的模拟曲线由潜在向量 z=[-2,15] 生成。每次运行项目代码时，控制生成曲线的参数会有所不同：有时 z[1] 的值可能主导生成曲线在 y 轴上的位置变化。
4. beta-VAE 生成 MNIST 图像
用于生成 MNIST 伪造图像的 beta-VAE 架构与生成二次曲线的架构基本相似，主要区别在于编码器 Q 和解码器 P 采用了深度卷积网络。我们知道仅含两个全连接层的神经网络处理 MNIST 图像能力有限。得到的真实与伪造 MNIST 图像(下图所示)本应相同，但可见模拟图像存在模糊及部分数字识别错误的问题。这些问题的根源在于为数据可视化设置 z_dim=2：当 MNIST 图像被编码器从 784 个像素压缩至仅含两个元素的潜在向量时，部分信息丢失导致解码器无法准确还原原始图像信息。
python
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
import torchvision.transforms as T
from torchvision.datasets import MNIST
from torchvision.utils import make_grid
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from tqdm import trange
n_epochs = 30
batch_size = 100
img_channels = 1
img_size = 28
img_dim = img_size*img_size
n_class = 10
z_dim = 2
lr = 5e-4
beta = 0.01
trainData = MNIST('./data/', train=True, download=True,
transform=T.Compose([T.ToTensor(), T.Normalize([0.5],[0.5])]))
n_samples = len(trainData)
train_dataloader = DataLoader(trainData, batch_size=batch_size, shuffle=True)
n_batch = len(train_dataloader)
for imgs, labels in train_dataloader:
print('imgs_batch.shape=', imgs.shape)
break
def denorm(img_tensors):
return img_tensors * 0.5 + 0.5
def show_imgs(imgs):
fig, ax = plt.subplots(figsize=(12,8))
input = make_grid(denorm(imgs[:64]), nrow=16, padding=2)
ax.imshow(input.permute(1,2,0), cmap='gray')
ax.set(xticks=[], yticks=[])
plt.show()
show_imgs(imgs)
def weights_init(m):
if(type(m) == nn.ConvTranspose2d or type(m) == nn.Conv2d):
nn.init.normal_(m.weight.data, 0.0, 0.02)
elif(type(m) == nn.BatchNorm2d):
nn.init.normal_(m.weight.data, 1.0, 0.02)
nn.init.constant_(m.bias.data, 0)
# Encoder Class
def basic_E(in_channels, out_channels, f=4, p=1):
return nn.Sequential(
nn.Conv2d(in_channels, out_channels,
kernel_size=f, stride=2, padding=p, bias=False),
nn.BatchNorm2d(out_channels),
nn.LeakyReLU(0.2, inplace=True))
class Encoder(nn.Module):
def __init__(self):
super().__init__()
self.net = nn.Sequential(
basic_E(1, 64, f=4, p=1),
basic_E(64, 128, f=4, p=2),
basic_E(128, 256, f=4, p=1),)
self.mu = nn.Conv2d(256, z_dim,
kernel_size=4, stride=1, padding=0, bias=True)
self.logvar = nn.Conv2d(256, z_dim,
kernel_size=4, stride=1, padding=0, bias=True)
def forward(self, images):
outputs = self.net(images)
mu = self.mu(outputs).squeeze()
logvar = self.logvar(outputs).squeeze() # logvar = ln(std^2)
std = torch.exp(0.5*logvar) # std=(e^(ln(std^2)))**0.5=(std^2)**0.5
epsilon = torch.randn_like(std) # epsilon: ~N(0,I)
z = mu + std*epsilon # noises constructed from ~N(mu, std)
return z, mu, logvar # z.shape = batch_size x z_dim
Q = Encoder().cuda()
Q.apply(weights_init)
def basic_D(in_channels, out_channels, f=4, p=1):
return nn.Sequential(
nn.ConvTranspose2d(in_channels, out_channels,
kernel_size=f, stride=2, padding=p, bias=False),
nn.BatchNorm2d(out_channels),
nn.ReLU(inplace=True))
class Decoder(nn.Module):
def __init__(self):
super().__init__()
self.net = nn.Sequential(
basic_D(z_dim, 256, f=4, p=0),
basic_D(256, 128, f=3, p=1),
basic_D(128, 64, f=4, p=1),
nn.ConvTranspose2d(64, img_channels,
kernel_size=4, stride=2, padding=1),
nn.Tanh())
def forward(self, z):
output = self.net(z.view(-1, z_dim, 1, 1))
return output
P = Decoder().cuda()
P.apply(weights_init)
def loss_function(fake, real, mu, logvar):
MSE = 0.5*torch.nn.functional.mse_loss(fake, real, reduction='sum')
Dkl = torch.sum(mu.pow(2) + logvar.exp() - logvar - 1)
return MSE, Dkl
optimizer = torch.optim.Adam([{'params': P.parameters()},
{'params': Q.parameters()}], lr=lr, betas=(0.5, 0.999))
def fit(epochs):
torch.cuda.empty_cache()
df = pd.DataFrame(np.empty([epochs, 3]),
index = np.arange(epochs),
columns=['loss','Dkl', 'MSE'])
for i in trange(epochs):
loss_total = 0.0; loss_prior=0; loss_rec=0
for real_images, _ in train_dataloader:
inputs = real_images.cuda()
z, mu, logvar = Q(inputs)
real = inputs.reshape(-1, img_dim)
fake = P(z).reshape(-1, img_dim)
MSE, Dkl = loss_function(fake, real, mu, logvar)
loss = MSE + beta*Dkl
loss_total += loss.item()
loss_prior += Dkl.item()
loss_rec += MSE.item()
optimizer.zero_grad()
loss.backward()
optimizer.step()
df.iloc[i, 0] = loss_total/n_samples
df.iloc[i, 1] = loss_prior/n_samples
df.iloc[i, 2] = loss_rec/n_samples
if i==0 or (i+1)%10==0:
print(
"Epoch={:4}, Loss={:.2f}, Dkl_Loss={:.2f}, Rec_Loss={:.2f}"
.format(i+1, df.iloc[i,0], df.iloc[i,1], df.iloc[i,2]))
show_imgs(real_images)
show_imgs(P(z).detach().cpu())
return df
train_history = fit(n_epochs)
下图展示了其他训练结果：包含潜在向量点的二维潜在空间、损失曲线以及通过 p ( x ∣ z ) p(x|z) p(x∣z) 在 z[0]∈[-10,10] 和 z[1]∈[-10,10] 范围内生成的伪造图像。将 β β β 值从 0 调整至 10 可观察到：潜在空间中的点会从大范围不规则区域被压缩至半径约3的圆形区域。虽然大幅增加 β β β 值确实能训练模型将 q ( z ∣ x ) q(z|x) q(z∣x) 强制逼近 N ( 0 , I ) N(0,I) N(0,I)，但由此生成的数字图像会模糊到难以辨识。我们需要 KL 散度在训练中发挥作用，但不宜过度。若将 z_dim 从 2 改为 20，将获得清晰度极高的 MNIST 模拟图像，但下图所示的潜在空间将无法可视化呈现。通过 VAE 技术，我们能够在无需图像数据集标签信息的情况下生成符合需求的模拟图像。
python
df = train_history
fig, ax = plt.subplots(figsize=(5,4))
df.plot(ax=ax, style=['-','-', '+'], color=['r','b','g'])
ax.grid(which='major', axis='both', color='g', linestyle=':')
ax.set(xlabel='epoch')
torch.cuda.empty_cache()
Q.eval()
x= np.empty(n_samples); y= np.empty(n_samples); c= np.empty(n_samples)
for i, (img, label) in enumerate(train_dataloader):
m = len(label)
real_img = img.cuda()
latent, _, _ = Q(real_img)
logits = latent.squeeze().detach().cpu()
x[i*m: (i+1)*m]= logits[:,0].numpy()
y[i*m: (i+1)*m]= logits[:,1].numpy()
c[i*m: (i+1)*m]= label
fig, ax = plt.subplots(figsize=(5,4))
plt.scatter(x=x, y=y, c=c, s=0.5, cmap="brg")
ax.grid(which='major', axis='both', color='k', linestyle=':')
ax.invert_yaxis()
ax.set(xlabel ='z [0]', ylabel ='z [1]',
title ='Latent Space (z_dim=2)')
cbar = plt.colorbar(label="label", orientation="vertical",shrink=0.9)
cbar.set_ticks(np.arange(10))
cbar.set_ticklabels(np.arange(10))
fig, ax = plt.subplots(figsize=(10,8))
big_image = np.empty((21*28,21*28))
b=10
for i in np.arange(0, 21, 1):
x_p = -b+i*0.1*b
for j in np.arange(0, 21, 1):
y_p = -b+j*b*0.1
noise = torch.tensor([[y_p, x_p]], dtype=torch.float).cuda()
latent_img = P(noise).detach().cpu()
big_image[28*i:28*(i+1), 28*j:28*(j+1)] = latent_img
plt.imshow(big_image, cmap='gray')
ax.set(xticks=[], yticks=[])
小结
本节介绍了变分自编码器 (Variational Autoencoder, VAE) 及其改进版本 beta-VAE 的理论原理与实现方法。VAE 通过编码器将高维数据压缩至低维潜在空间，并利用解码器重构数据，其核心在于优化重构损失与 KL 散度之间的平衡。beta-VAE 通过引入超参数 β β β 调节这一平衡，增强了学习解耦表示的能力。在二次曲线生成任务中，beta-VAE 成功将带噪曲线压缩至二维潜在空间，并实现了平滑曲线的重构；在 MNIST 数据集上，尽管二维潜在空间限制了生成图像的清晰度，但模型仍能有效学习数字的语义特征。实验表明，beta-VAE 系列模型在无监督特征学习和数据生成方面具有独特优势，但其生成质量受潜在维度与 β β β 值影响显著，需根据具体任务权衡取舍。
系列链接
PyTorch计算机视觉（1）------计算机视觉的数学工具
PyTorch计算机视觉（2）------神经网络模型训练与PyTorch基础
PyTorch计算机视觉（3）------卷积神经网络（CNN）详解与实现
PyTorch计算机视觉（4）------迁移学习（Transfer Learning）详解与实现
PyTorch计算机视觉（5）------生成对抗网络（Generative Adversarial Network，GAN）
PyTorch计算机视觉（6）------深度卷积对抗神经网络（DCGAN）