AutoEncoder(AE)中文叫自编码器,是一种介于监督学习和无监督学习之间的模型。其学习方式现在被称为自监督学习,类似于Bert模型,它的训练需要标签数据,但是这个标签数据是自己生成的,因此叫做自监督学习。
AutoEncoder可以实现很多有用的事情,比如数据降维、图像降噪、风格迁移等。今天我们用AutoEncoder实现人脸生成的例子,这个实际上不是非常常见的例子。
AutoEncoder的网络结构非常简单,由Encoder和Decoder两部分组成,Encoder对数据不断降维,而Decoder对降维后的数据不断升维。最后的希望Encoder的输入和Decoder的输出越接近越好。如下图:
根据问题的不同Encoder和Decoder的结构也可以自由更改。比如针对简单的问题,Encoder和Decoder可以用全连接实现,逐层降维而后逐层升维。如果处理图像数据,可以用卷积神经网络来实现Encoder和Decoder。在本文中,我们处理的是人脸图像,因此选取卷积神经网络实现Encoder和Decoder。
在开始生成人脸之前,需要用人脸数据训练一个AutoEncoder网络。这里使用CelebA数据集作为训练数据,数据下载地址为:mmlab.ie.cuhk.edu.hk/projects/Ce… 。选择In The Wild Images即可,不过默认使用的是Google云盘,如果不能访问可以选择百度网盘下载:pan.baidu.com/s/1eSNpdRG?at=1678847322615#list/path=%2F 。
这里使用PyTorch来实现,首先看看需要用到的模块:
import torch
import numpy as np
from torch import nn
from PIL import Image
from torch import optim
from torchvision.utils import make_grid
from torch.utils.data import DataLoader, Dataset
import matplotlib.pyplot as plt
其中make_grid的作用是把多张图片合并成一张,方便后续显示。
为了方便数据读取,我们实现一个FaceDataset类,实现__getitem__
和__len__
方法,代码如下:
class FaceDataset(Dataset):
def __init__(self, data_dir="/home/zack/Files/datasets/img_align_celeba", image_size=64):
self.image_size = image_size
# 图像预处理
self.trans = transforms.Compose([
transforms.Resize(image_size),
transforms.CenterCrop(image_size),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
# 保持所有图片的路径
self.image_paths = []
# 读取根目录,把所有图片路径放入image_paths
for root, dirs, files in os.walk(data_dir):
for file in files:
self.image_paths.append(os.path.join(root, file))
def __getitem__(self, item):
# 读取图片,并预处理
return self.trans(Image.open(self.image_paths[item]))
def __len__(self):
return len(self.image_paths)
有了上面的Dataset我们就可以获取DataLoader对象了:
dataset = FaceDataset()
# 创建DataLoader对象,设置batch_size为64
dataloader = DataLoader(dataset, 64)
我们可以测试一下:
for data in dataloader:
print(data.shape)
break
# 输出
# torch.Size([64, 3, 64, 64])
下一步就是构建网络了。
这里使用卷积神经网络实现Encoder和Decoder,代码如下:
class FaceAutoEncoder(nn.Module):
def __init__(self, encoded_dim=1024):
super(FaceAutoEncoder, self).__init__()
# [b, 3, 64, 64] --> [b, encoded_dim, 1, 1]
self.encoder = nn.Sequential(
nn.Conv2d(3, 64, 4, 2, 1, bias=False),
nn.LeakyReLU(0.2, inplace=True),
nn.Conv2d(64, 64 * 2, 4, 2, 1, bias=False),
nn.BatchNorm2d(64 * 2),
nn.LeakyReLU(0.2, inplace=True),
nn.Conv2d(64 * 2, 64 * 4, 4, 2, 1, bias=False),
nn.BatchNorm2d(64 * 4),
nn.LeakyReLU(0.2, inplace=True),
nn.Conv2d(64 * 4, 64 * 8, 4, 2, 1, bias=False),
nn.BatchNorm2d(64 * 8),
nn.LeakyReLU(0.2, inplace=True),
nn.Conv2d(64 * 8, encoded_dim, 4, 1, 0, bias=False),
nn.LeakyReLU(0.2, inplace=True)
)
# [b, encoded_dim, 1, 1] - > [b, 3, 64, 64]
self.decoder = nn.Sequential(
nn.ConvTranspose2d(encoded_dim, 64 * 8, 4, 1, 0, bias=False),
nn.BatchNorm2d(64 * 8),
nn.ReLU(True),
nn.ConvTranspose2d(64 * 8, 64 * 4, 4, 2, 1, bias=False),
nn.BatchNorm2d(64 * 4),
nn.ReLU(True),
nn.ConvTranspose2d(64 * 4, 64 * 2, 4, 2, 1, bias=False),
nn.BatchNorm2d(64 * 2),
nn.ReLU(True),
nn.ConvTranspose2d(64 * 2, 64, 4, 2, 1, bias=False),
nn.BatchNorm2d(64),
nn.ReLU(True),
nn.ConvTranspose2d(64, 3, 4, 2, 1, bias=True),
nn.Tanh()
)
def forward(self, x):
x = self.encoder(x)
x = self.decoder(x)
return x
Encoder部分主体为卷积层,同设置stride为2达到缩小特征图的效果,因此省去了pooling层。Pytorch处理图像时,四维数据的各个维度分别表示batch_size、通道数、高、宽
。Encoder的输入形状为(batch_size * 3 * 64 * 64)
,输出形状为(batch_size * encoded_dim * 1 * 1)
,其中encoded_dim为图像编码后的维度。
Decoder部分和Encoder类似,把卷积改为转置卷积,不断减少通道数,扩大特征图尺寸。Decoder输入形状为(batch_size * encoded_dim * 1 * 1)
,正好和Encoder输出一致,输出形状为(batch_size * 3 * 64 * 64)
,正好和Encoder的输入一致。因为图像预处理后的范围存在负值,因此输出层使用Tanh而不是Sigmoid。
下面就可以开始训练了,在加载数据是,我们的Dataset对象只返回了一个值,这个值既是特征值也是目标值,因此叫自监督学习。代码如下:
# 使用gpu
device = "cuda:0" if torch.cuda.is_available() else "cpu"
# 超参数
lr = 0.0001
batch_size = 64
image_size = 64
encoded_dim = 1024
epochs = 20
# 加载数据
dataset = FaceDataset(image_size=image_size)
dataloader = DataLoader(dataset, batch_size)
# 构建模型
model = FaceAutoEncoder(encoded_dim=encoded)
# 定义loss和优化器
optimizer = optim.Adam(model.parameters(), lr=lr)
criterion = nn.MSELoss()
# 训练
model.to(device)
if not model.training:
model.train()
for epoch in range(epochs):
for idx, data in enumerate(dataloader):
data = data.to(device)
# 目标值和特征值为同一批数据
target = data
# 正向传播
output = model(data)
loss = criterion(output, target)
# 反向传播
optimizer.zero_grad()
loss.backward()
optimizer.step()
# evaluate
if idx % 300 == 0:
with torch.no_grad():
# 编码解码
outputs = model(data)
# 讲解码后的图片按8行8列组成一张
outputs = make_grid(outputs.cpu(), normalize=True)
# pytorch中的图像是通道在第一维,这里把通道放到最后一个维度
outputs = outputs.numpy().transpose((1, 2, 0))
plt.imshow(outputs)
plt.show()
print("epoch: %s, train_loss: %.4f" % (
epoch, loss.item()
))
# 每个epoch都保存一次模型
torch.save(model.state_dict(), 'face_auto_encoder.pth')
训练完成后得到效果如下图,从左到右训练轮次依次增加。最后一张图可以明显看到人脸情况,说明我们的AutoEncoder能很好的学习如何编码一张人脸。
不过现在我们还不是在生成人脸,而是在对人脸图片编码和解码。下面我们来看看如何用AutoEncoder生成人脸。
AutoEncoder分为Encoder和Decoder两个部分,Encoder可以把人脸图像转换成1024维的向量,而Decoder可以由一个1024维的向量生成一个人脸图像。由此,可以把生成人脸的问题看做是获取一个1024维向量的问题。
现在我们做一个测试,如果给Decoder一个随机的1024维向量,能否给我们返回一个人脸图像,我们用下面代码进行测试:
zdim = 1024
model = FaceAutoEncoder(encoded_dim=zdim)
model.load_state_dict(torch.load('face_auto_encoder.pth'))
model.eval()
# 生成符合标准整体分布的数据
z1 = torch.randn(64, 1024, 1, 1)
# 生成符合均分分布的数据
z2 = torch.rand(64, 1024, 1, 1)
z = [z1, z2]
with torch.no_grad():
for encoded in z:
# 使用解码器生成人脸
outputs = model.decoder(encoded)
outputs = torch.clamp(outputs, 0, 255)
grid = make_grid(outputs).numpy().transpose((1, 2, 0))
plt.imshow(grid)
plt.show()
下面是得到的结果,左边是用正太分布的向量生成的,右边是用均匀分布的向量生成的,完全看不出来人脸。
现在我们做一个假设,即假设人脸向量符合一个多元正太分布,此时我们只需要求出这个分布的均值和协方差矩阵就可以得到这个分布。得到这个分布后,就可以用这个分布采样人脸向量。此时采样到的人脸向量很有可能可以解码出一个人脸。
那么均值和协方差应该怎么求呢?一种非常简单的办法就是通过统计得到,代码如下:
zdim = 1024
# 加载数据
dataloader = DataLoader(
FaceDataset(image_size=64),
batch_size=128
)
# 构建模型
model = FaceAutoEncoder(encoded_dim=zdim)
model.load_state_dict(torch.load('face_auto_encoder.pth'))
# 生成mean和cov
mean = np.zeros((zdim,), dtype=np.float32)
cov = np.zeros((zdim, zdim), dtype=np.float32)
device = "cuda" if torch.cuda.is_available() else "cpu"
model.eval()
model = model.to(device)
with torch.no_grad():
for idx, data in enumerate(dataloader):
data = data.to(device)
# 对人脸数据进行编码
encoded = model.encoder(data).view(128, -1)
# 求人脸向量的均值并累加
mean += encoded.mean(axis=0).cpu().numpy()
# 求人脸向量的协方差矩阵并累加
cov += np.cov(encoded.cpu().numpy().T)
if idx % 50 == 0:
print(f"\ridx: {idx}/{len(dataloader)}", end="")
# 归一化
mean /= (idx + 1)
cov /= (idx + 1)
np.savez('face_distribution.npz', mean, cov)
这里可以用np.cov求协方差矩阵,最后把均值和协方差矩阵保存成一个npz文件。
现在我们以及估计出了人脸向量对应的正太分布,可以用np.random.multivariate_normal
来对这个正太分布进行采样,然后把采样结果交给Decoder部分。具体代码如下:
zdim = 1024
model = FaceAutoEncoder(encoded_dim=zdim)
model.load_state_dict(torch.load('face_auto_encoder.pth'))
model.eval()
# 加载人脸分布
distribution = np.load('face_distribution.npz')
mean = distribution['arr_0']
cov = distribution['arr_1']
# 生成编码向量
batch_size = 64
z = np.random.multivariate_normal(
mean,
cov,
batch_size
).astype(np.float32)
# 解码
with torch.no_grad():
encoded = torch.from_numpy(z).view(batch_size, zdim, 1, 1)
outputs = model.decoder(encoded)
# 把像素值映射到0-255
outputs = torch.clamp(outputs, 0, 255)
grid = make_grid(outputs).numpy().transpose((1, 2, 0))
plt.imshow(grid)
plt.show()
下面是生成的一部分人脸:
虽然效果不是那么好,但是也生成了人脸。不过这个效果还可以继续改进,在后续的文章中再继续给大家分享。
阅读量:486
点赞量:0
收藏量:0