实时视频神经风格转移
这篇文章详细介绍在一个受限的移动环境中创建的一个项目,该项目将艺术风格应用于实时视频。我们将描述风格转移技术、使用的机器学习模型架构、在受约束的硬件上实时执行风格转移的挑战和限制、为优化这些约束的风格转移系统而采取的措施,以及创建应用CoreML和Swift来包装机器学习模型的iOS应用程序。
首先,在背景部分,我们将简要地概述一下什么是神经风格转移,实时地执行视频风格转移的挑战和约束,并概述我们在处理这个项目中的总体方法。接下来,在机器学习模型中,我们将讲述我们训练的神经网络架构和训练过程。在机器学习模型压缩中,我们将会透露一些用于模型压缩的方法,比如迁移学习和知识蒸馏,以及这些方法如何提高我们的网络效率的结果。接下来,在移动环境中的风格转换中,我们将描述将神经网络转换成CoreML格式并将其包装在iOS应用程序中的过程。最后,在结果与接下来的步骤中,我们将提供一些我们的视频风格转移的示例。
背景
神经风格转移是用于将一个图像的艺术风格与另一图像的内容结合起来的机器学习技术。神经风格转移最早由Gatys等人在2015年的A Neural Algorithm of Artistic Style中描述。其基本思想是利用经过预处理的深度卷积神经网络(通常是经过图像分类或目标检测训练的神经网络)学习到的特征表示来获取任何图像的风格和内容的单独表示。一旦找到这些表示,我们就可以尝试优化生成的图像,以重新组合不同目标图像的内容和风格。
蒙娜丽莎Pablo Picasso立体主义风格的风格转移示例
从2015年以来,在风格转移上取得了很大进展,以便更快地进行风格转移训练和推理,并将风格转移技术从静止图像扩展到视频甚至音频和其它媒体。在视频上进行风格转换有多种方法,最简单的方法是在视频的每一个连续帧上重复循环静止图像风格转换过程。该方法的其他改进方法是对输出视频帧与帧之间的平滑度进行优化。我们选择了最简单的视频风格转换方法,在输入视频的每一帧上连续进行图像到图像的风格转换。
除了要使风格转换看起来很好之外,实时风格化视频还对神经网络推理的速度定义了额外的约束,因为我们努力最小化延迟和最大化输出视频的帧速率,使其看起来平滑一致。在iOS应用程序中包装风格转移流程产生了另外的挑战,因为网络必须足够小并且足够快,以使用iPhone的受限硬件来存储和处理。
我们采取的方法是:将一个预先训练好的图像风格传递神经网络,通过传递学习和知识提炼对其进行压缩和优化,训练不同的风格,将经过压缩和训练的模型转换为CoreML,最后创建一个iOS应用程序,从手机摄像头捕捉视频,对每一帧进行网络推理,并将程式化的结果显示在手机屏幕上。
机器学习模型
架构
我们使用的风格转移模型由Johnson 等人定义。该模型是用于图像神经网络的深度卷积图像,具有三个卷积层,五个残差块和三个反卷积层。该模型的实现可以在PyTorch库中找到这里(https://github.com/jsigee87/real-time-style-transfer/tree/master/webcam-app)。
训练
模型架构和训练过程概述
由于风格转移的目标是优化输出图像以包含类似于指定的“风格目标”图像(例如pablo picasso立体主义风格)的风格,并包含与输入图像类似的内容,训练损失在两个单独的部分中计算:内容损失和风格损失,它们都是从预训练用于图像分类的卷积神经网络中提取的。
内容丢失表示风格转移网络输出图像的内容与输入图像的内容有多相似 。由于图像分类卷积神经网络被迫学习高级特征/抽象表示或图像的“内容”以便精确地对它们进行分类,因此可以容易地从预训练图像分类网络的层激活中获取图像内容的表示。在该实现中使用的预训练网络是VGG-16。在图像分类网络中,较深层倾向于学习越来越高级别的特征,因此内容损失可以定义为层输出之间的欧几里德距离(当y ^和yc通过VGG-16时),来自VGG-16的j。可以根据需要选择j(从中提取内容特征的层),并且所选择的层越深,输出图像的外观就越抽象。
类似地,尽管稍微复杂一点,“风格目标”(y_s)和网络输出y ^之间的风格损失也被计算为从VGG-16的层输出中提取的特征之间的距离。这里的主要区别在于,我们不是直接从VGG-16的层激活中获取特征表示,而是首先将这些特征表示转换为称为Gram矩阵的东西。该Gram矩阵是矢量的所有可能内积的矩阵。在没有进入太多数学的细节的情况下,基本上Gram矩阵所做的是从VGG的层激活中捕获所有信息,但它会删除与信息在图像中空间表示的位置有关的所有信息。Gram Matrix表示图像的一般内容,无需担心图像中某些对象/元素的位置。换句话说,Gram矩阵表示图像的风格。因此,风格损失就是VGG层i激活的Gram矩阵之间的(Frobenius)距离,当y^从y_s通过VGG层i激活的Gram矩阵是。风格损失实际上被计算为VGG中多个层i的风格损失的总和,如上图所示。
训练网络以同时最小化内容损失和风格损失。对于每种艺术风格,必须训练单独的网络。训练的输入/内容图像y_c可以是任意图像,使得容易获得充足的训练数据。在训练之后,模型可以通过以图像/视频帧作为输入执行训练网络的forward pass来使图像(以及因此视频帧)风格化。但是,网络规模相当大,因此feedforward可能需要花费大量时间来设置新图像的风格。
模型压缩
由于我们希望我们的风格转移实时发生,我们希望将风格转移所花费的时间尽可能地低。由于每个帧必须通过训练过的传输网络ed forward,所以我们最感兴趣的是使我们的模型尽可能小,同时仍然保留在训练过程中学习到的大部分函数f_w。虽然神经网络经常受益于在训练过程中使用可训练参数来学习近似函数,但实际的计算过程并不需要所有的网络参数。有几种常用的方法可以将经过训练的网络缩小到较小的规模。我们重点研究的两种方法是转移学习和知识精馏。为了确定哪些较小的体系结构可以最大限度地减少推理时间,我们还使用不同的较小的模型体系结构进行了一系列的实验,这些模型体系结构使用虚拟权重初始化,并对feed forward passes进行了计时。
迁移学习
迁移学习是机器学习研究的一个一般领域,用于将学到的知识用于一个类似但不同的任务。通常,这是通过从为第一个任务训练的网络中获取权重来完成的,并使用它们初始化用于学习新任务的网络。这里的基本思想是,网络可以从第一个任务中学到的一些一般知识中获益,而不必从头开始学习所有的权重。
在我们的例子中,学习的任务实际上是相同的,但网络规模是不同的。我们在这里可以做的事情是有限的,因为为了能够用原始网络的权重初始化我们较小的网络权重,每层的权重必须是相同的形状。这限制了我们仅能够一次移除整个层/块,而且,仅能够移除具有相同输入和输出尺寸(并因此具有重量形状)的网络部分。在我们的例子中,符合此标准的网络的唯一部分是中间的残差块。由于原始网络中有五个残差块,我们删除了中间三个残差块,并将网络其余部分的权重初始化为原始网络的权重。这样做的目的是为了学习与以前相同的函数,所以网络只需要学习从一个基本相似的初始函数中补偿丢失的残差块。因此,具有初始权值的网络在训练中收敛速度比随机初始化的网络快许多数量级。
知识蒸馏
模型压缩的另一种方法是知识蒸馏,其中训练较小的“student”网络以模拟每个训练示例的较大的预训练“teacher”网络的输出。由于较大的神经网络已经在学习函数的高级特征方面做了大量工作,因此较小的网络可以从其知识中受益,而无需经历几乎同样多的训练示例。我们采用开源知识蒸馏代码进行图像分类任务,为我们的风格转移任务进行知识蒸馏。
我们的任务的损失函数既衡量student和teacher模型输出之间的差异,也包括原始训练代码(在上面的训练模型部分中描述)的原始内容和风格损失,以训练较小的网络。
推理时间实验
为了更好地了解模型压缩如何影响我们神经网络的推理时间,我们构建了较小的模型体系结构变量,随机初始化权重,只是为了查看哪些体系结构具有最小的前馈推理时间。尽管这些随机初始化的“虚拟模型”的输出是随机的,但是看看推理时间如何随参数体系结构的变化而变化是有用的,因此我们可以知道哪些体系结构可能是最值得训练的。我们的实验结果如下所示。每个实验进行100次试验。
不同架构修改的推理时间
所有实验均在运行Ubuntu Linux 16.04 LTS的计算机上进行,具有以下规范:
- 处理器:英特尔®酷睿™i7-6500U CPU @ 2.50GHz×4
- 内存:8 G
移动环境中的风格转移
转换为CoreML
为了在通用计算机上训练深度神经网络并将它们部署在iOS设备上以执行实时风格化,我们必须找到一种方法将我们的神经网络从其原始PyTorch格式转换为本机Apple CoreML格式,以管理深度学习网络。作为执行此转换的中间件,我们使用Open Neural Network Exchange(ONNX),这是一种开发通用深度学习网络存储格式的开源工具。目前,他们提供转换工具和示例,用于获取Pytorch模型,将它们移植到ONNX,并将它们转换为CoreML模型格式。通过将PyTorch的这个管道设置为ONNX到CoreML,我们能够迅速将所有经过训练的PyTorch模型转换为Apple设备可用的格式。
实时捕获和推理
应用程序使用的摄像头是iPhone X和iPhone 6s上的后置摄像头。一个直接的折衷是选择一个足够清晰的分辨率来区分捕获图像的内容,而不是选择一个足够低的分辨率来使用CoreML模型进行更快的推理。对于无符号32位RGB(3通道),我们选择640像素×480像素的分辨率作为图像分辨率和推理时间之间的良好平衡。帧以每秒30帧(FPS)捕获。
AV Capture是一个围绕低级设备驱动程序的Swift包装器,用于连接Apple设备的视听输入。为了从相机实时获取图像帧并将它们带入内存,我们建立了一个捕获会话,以配置具有所需分辨率的后置iPhone相机,并设置一个异步先入先出缓冲区来保存传入帧,直到我们的程序准备好接收帧为止。当程序准备就绪时,它会将RGB图像帧拉出缓冲区,检查RGB帧是否具有所需的通道,宽度和高度(3,640,480),并通过CoreML模型运行它以获得程式化的形象。最后,将风格化图像异步调度到主线程以在显示器上呈现。
结果
总之,我们使用模型压缩技术(如转移学习和知识蒸馏)训练视频风格转移模型,创建iOS应用程序,并显著缩短机器学习模型推理时间。以下是我们的应用程序中的视频的一些示例。
Undie style
Undie风格
Tron style
Green swirly style
可以在GitHub存储库中找到此项目的模型和应用程序代码(https://github.com/jsigee87/real-time-style-transfer/tree/master/webcam-app)。