图片风格转换应用程序:使用CoreML创建Prisma
2017年10月13日 由 xiaoshan.xiang 发表
528767
0
WWDC 2017让我们了解了苹果公司对机器学习的看法以及它在移动设备上的应用。CoreML框架使得将ML模型引入iOS应用程序变得非常容易。
大约一年前,我们在iOS和Android上实现了自己的神经网络推理机,现在我们很兴奋,因为苹果已经推出了类似的本机版本。在这篇文章中,我将向你展示如何使用只有开源模型和CoreML的方式实现一个小型的风格转换应用程序。
最近,我们在GitHub上共享了一个工具,这个工具将Torch7模型转换为CoreML。
CoreML:
https://github.com/prisma-ai/torch2coreml
作品地址:
http://cs.stanford.edu/people/jcjohns/eccv16/
repo地址:
https://github.com/jcjohnson/fast-neural-style
首先,把我们正在做的事情正则化。我们需要获取图像到图像(image-to-image)神经网络的CoreML(.mlmodel文件)模型,将一些“风格”应用于3个通道的图像输入并且得到输出图像。我们将利用其中一些模型作为示例来创建一个小型应用程序。
要求
需要安装macOS、Xcode 9、Torch7、PyTorch和torch2coreml。在描述每个组件安装过程的文本中还有一些链接。
步骤1:下载预先训练的模型
复制快速神经风格(fast-neural-style)的repo:
git clone https://github.com/jcjohnson/fast-neural-style
这个repo包含用于训练新模型的代码,并使用Torch7来运行推断(在图片上应用风格转换。需要正确地安装和配置Torch7。
Torch7地址:
http://torch.ch/docs/getting-started.html(本教程适用于没有CUDA的CPU版本)
然后下载预先训练的风格转换模型:
cd fast-neural-style && bash models/download_style_transfer_models.sh
下载的模型将出现在“model / eccv16”和“model / instance_norm”的子目录中。我们将使用“instance_norm”(带有实例正则化(InstanceNormalization)层的神经网络)模型,本教程中的所有内容都与“eccv16”模型兼容。你也可以训练自己的模型,该模型在快速神经样式(fast-neural-style)的repo中被描述,并用于进一步的实验。
步骤2:准备转换模型
如果我们深入研究快速神经风格(fast-neural-style)模型,我们会发现该模型存储在具有相应的预处理参数的Lua表格中,我们需要提取它。本教程的这一步包含了一些简单的Lua代码。
local model = torch.load(model_path).model
下一个问题是在模型中定制Torch7层(由Lua实现)。以下不是标准torch.nn程序包分布的一部分:
- ShaveImage;
- TotalVariation;
- 实例规范化(InstanceNormalization)。
看一下这些层的内部,并修改模型以准备转换。我们只对正向传递感兴趣,因此只对“updateOutput”方法进行了研究。
ShaveImage是一个简单的裁剪输入的层,它可以从左、右、顶部和图像底部裁剪一些像素。我们可以简单地把它改成标准torch的 SpatialZeroPadding层。
ShaveImage地址:
https://github.com/jcjohnson/fast-neural-style/blob/master/fast_neural_style/ShaveImage.lua
SpatialZeroPadding地址:
https://github.com/torch/nn/blob/master/SpatialZeroPadding.lua
如果我们看一下TotalVariation层,我们会看到它只是在前进阶段的一个无op层。因此,我们可以安全地从模型中删除这个层。
TotalVariation地址:
https://github.com/jcjohnson/fast-neural-style/blob/master/fast_neural_style/TotalVariation.lua
我们不能修改实例正则化(InstanceNormalization)层,因为它没有类似torch.nn的程序包,我们可以将它转换为相应的CoreML层,因为最新的coremltools python程序包支持该层。因此,我们将这个层保留在模型中。
完整的脚本准备模型:
https://github.com/prisma-ai/torch2coreml/blob/master/example/fast-neural-style/prepare_model.lua
步骤3:实现CoreML的转换器
在这一步骤中,我们将使用torch2coreml python程序包。我们编写一些python代码。我们从前面的步骤中得到了一些模型,这些模型是由Lua脚本编写的。我们将为快速神经样式(fast-neural-style )的torch模型实现一个python转换器。
必须安装PyTorch。安装步骤:
http://pytorch.org/。
还需要安装最新版本的torch2coreml python程序包,用于python 2.7(coremltools库只支持python版本作为撰写时间):
[sudo] pip install -U torch2coreml
torch2coreml库可以转换模型,该模型来自文件或已经加载的PyTorch模型(它在内部使用PyTorch来推断某些层的输入/输出形状并解析.t7文件)。我们需要使用PyTorch来加载Torch7模型,因为我们需要在python中使用存储在Lua模型中的参数来实现定制的实例正则化(InstanceNormalization)层。
from torch.utils.serialization import load_lua
model = load_lua(path, unknown_classes=True)
当模型中有实例规范化(InstanceNormalization)层时,我们需要打开“unknown_classes”引数。否则它会引发异常。
来自torch2coreml的“Convert”函数在PyTorch模型上运行推理,这就是为什么我们需要用它的工作模拟来替换实例规范化(InstanceNormalization)。在PyTorch中,我们找到InstanceNorm3d类,它和Justin Johnson在Lua的实现完全一样。因此,让我们使用这个类来实现PyTorch传统模块(我们只在PyTorch中使用Torch7等价代码)。
InstanceNorm3d地址:
https://github.com/pytorch/pytorch/blob/master/torch/nn/modules/instancenorm.py#L127
from torch.legacy.nn import Module
from torch.autograd import Variable
from torch.nn import InstanceNorm3d
class InstanceNormalization(Module):
def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=False):
super(Module, self).__init__()
if momentum is None:
momentum = 0.1
self._instance_norm = InstanceNorm3d(
num_features,
eps=eps,
momentum=momentum,
affine=True
)
@property
def eps(self):
return self._instance_norm.eps
@property
def weight(self):
return self._instance_norm.weight.data
@weight.setter
def weight(self, value):
self._instance_norm.weight.data = value
@property
def bias(self):
return self._instance_norm.bias.data
@bias.setter
def bias(self, value):
self._instance_norm.bias.data = value
def updateOutput(self, input):
return self._instance_norm.forward(Variable(input, volatile=True)).data
接下来,我们用新实现的类的实例替换模型中的每个未知的实例正则化(InstanceNormalization)对象。从这一点来看,我们有完全工作的PyTorch模型,它已经准备好被转换了。
我们没有实现在torch2coreml库中转换自定义的实例正则化(InstanceNormalization)层,但是幸运的是,它有一种机制,可以为未知的层添加自己的转换函数。
def convert_instance_norm(builder, name, layer, input_names, output_names):
if not isinstance(layer, InstanceNormalization):
raise TypeError('Unsupported type {}'.format(layer,))
epsilon = layer.eps
weight = layer.weight.numpy()
bias = None
if layer.bias is not None:
bias = layer.bias.numpy()
builder.add_batchnorm(
name=name,
channels=weight.shape[0],
gamma=weight,
beta=bias,
compute_mean_var=True,
instance_normalization=True,
input_name=input_names[0],
output_name=output_names[0],
epsilon=epsilon
)
return output_names
现在,将准备好的torch7模型转换成CoreML。使用完整的python转换器脚本。它有一些额外的修改,例如,修复空间全卷积层(SpatialFullConvolution)的加载。
python转换器脚本:
https://github.com/prisma-ai/torch2coreml/blob/master/example/fast-neural-style/convert-fast-neural-style.py
.mlmodel文件可以嵌入到iOS应用程序内部,运行生成的python CoreML模型来测试是否在图像上应用了风格转换。为此,我创建了一个简单的脚本:
import argparse
from PIL import Image
from coremltools.models import MLModel
def main():
parser = argparse.ArgumentParser(
description='Stylize image using CoreML'
)
parser.add_argument('-input', required=True, help='Path to input image')
parser.add_argument('-output', required=True, help='Output path')
parser.add_argument('-model', required=True, help='CoreML model path')
args = parser.parse_args()
image = Image.open(args.input)
net = MLModel(args.model)
stylized_image = net.predict({'inputImage': image})['outputImage']
stylized_image = stylized_image.convert('RGB')
stylized_image.save(args.output)
if __name__ == "__main__":
main()
请注意使用输入大小与模型兼容的图像,或者你可以在调用MLModel的“predict”方法之前简单地添加图像扩展。
马赛克模型的输出示例:
步骤4:iOS应用程序
利用上一步使用的4个风格转换模型实现了一个简单的应用程序。由于它超出了当前文章的范围,你可以在Apple教程和文档中找到使用CoreML的详细解释。它非常简单直观。你可以找到完整的iOS应用程序源代码(本教程的第一个截图是这个应用程序)。
程序源代码地址:
https://github.com/prisma-ai/torch2coreml/tree/master/example/fast-neural-style/ios
结论
我们使用torch2coreml软件包将原始Justin Johnson的快速神经风格(fast-neural-style)模型转换为CoreML。 获得的模型可以在iOS和macOS应用程序中使用。你可以将torch2coreml包用于风格转换和其他模型。
本教程的完整的源代码可以在这里找到:
https://github.com/prisma- ai/torch2coreml/tree/master/example/fast-neural-style。