变分自编码器（VAE）编程实现
一、简介
本文使用VAE的方法去学习MNIST手写数字图像。
使用的是Pytorch深度学习框架。
MNIST手写数字图像大小是1×28×28，将其展平为一个784维向量，这个向量就是输入x。
然后经过编码器得到潜在空间的特征。编码器由全连接层组成。首先经过一个512维的全连接层，然后经过一个256维的全连接层，最后分别经过一个2维的全连接层，得到潜在特征的均值和方差。
再次将编码器得到的均值和方差通过公式变换得到二维的潜在特征z。
最后通过解码器得到由z生成的图像y。解码器由全连接层组成。首先经过一个256维的全连接层，然后经过一个512维的全连接层，最后经过一个784维的全连接层得到生成的图像y。
我们的目的是生成的图像y和输入图像x是近似相同的，越越相似越好。
二、编程实现
根据上述写VAE模型。
import torch
import torch.nn as nn
class VAE(nn.Module):
def __init__(self, x_dim, h_dim1, h_dim2, z_dim):
super(VAE, self).__init__()
# 编码器
self.fc1 = nn.Linear(x_dim, h_dim1)
self.fc2 = nn.Linear(h_dim1, h_dim2)
self.fc31 = nn.Linear(h_dim2, z_dim)
self.fc32 = nn.Linear(h_dim2, z_dim)
# 解码器
self.fc4 = nn.Linear(z_dim, h_dim2)
self.fc5 = nn.Linear(h_dim2, h_dim1)
self.fc6 = nn.Linear(h_dim1, x_dim)
# 编码器
def encoder(self, x):
h = nn.ReLU()(self.fc1(x))
h = nn.ReLU()(self.fc2(h))
return self.fc31(h), self.fc32(h)
# 换算z的公式
def sampling(self, mu, log_var):
std = torch.exp(0.5 * log_var)
eps = torch.randn_like(std) # 生成std形状大小的随机标准正态分布的噪声
return eps.mul(std).add_(mu) # 得到z
# 解码器
def decoder(self, z):
h = nn.ReLU()(self.fc4(z))
h = nn.ReLU()(self.fc5(h))
return nn.Sigmoid()(self.fc6(h))
def forward(self, x):
mu, log_var = self.encoder(x.view(-1, 784))
z = self.sampling(mu, log_var)
return self.decoder(z), mu, log_var
写测试程序。我是用多gpu进行训练，当然可以修改为使用单个gpu训练，或者使用cpu训练。
import os
import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
from torchvision import datasets, transforms
from vae import VAE
import torch.nn.functional as F
# 设置gpu
os.environ['CUDA_VISBLE_DEVICES'] = '0, 1'
# 数据集预处理
transforms = transforms.Compose([
transforms.ToTensor()
])
# 下载数据集
trainDataset = datasets.MNIST(root='./mnist_data/', train=True, transform=transforms, download=True)
testDataset = datasets.MNIST(root='./mnist_data/', train=False, transform=transforms, download=False)
# 加载数据集
trainLoader = torch.utils.data.DataLoader(dataset=trainDataset, batch_size=128, shuffle=True, num_workers=8)
testLoader = torch.utils.data.DataLoader(dataset=testDataset, batch_size=128, shuffle=False, num_workers=8)
# 创建模型
vae = VAE(784, 512, 256, 2)
if torch.cuda.is_available():
vae = nn.DataParallel(vae)
vae.cuda()
# 优化
optimizer = torch.optim.SGD(vae.parameters(), lr=1e-4, weight_decay=5e-4)
# 损失函数
def lossFunction(recon_x, x, mu, log_var):
bce = F.binary_cross_entropy(recon_x, x.view(-1, 784), reduction='sum')
kld = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp())
return bce + kld
# 训练函数
def train(epoch):
vae.train()
for batch_idx, (data, _) in enumerate(trainLoader):
data = data.cuda()
recon_batch, mu, log_var = vae(data)
loss = lossFunction(recon_batch, data, mu, log_var)
optimizer.zero_grad()
loss.backward()
optimizer.step()
if batch_idx % 128 == 0:
print('Train Epoch: {} [{}/{} ({:0f}%)]\tLoss: {:.6f}'.format(epoch, batch_idx * len(data),
len(trainLoader.dataset), 100. * batch_idx / len(trainLoader), loss.item() / len(data)))
# 测试函数
def test():
vae.eval()
testLoss = 0
with torch.no_grad():
for data, _ in testLoader:
data = data.cuda()
recon, mu, log_var = vae(data)
testLoss += lossFunction(recon, data, mu, log_var).item()
testLoss /= len(testLoader.dataset)
print('=====> Test set loss: {:.4f}'.format(testLoss))
# 训练
for epoch in range(500):
train(epoch)
test()
# 潜在变量z通过解码器生成图片
with torch.no_grad():
n = 20
digitSize = 28
figure = np.zeros((digitSize * n, digitSize * n))
gridX = np.linspace(-4, 4, n)
gridY = np.linspace(-4, 4, n)
for i, xi in enumerate(gridX):
for j, yi in enumerate(gridY):
zSample = torch.Tensor([[yi, xi]]).cuda()
xDecoder = vae.module.decoder(zSample).cuda()
digit = xDecoder[0].cpu().detach().numpy().reshape(digitSize, digitSize)
figure[(n-i-1)*digitSize: (n-i) * digitSize, j*digitSize:(j+1) *digitSize] = digit
plt.figure(figsize=(10, 10))
plt.imshow(figure)
plt.savefig("Result")
plt.show()
三、实验结果
最后对训练的模型进行测试。潜在特征是个二维向量，即只有两个数。我们通过给不同的z来生成图像。
如图是实验的结果。横纵坐标分别代表潜在特征z的二维向量的其中一个值。最后发现，同一个数字的图像，潜在特征z的值是相近的。
同一个数字的图像在潜在空间上是聚集在一起的。