PyTorch实现人脸数据集处理教程

释放双眼,带上耳机,听听看~!
本教程介绍了如何使用PyTorch对人脸数据集进行处理,包括模块导入、数据加载和预处理。通过实现FaceDataset类和使用DataLoader对象,可以方便地处理人脸图像数据。

 本文正在参加mmlab.ie.cuhk.edu.hk/projects/Ce… 。选择In The Wild Images即可,不过默认使用的是Google云盘,如果不能访问可以选择百度网盘下载:pan.baidu.com/s/1eSNpdRG?at=1678847322615#list/path=%2F

3.1 模块导入

这里使用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的作用是把多张图片合并成一张,方便后续显示。

3.2 加载数据

为了方便数据读取,我们实现一个FaceDataset类,实现__getitem____len__方法,代码如下:

class FaceDataset(Dataset):
    def __init__(self, data_dir="/home/zack/Files/datasets/img_align_celeba", image_size=64):
        self.jpg># 图像预处理
        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])

下一步就是构建网络了。

3.3 构建模型

这里使用卷积神经网络实现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。

3.4 模型训练

下面就可以开始训练了,在加载数据是,我们的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能很好的学习如何编码一张人脸。

PyTorch实现人脸数据集处理教程

不过现在我们还不是在生成人脸,而是在对人脸图片编码和解码。下面我们来看看如何用AutoEncoder生成人脸。

四、AutoEncoder生成人脸

AutoEncoder分为Encoder和Decoder两个部分,Encoder可以把人脸图像转换成1024维的向量,而Decoder可以由一个1024维的向量生成一个人脸图像。由此,可以把生成人脸的问题看做是获取一个1024维向量的问题。

4.1 使用随机向量

现在我们做一个测试,如果给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()

下面是得到的结果,左边是用正太分布的向量生成的,右边是用均匀分布的向量生成的,完全看不出来人脸。

PyTorch实现人脸数据集处理教程

4.2 估计人脸向量的分布

现在我们做一个假设,即假设人脸向量符合一个多元正太分布,此时我们只需要求出这个分布的均值和协方差矩阵就可以得到这个分布。得到这个分布后,就可以用这个分布采样人脸向量。此时采样到的人脸向量很有可能可以解码出一个人脸。

那么均值和协方差应该怎么求呢?一种非常简单的办法就是通过统计得到,代码如下:

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文件。

4.3 生成人脸

现在我们以及估计出了人脸向量对应的正太分布,可以用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()

下面是生成的一部分人脸:

PyTorch实现人脸数据集处理教程

虽然效果不是那么好,但是也生成了人脸。不过这个效果还可以继续改进,在后续的文章中再继续给大家分享。

本文正在参加 人工智能创作者扶持计划

本网站的内容主要来自互联网上的各种资源,仅供参考和信息分享之用,不代表本网站拥有相关版权或知识产权。如您认为内容侵犯您的权益,请联系我们,我们将尽快采取行动,包括删除或更正。
AI教程

机器学习的分布式训练方法及其优势-亚马逊云科技

2023-12-18 19:04:14

AI教程

清华大学 TSAIL 团队发布文生3D新算法ProlificDreamer,无需3D数据即可生成高质量3D内容

2023-12-18 19:14:14

个人中心
购物车
优惠劵
今日签到
有新私信 私信列表
搜索