Amazon SageMaker 使用自编码器完成人脸生成

| 作者介绍：ZackSock，iVX开发工程师，高校人工智能、数据挖掘讲师，csdn博客专家、2020年腾讯云年度最佳作者。专注自然语言处理相关应用的研究。微信公众号：新建文件夹X | | | --- | --- | # 一、前言 最近受邀参与了亚马逊云科技【云上探索实验室】活动，体验了一下 Amazon SageMaker 平台，训练了一个人脸的自编码器。对比本地的训练时间，速度提升比较明显。本文将围绕 Amazon SageMaker 和自编码器进行。 自编码器是一个非常简单的网络，早在上世纪90年代就提出了自编码器的概念。当时使用受限的玻尔兹曼机分层训练，在硬件强大的今天可以实现端到端的训练。自编码器有许多变种，比如变分自编码器、去噪自编码器、正则自编码器等。由于自编码器采用的是自监督学习，因此使用自编码器可以使用较低成本得到较好的效果。 今天我们将训练一个自编码器。自编码器训练完成后可以进行各种有趣的实验，像编辑一个人的表情、让人脸A渐渐变成人脸B、让一个人从小慢慢变老、生成人脸等。  在本文我们会做一部分实验，实现人脸渐变和生成人脸的操作。 ## 1.1 什么是Amazon SageMaker Amazon SageMaker 是一个完全托管的机器学习服务平台，包含了机器学习的各个流程，从标注到部署。开发人员可以快速构建模型并训练，还可以部署到托管环境。Amazon SageMaker 提供了 Jupyter 笔记本，而且可以执行各种流行框架，不止是 MXNet，还可以使用 yTorch、TensorFlow 等主流框架。 ## 1.2 有什么特别的 在数据标注时，Ground Truth 可以用来团队标注，在标注一定数据后，Ground Truth 可以自动标注，当对标注不确定时才会让人工进行标注。 Amazon SageMaker 提供了数据存储，模型部署等服务，完成这些操作都可以一键式完成。 同时Amazon SageMaker提供了许多高阶 API，可以使用 Amazon Auto Pilot 自动训练模型与调优模型。同时还可以对模型的情况进行监控，以便更好地改善模型。 # 二、机器学习流程 ## 2.1 机器学习整体流程 在开始使用 Amazon SageMaker 完成机器学习任务前，先熟悉一下机器学习的流程。机器学习的流程分为一下几个步骤： 1. 获取数据 2. 数据清洗 3. 将数据处理成模型的输入 4. 训练模型 5. 评估模型 6. 部署模型 7. 监控、收集数据、评估模型 上述步骤整体是一条直线，但是经常会回到前面几个步骤重新往下执行。如下图所示：  ### （1）数据处理 1-3为数据处理部分，主要是数据从无到有、从杂乱到规范，在此步骤中，我们会用到诸如爬虫、正则、归一化、标准化等技术。处理后的数据通常表示为一个张量（多维数组），不同类型的数据的形状有所不同。图像数据通常为 `n_sample × channel × height × width`，各维度含义分别为样本数、通道数、高、宽，有时候通道数会放在最后一个维度。表格数据被处理为 `n_sample × n_feature`，分别为样本数、特征数。文本数据则是 `n_sample × sequence_len × 1`，分别为样本数和文本长度。还有时间序列、视频、立体图片等数据都会处理成固定形状的张量。 数据处理成张量后还可以做一些特征工程，比如特征选择、标准化、归一化等。这一步操作有利于模型训练。 ### （2）训练模型 4-5为训练评估部分，在处理好数据后就可以开始训练模型了。训练模型需要我们确定几样东西，分别是模型结构（选取上面算法）、对应的超参数、优化器等。在此过程通常会训练多个模型，选取一个最优模型作为最终模型。 为了选择较优的超参数，会使用到网格搜索交叉验证。在 sklearn 模块中，提供了对应的实现。SageMaker 中也有类似更强的功能，后面有具体体验。 ### （3）部署及维护 当选取一个最优模型后，既可以将模型部署上线，模型的部署可以作为 API、Web、小程序、APP 等。在部署上线后，模型可能会出现各种问题，也会慢慢落后，因此还需要对模型进行监控维护。在上线后，可以继续收集一些用户授权的数据，然后把这些数据重复之前的步骤，对模型进行迭代优化。 ## 2.2 SageMaker 中机器学习流程 SageMaker 中机器学习流程于上面一致，下面我们来实际看看各个步骤如何操作。 ### （1）数据处理 Amazon SageMaker 中可以创建 Jupyter Notebook，创建的 Notebook 可以执行 pip、wget 等指令。我们可以使用以往使用的所有方式处理数据，也可以使用 SageMaker 自带的 SDK。在 SageMaker 中有一个 sklearn 子模块，可以用来处理数据。 SKLearnProcessor 可以执行 sklearn 脚本来处理数据，首先需要创建SKLearnProcessor对象，然后调用 run 方法对数据进行处理，示例代码如下： ```js from sagemaker import get_execution_role from sagemaker.sklearn.processing import SKLearnProcessor from sagemaker.processing import ProcessingInput, ProcessingOutput role = get_execution_role() sklearn_processor = SKLearnProcessor( framework_version="0.20.0", role=role, instance_type="ml.t3.medium", instance_count=1 ) sklearn_processor.run( code="preprocessing.py", inputs=[ ProcessingInput(source="s3://your-bucket/path/to/your/data", destination="/opt/ml/processing/input"), ], outputs=[ ProcessingOutput(output_name="train_data", source="/opt/ml/processing/train"), ProcessingOutput(output_name="test_data", source="/opt/ml/processing/test"), ], arguments=["--train-test-split-ratio", "0.2"], ) preprocessing_job_description = sklearn_processor.jobs[-1].describe() ``` 在创建调用 run 方法时，需要先创建好用于处理数据的 `preprocessing.py` 文件，对应 code 参数。`preprocessing.py` 命令行参数由 arguments 给出。 `preprocessing.py` 可以用 sklearn 来做具体处理。处理完成后，处理结果等信息保存在 preprocessing_job_description 中，可以通过 preprocessing_job_description['Outputs'] 拿到处理结果。 ### （2）训练模型 模型的训练与处理数据一样，需要准备好对应的训练脚本 `train.py`。处理好数据并准备好 `train.py` 脚本后，使用下面的代码就可以进行训练了，这里的实例类型要和数据处理的实例类型一致。 ```js from sagemaker.sklearn.estimator import SKLearn sklearn = SKLearn( entry_point="train.py", framework_version="0.20.0", instance_type="ml.t3.medium", role=role ) sklearn.fit({"train": preprocessed_training_data}) training_job_description = sklearn.jobs[-1].describe() model_data_s3_uri = "{}{}/{}".format( training_job_description["OutputDataConfig"]["S3OutputPath"], training_job_description["TrainingJobName"], "output/model.tar.gz", ) ``` 评估模型的代码也是一样的风格： ```js sklearn_processor.run( code="evaluation.py", inputs=[ ProcessingInput(source=model_data_s3_uri, destination="/opt/ml/processing/model"), ProcessingInput(source=preprocessed_test_data, destination="/opt/ml/processing/test"), ], outputs=[ProcessingOutput(output_name="evaluation", source="/opt/ml/processing/evaluation")], ) evaluation_job_description = sklearn_processor.jobs[-1].describe() ``` ### （3）部署 构建和训练模型后，您可以将模型部署至终端节点，以中获取预测推理结果。 ## 2.3 实战 # 三、自编码器 下面使用 SageMaker 完成自编码相关的实验。 ## 3.1 自编码器介绍 自编码器是一个非常简单网络，通常由编码器和解码器两个部分组成。编码器解码器的结构可以用全连接，也可以用卷积，或者其它一些网络。在早期自编码器的编码器解码器需要分开训练，而现在通常是端到端的训练。 编码器部分会将输入逐步降维，最后得到一个固定长度的向量，这个向量可以作为输入数据的编码。解码器部分接收编码器的输出，结果解码器会得到一个形状与编码器输入一样的数据。自编码器训练的目的就是输出与数据尽可能接近。 整体上看自编码器使用的是监督学习方法，但是目标值和特征值是一样的，像这种标签由数据自身给出的学习方法又被称为自监督学习。 自编码器结构简单，但是有一些非常好的性质。比如训练简单，不需要人工标记数据等。使用自编码器可以对数据进行降维，创建新数据等。假设对人脸图像数据进行编码，编码器会得到1024维的向量。那么编码器得到的这个向量每一个维度可能代表着一种特征，比如第n个维度可能代表表情、第k个维度可能代表性别等。如果能知道这些信息，就可以做一些有趣的事情。 ## 3.2 环境准备 这里使用 Amazon SageMaker 的笔记本实例进行实验，创建笔记本实例，创建时使用默认选项即可。在创建时需要记住使用的实例类型，后续训练需要对应正确的类型。 创建好环境后，可以在笔记本中运行下面代码获取当前角色以及 S3 桶： ```js import sagemaker import os sess = sagemaker.Session() role = sagemaker.get_execution_role() bucket = sagemaker_session.default_bucket() ``` 下面就可以开始数据的准备了。 ## 3.3 数据处理 这里使用 PyTorch 完成数据的处理和训练，下面是需要用到的一些模块： ```js from torch import nn from torch import optim from torch.utils.data import Dataset, DataLoader from torchvision import datasets from torchvision import transforms ``` 本实验使用的是 CelebA Dataset，可以在http://mmlab.ie.cuhk.edu.hk/projects/CelebA.html?trk=cndc-detail 下载。数据集里面包含了10万修正好的人脸图片。将图片下载到电脑上后，再把数据上传到笔记本中。为了方便加载，创建一个 Dataset 类，完成数据的加载： ```js class FaceDataset(Dataset): def __init__(self, data_dir="./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 = [] if os.path.isdir(data_dir): for root, dirs, files in os.walk(data_dir): for file in files: self.image_paths.append(os.path.join(root, file)) elif os.path.isfile(data_dir): with open(data_dir, encoding='utf-8') as f: self.image_paths = f.read().strip().split() def __getitem__(self, item): return self.trans(Image.open(self.image_paths[item])) def __len__(self): return len(self.image_paths) ``` 下面创建 DataLoader 用于加载数据： ```js dataset = FaceDataset(data_dir="./datasets/img_align_celeba") dataloader = DataLoader(dataset, 64) ``` 另外，可以通过下面代码将数据上传至S3桶： ```js inputs = sagemaker_session.upload_data(path="datasets/img_align_celeba", bucket=bucket, key_prefix='sagemaker/img_align_celeba') ``` 数据准备好后，需要编写模型以及训练脚本。 ## 3.2 模型训练 本文以人脸数据训练一个自编码器，而后用这个自编码器做一些其它事情，首先训练一个自编码器。网络结构由卷积和转置卷积组成，代码如下： ```js class FaceAutoEncoder(nn.Module): def __init__(self, encoded_dim=1024): super(FaceAutoEncoder, self).__init__() # [b, 3, 64, 64] --> [b, 1024, 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, 1024, 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 ``` 编码器部分将64×64的图片转换成1024维的向量，而解码器则是利用转置卷积将1024维的向量转换成64×64的图像。现在我们的目标是模型输入与数据越接近越好，所有可以用均方误差作为损失函数，同时输入和目标为同一批数据。下面编写一个训练脚本 `autoencoder.py`，代码如下： ```js def train(args): device = "cuda" if torch.cuda.is_available() else "cpu" dataset = FaceDataset() dataloader = DataLoader(dataset, 64) model = FaceAutoEncoder().to(device) optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=0.9) for epoch in range(1, args.epochs + 1): model.train() for batch_idx, data in enumerate(dataloader, 1): data = data.to(device) optimizer.zero_grad() output = model(data) loss = F.mse_loss(output, data) loss.backward() optimizer.step() save_model(model, args.model_dir) def model_fn(model_dir): device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = torch.nn.DataParallel(AutoEncoder()) with open(os.path.join(model_dir, "model.pth"), "rb") as f: model.load_state_dict(torch.load(f)) return model.to(device) def save_model(model, model_dir): path = os.path.join(model_dir, "model.pth") torch.save(model.decoder.cpu().state_dict(), path) if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument( "--epochs", type=int, default=10, metavar="N", help="number of epochs to train (default: 10)", ) parser.add_argument( "--lr", type=float, default=0.01, metavar="LR", help="learning rate (default: 0.01)" ) train(parser.parse_args()) ``` 我们可以使用 argparse 设置一些参数。脚本准备好后，可以开始进行训练，因为使用的是 PyTorch，所以要改用 sagemaker.pytorch 中的 PyTorch 来进行训练，代码如下： ```js from sagemaker.pytorch import PyTorch estimator = PyTorch( entry_point="autoencoder.py", role=role, py_version="py38", framework_version="1.11.0", instance_count=1, instance_type="ml.c5.2xlarge", hyperparameters={"epochs": 4}, ) estimator.fit({"training": inputs}) ``` 其中 entry_point 就是前面的训练脚本。然后等待模型训练，训练完成后可以得到模型文件 model.pth，这里只包含 ecoder 部分。下面是输出的结果：  训练完成后使用下面一句就可以对模型进行部署： ```js predictor = estimator.deploy(initial_instance_count=1, instance_type="ml.m4.xlarge") ``` 不过自编码器本身只是还原原有数据而已，要生成人脸需要使用 Decoder 部分，进行推理。 # 四、使用自编码器实现人脸渐变 ## 4.1 人脸渐变原理 在前面提到自编码器的一个优点就是可以把对输入的操作改为对编码的操作，现在我们训练了一个人脸的自编码。假设人脸A被编码成z1，人脸B被编码成z2，现在想让人脸由A到B渐变。现在可以把这个问题转换为向量z1和z2之间的渐变，向量的渐变可以直接使用插值算法，我们在两个向量见插入n个向量，再把这些向量输入解码器，得到的人脸图像就是介于A和B之间的人脸。现在人脸渐变就变为了插值。 ## 4.2 人脸渐变的实现 首先实现插值算法，插值的实现很简单，具体代码如下： ```js def interpolate(x1, x2, num): result = torch.zeros((num, *x1.shape)) step = (x2 - x1) / (num - 1) for i in range(num): result[i] = x1 + step * i return result ``` 上面函数输入两个长度一样的向量，输出 num 个向量。这 num 个向量将作为 Decoder 的输入。接下来使用 Decoder 部分进行推理： ```js # 加载数据集 dataloader = DataLoader( FaceDataset(data_dir='/home/zack/Files/datasets/img_align_celeba', image_size=64), batch_size=2 ) model = FaceAutoEncoder() model.load_state_dict(torch.load('../outputs/face_auto_encoder.pth')) model.eval() with torch.no_grad(): for idx, data in enumerate(dataloader): # 对人脸编码 encoded1 = model.encoder(data[0].reshape((1, 3, 64, 64))) encoded2 = model.encoder(data[1].reshape((1, 3, 64, 64))) # 对人脸编码进行插值 encoded = interpolate(encoded1[0], encoded2[0], 64) # 解码成人脸 outputs = model.decoder(encoded).reshape((64, 3, 64, 64)) outputs = make_grid(outputs, normalize=True) plt.imshow(outputs.numpy().transpose((1, 2, 0))) plt.show() ``` 下面是实现的效果：  可以看到人脸转变的很自然。现在可以将这一部分部署，部署代码如下： ```js from sagemaker.pytorch import PyTorchModel pytorch = PyTorchModel( model_data=model_data, role=role, entry_point="inference.py", source_dir="code", framework_version="1.3.1", py_version="py3", sagemaker_session=sagemaker_session, ) predictor = pytorch.deploy( initial_instance_count=1, instance_type='ml.m5.large', endpoint_name=endpoint_name, wait=True, ) ``` 其中 `inference.py `是推理脚本，具体代码如下： ```js def predict_fn(input_data, model): device = torch.device("cuda" if torch.cuda.is_available() else "cpu") input_data = input_data.to(device) with torch.no_grad(): inputs = interpolate(input_data[:, 0], input_data[:, 1]) return model.decoder.forward(inputs) def model_fn(model_dir="model.pth"): loaded_model = torch.jit.load(model_dir, map_location=torch.device("cpu")) loaded_model = loaded_model.eval() return loaded_model ``` 然后使用 `predictor.predict` 传入2 × 3 × 64 × 64的张量就可以完成模型的推理，返回插值后的人脸图像。 # 五、自编码器生成人脸 自编码器除了实现人脸渐变外，还可以用来生成人脸，关键点依然在编码部分。 ## 5.1 人脸分布 自编码器生成人脸的原理比较简单，在训练自编码器时，把人脸编码成一个长度为1024维的向量。现在我们假设人脸服从高斯分布，如果能求出均值和方差，就可以知道这个高斯分布的具体样子。在知道高斯分布的具体表达式后，就可以对从中采样人脸向量，把这个向量交给 decoder 就可以生成人脸。 均值和方差可以用统计的方式获取，具体代码如下，把结果保存为一个npz文件： ```js 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): try: 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="") except Exception as e: pass mean /= (idx + 1) cov /= (idx + 1) np.savez('face_distribution.npz', mean, cov) ``` ## 5.2 生成人脸 生成人脸的操作就是从前面的高斯分布中进行采样，然后把采样的向量交给 decoder 进行编码即可。具体代码： ```js # 加载人脸分布 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, 1024, 1, 1) outputs = model.decoder(encoded) outputs = torch.clamp(outputs, 0, 255) grid = make_grid(outputs).numpy().transpose((1, 2, 0)) plt.imshow(grid) plt.show() ``` 下面是生成的一些人脸，五官部分可以看的很清楚，但是背景部分有一些模糊。  现在前面 Amazon SageMaker 部署的代码可以复用，只需要修改推理的代码即可。把 inference.py 中 predict_fn 函数修改成下面的样子： ```js def predict_fn(input_data, model): device = torch.device("cuda" if torch.cuda.is_available() else "cpu") z = torch.randn(1, 1024, 1, 1).to(device) with torch.no_grad(): return model.decoder.forward(z) ``` ## 5.3 生成相关的模型 在前面，我们使用 Amazon SageMaker，围绕自编码器实现了人脸渐变和人脸生成的实验。自编码器并非主流的生成模型，对于图像生成，现在更流行 Stable Diffusion 模型，相比之下 Stable Diffusion 的能力更强，生成的图像更逼真。Amazon 官方给出了 Stable Diffusion 相关实验：https://catalog.us-east-1.prod.workshops.aws/workshops/3b86fa89-da3a-4e5f-8e77-b45fb11adf4a/zh-CN?trk=cndc-detail 。我们可以使用 Amazon SageMaker 直接运行，详细的实践可以参考实验手册，下面给出跑出来的一些示例图片。 比如文本输入：拿着长剑的猫，可以得到下面的图片  或者：梵高星空，可以得到下面图片  每次生成的图片是不一样的，大家可以多尝试一些有趣的东西。 # 六、总结 ## 6.1 算法总结 自编码器的结构非常简单，训练起来相比 GAN 等要更容易，但是功能特别强大。在前面的实验中，我们都是把对结果的调整转为对编码的处理。在自编码器中有一个非常理想的设想，就是每个维度控制一个特征，但是实际上没有这么简单。自编码器还有一些其它变形，比如变分自编码器，ALAE 等，可以实现更加强大的功能，甚至可以媲美 StyleGAN，生成逼真的网络。 除了直接使用自编码器，在其它网络中也可以插入与自编码器类似的结构，比如 Stable Diffusion 中就有 Unet 网络，也是一种编码解码结构。 ## 6.2 Amazon SageMaker 总结 用于机器学习的平台有很多，Amazon SageMaker 是比较全面，且面向应用的一个。在 Amazon SageMaker 中包括了机器学习的各个流程，在以往 Python 中的开发习惯可以在 Amazon SageMaker 中完全适用，另外 Amazon SageMaker 提供了更高阶的 API，可以让用户更专注于算法的研究。 Amazon SageMaker 支持 Sklearn、PyTorch、TensorFlow、Hugging Face 等，对这些主流模块和框架都有相应的封装。另外 Amazon SageMaker 提供了非常便捷的部署方式。 为了方便训练模型，Amazon SageMaker 中提供了 Amazon AutoPilot 可以自动对各种模型以及各组超参数进行搜索，训练最优模型。 关于云上探索实验室的活动可以参考 https://dev.amazoncloud.cn/experience?trk=cndc-detail 里面有许多有趣的实验案例。通过云上探索实验室，开发者可以学习实践云上技术，同时将自己的技术心得分享给其他开发者小伙伴。一同创造分享，互助启发，玩转云上技术。云上探索实验室不仅是体验的空间，更是分享的平台”