(八)构建一个Docker容器来训练Deep Fake Autoencoders

最新推荐文章于 2024-05-08 18:41:38 发布

寒冰屋

最新推荐文章于 2024-05-08 18:41:38 发布

阅读量466

点赞数

分类专栏：人工智能 python 文章标签： docker 深度学习 tensorflow

原文链接：https://www.codeproject.com/Articles/5298029/Building-a-Docker-Container-to-Train-Deep-Fake-Aut

版权

人工智能同时被 2 个专栏收录

564 篇文章 57 订阅

订阅专栏

python

261 篇文章 9 订阅

订阅专栏

构建Docker镜像并将其推送到Google AI 平台

将训练作业提交并监控到Google AI Platform

在GCS存储桶中查找和下载经过训练的模型

下载项目文件 - 75.5 MB

深度伪造——使用深度学习在视频中将一个人的脸换成另一个人的脸——是当今使用人工智能的最有趣和最可怕的方式之一。

虽然深度伪造可用于合法目的，但它们也可用于虚假信息。能够轻松地将某人的脸换成任何视频，我们真的可以相信我们的眼睛告诉我们的吗？政治家或演员做或说令人震惊的事情的真实视频可能根本不是真实的。

在本系列文章中，我们将展示深度伪造的工作原理，并展示如何从头开始实现它们。然后我们将看看DeepFaceLab，它是一种多合一的Tensorflow 驱动的工具，通常用于创建令人信服的深度伪造。

在之前的文章中，我们讨论了深度伪造的自动编码器以及使它们工作所需的条件。如果你刚刚加入我，我最近解释了先决条件，然后跳到创建容器和提交到谷歌云平台的代码，所以请不要跳过它，否则这里列出的步骤可能无法按预期工作。在本文中，我将向您展示如何创建、测试和提交引用的容器。让我们开始吧！

我们的Docker容器的结构

在我们的例子中，我们的Docker容器并不复杂，只包含几个文件：

/deepfakes
|
| -- Dockerfile
| -- config.yaml
| -- model.py
| -- task.py
| -- data_utils.py

Dockerfile规定了使用正确镜像的必要指南，安装我们训练顺利运行的先决条件，将文件复制到Docker镜像，并定义此容器的入口点。你可以在这里找到AI平台的Dockerfile基础知识。
config.yaml包含定义训练层所需的信息。它本质上是一个文件，说明训练模型需要多少计算资源。您可以在此处找到有关如何构建、理解和定义您需要使用的层的步骤。如果您计划实施基于GPU的训练，您可以在此处找到定价。
task.py处理与模型训练工作流相关的一切。它调用所需的函数来获取通过CLI传递的参数、解压缩、预处理和拆分从GitHub 存储库获取的数据。它设置GPU、构建模型、开始训练并将生成的自动编码器保存在Google Storage中。
model.py只包含一个构建我们的模型并在编译后返回它的函数。
data_utils.py包含数据加载、数据集创建和模型保存所需的功能。

让我们进入实际代码！但首先，在您的机器中创建一个包含所有这些的文件夹。

编码Dockerfile

正如我在Docker容器概述文章中（第六篇）提到的，这个文件说明了构建我们的容器所需的一切。这是使此容器在Google AI Platform中可执行所需的Dockerfile：

FROM gcr.io/deeplearning-platform-release/tf2-gpu.2-3
WORKDIR /root
 
RUN pip install pandas numpy google-cloud-storage scikit-learn opencv-python
RUN apt-get update; apt-get install git -y; apt-get install -y libgl1-mesa-dev
RUN git clone https://github.com/sergiovirahonda/DeepFakesDataset.git
 
COPY model.py ./model.py
COPY data_utils.py ./data_utils.py
COPY task.py ./task.py
 
ENTRYPOINT ["python","task.py"]

第一行表明我们将用基于 Ubuntu 的TensorFlow 2 GPU支持的镜像。第二行定义哪个将成为工作目录。第三行和第四行安装所需的依赖项。第5行克隆包含原始数据集的repo，从第6行到第8行，我们复制模型训练所需的.py文件。最后，我们指出训练过程从task.py文件开始，并且必须用Python解释。将这些命令复制并粘贴到您的Dockerfile中。

定义config.yaml文件

该文件将取决于您的需要。如果你想加速你的训练过程，它可能会变得更复杂，但对于我们的例子，文件是这样的：

trainingInput:
    scaleTier: CUSTOM
masterType: standard_gpu

我使用的是单个标准GPU机器，但您可以根据需要添加更多机器并实现分布式训练。

编写task.py文件

正如我之前提到的，这个文件将处理与训练工作相关的所有事情。它是容器入口并调用所有外围文件和函数：

import tensorflow as tf
from tensorflow.keras.callbacks import ModelCheckpoint
import argparse
import data_utils
import model
import sys
import os
from google.cloud import storage
import datetime
 
def train_model(args):
X_train_a, X_test_a, X_train_b, X_test_b = data_utils.load_data(args)
 
    print("Num GPUs Available: ", len(tf.config.experimental.list_physical_devices('GPU')))
    tf.config.set_soft_device_placement(True)
    tf.debugging.set_log_device_placement(True)
 
    # If you're going to use MultiGPU - Create a MirroredStrategy
    # strategy = tf.distribute.MirroredStrategy()
    # print("Number of devices: {}".format(strategy.num_replicas_in_sync))
 
    # with strategy.scope():
    #     autoencoder_a = model.autoencoder_model()
    #     autoencoder_b = model.autoencoder_model()
 
    # If you're going to use MultiGPU - Comment out the next two lines
    autoencoder_a = model.autoencoder_model()
    autoencoder_b = model.autoencoder_model()
 
    print('Starting autoencoder_a training', flush=True)
    checkpoint_1 = ModelCheckpoint("autoencoder_a.hdf5", monitor='val_loss', verbose=1,save_best_only=True, mode='auto', save_freq="epoch")
    autoencoder_a.fit(X_train_a, X_train_a,epochs=args.epochs,batch_size=128,shuffle=True,validation_data=(X_test_a, X_test_a),callbacks=[checkpoint_1])
    print('Training A has ended. ',flush=True)
 
    print('Starting autoencoder_b training', flush=True)
    checkpoint_2 = ModelCheckpoint("autoencoder_b.hdf5", monitor='val_loss', verbose=1,save_best_only=True, mode='auto', save_freq='epoch')
    autoencoder_b.fit(X_train_a, X_train_a,epochs=args.epochs,batch_size=128,shuffle=True,validation_data=(X_test_b, X_test_b),callbacks=[checkpoint_2])
    print('Training B has ended. ',flush=True)
 
    print('Proceeding to save models into GCP.',flush=True)
 
    data_utils.save_model(args.bucket_name,'autoencoder_a.hdf5')
    print('autoencoder_a saved at GCP Storage', flush=True)
    data_utils.save_model(args.bucket_name,'autoencoder_b.hdf5')
    print('autoencoder_b saved at GCP Storage', flush=True)
 
 
def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--bucket-name',
                        type=str,
                        help='GCP bucket name')
    parser.add_argument('--epochs',
                        type=int,
                        default=10,
                        help='Epochs number')
    args = parser.parse_args()
    return args
 
def main():
 
    args = get_args()
    train_model(args)
 
if __name__ == '__main__':
main()

请记住，如果您决定增加训练机器的数量并使用多GPU，您需要为此采取一种策略。一个非常常见的方法是MirroredStrategy。

编码model.py文件

我之前已经完成了几篇文章的模型开发，这个文件基本上回收了构建我们的自动编码器所需的代码。它只包含一个返回我们已经编译的模型的函数：

from tensorflow.keras.models import load_model
from tensorflow.keras import Model
from tensorflow.keras import layers
from tensorflow.keras import optimizers
import tensorflow as tf
import numpy as np
 
 
def autoencoder_model():
    #Making encoder:
 
    input_img = layers.Input(shape=(120, 120, 3))
    x = layers.Conv2D(256,kernel_size=5, strides=2, padding='same',activation='relu')(input_img)
    x = layers.MaxPooling2D((2, 2), padding='same')(x)
    x = layers.Conv2D(512,kernel_size=5, strides=2, padding='same',activation='relu')(x)
    x = layers.MaxPooling2D((2, 2), padding='same')(x)
    x = layers.Conv2D(1024,kernel_size=5, strides=2, padding='same',activation='relu')(x)
    x = layers.MaxPooling2D((2, 2), padding='same')(x)
    x = layers.Flatten()(x)
    x = layers.Dense(9216)(x)
    encoded = layers.Reshape((3,3,1024))(x)
 
    encoder = Model(input_img, encoded,name="encoder")
    encoder.summary()
 
    #Making decoder:
    decoder_input= layers.Input(shape=((3,3,1024)))
    x = layers.Conv2D(1024,kernel_size=5, strides=2, padding='same',activation='relu')(decoder_input)
    x = layers.UpSampling2D((2, 2))(x)
    x = layers.Conv2D(512,kernel_size=5, strides=2, padding='same',activation='relu')(x)
    x = layers.UpSampling2D((2, 2))(x)
    x = layers.Conv2D(256,kernel_size=5, strides=2, padding='same',activation='relu')(x)
    x = layers.Flatten()(x)
    x = layers.Dense(np.prod((120, 120, 3)))(x)
    decoded = layers.Reshape((120, 120, 3))(x)
 
    decoder = Model(decoder_input, decoded,name="decoder")
    decoder.summary()
 
    #Making autoencoder
    auto_input = layers.Input(shape=(120,120,3))
    encoded = encoder(auto_input)
    decoded = decoder(encoded)
 
    autoencoder = Model(auto_input, decoded,name="autoencoder")
    autoencoder.compile(optimizer=optimizers.Adam(lr=5e-5, beta_1=0.5, beta_2=0.999), loss='mae')
    #autoencoder.compile(optimizer=optimizers.Adam(lr=5e-5, beta_1=0.5, beta_2=0.999), loss='mse')
    autoencoder.summary()
 
    return autoencoder

编码我们的data_utils.py文件

我们的最后一个文件由task.py调用以预处理从GitHub 存储库克隆的数据并将其拆分返回，同时还将模型保存到我们的Google Cloud Storage存储桶中：

import datetime
from google.cloud import storage
import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn import preprocessing
import os
import zipfile
import cv2
import sys
 
def load_data(args):
    sys.stdout.write("Starting data processing")
    print('Current directory:', flush=True)
    print(os.getcwdb(), flush=True)
    entries = os.listdir('/root/DeepFakesDataset')
    print(entries)
    print('Extracting files', flush=True)
 
    file_1 = '/root/DeepFakesDataset/trump.zip'
    file_2 = '/root/DeepFakesDataset/biden.zip'
    extract_to = '/root/DeepFakesDataset/'
 
    with zipfile.ZipFile(file_1, 'r') as zip_ref:
        zip_ref.extractall(extract_to)
    
    with zipfile.ZipFile(file_2, 'r') as zip_ref:
        zip_ref.extractall(extract_to)
 
    print('Files extracted', flush=True)
 
    faces_1 = create_dataset('/root/DeepFakesDataset/trump')
    print('First dataset created', flush=True)
    faces_2 = create_dataset('/root/DeepFakesDataset/biden')
    print('Second dataset created', flush=True)
 
    print("Total President Trump face's samples: ",len(faces_1))
    print("Total President Biden face's samples: ",len(faces_2))
 
    X_train_a, X_test_a, y_train_a, y_test_a = train_test_split(faces_1, faces_1, test_size=0.20, random_state=0)
    X_train_b, X_test_b, y_train_b, y_test_b = train_test_split(faces_2, faces_2, test_size=0.15, random_state=0)
 
    print(X_train_a.shape)
    print(X_train_a[0])
 
    return X_train_a, X_test_a, X_train_b, X_test_b
 
def create_dataset(path):
    images = []
    for dirname, _, filenames in os.walk(path):
        for filename in filenames:
            image = cv2.imread(os.path.join(dirname, filename))
            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            image = image.astype('float32')
            image /= 255.0
            images.append(image)
    images = np.array(images)
    return images
 
def save_model(bucket_name, best_model):
    print('inside save_model, about to save it to GCP Storage',flush=True)
    bucket = storage.Client().bucket(bucket_name)
    print(bucket)
    blob1 = bucket.blob('{}/{}'.format(datetime.datetime.now().strftime('model_%Y%m%d_%H%M%S'),best_model))
blob1.upload_from_filename(best_model)

构建Docker镜像并将其推送到Google AI 平台

我们真的快要完成了！还有几个命令。保持良好的工作！

启动终端窗口，浏览到存储容器文件的文件夹并运行以下命令：

docker build -f Dockerfile -t $IMAGE_URI ./

此命令将需要几分钟才能完成。它将启动Docker镜像下载并按照Dockerfile中所述的所有说明进行操作。最后，它将在您的本地注册表中创建容器。

要将Docker镜像推送到GCP，请运行docker push $IMAGE_URI.。需要几分钟才能完成。本质上，它会将您刚刚构建的镜像推送到GCP注册表，以便您稍后可以使用它来启动训练作业，这是我们接下来要做的事情。一旦推送过程结束，就该最终提交作业以训练我们的模型了。

顺便说一句，如果您在将镜像推送到GCP时遇到问题，该资源可以帮助您了解发生了什么。

将训练作业提交并监控到Google AI Platform

要将训练作业提交到AI Platform，请发出以下命令，您唯一需要调整的是--epochs参数。将其设置为达到良好指标数字所需的迭代次数。请记住，epochs数越高，您的自动编码器就越好。按回车键后，它应该显示作业已正确提交。

gcloud ai-platform jobs submit training $JOB_NAME --region $REGION --master-image-uri $IMAGE_URI --config config.yaml -- --epochs=10000 --bucket-name=$BUCKET_NAME

要监控您的训练作业并了解那里发生了什么，请发出命令gcloud ai-platform jobs describe $JOB_NAME，这将自动返回可以查看训练日志的URL。另一种方法是使用培训工作页面并选择您提交的工作。请记住，该过程可能需要几个小时，如果您仅使用一个GPU在单个实例上训练模型，则可能需要更多时间。训练结束后，您应该可以在之前创建的存储桶中找到模型文件。