0°
PyTorch入门实战教程

在 iOS 上用 PyTorch 和 CoreML 实现图像风格迁移 (Neural Style Transfer)

2017 年 6 月,Apple 推出了 CoreML,这是一个旨在将机器学习模型集成到 iOS 应用程序中的框架,为开发人员提供了大量的可能性,包括图像分析到 NLP(自然语言处理),决策树学习等等。

在 2018 年春季,我开始深入研究 CoreML,并学习了神经风格转移。 这篇文章,我将讲解一下大致的实现过程。

训练一个图像风格迁移模型

我们使用 PyTorch 官方提供的 Fast Neural Style 代码(地址是 https://github.com/pytorch/examples/tree/master/fast_neural_style)来训练模型,我们使用默认参数,并且每 2000 次循环保存一个 checkpoint:

将模型导出到 CoreML

PyTorch 本身是不支持直接导出 CoreML 模型的(至少目前是这样),但我们可以借助 ONNX(Open Neural Network Exchange)的帮助,来先转换成 .onnx 格式,然后再转换成 CoreML。

PyTorch → ONNX→ CoreML

文件 neural_style.py 中已经有输出位 ONNX 格式的代码了,我们只需要完成第二步就行。我们首先定义一个转换器:

接下来我们进行转换:

当然,如果你不想自己训练,可以在这里下载训练好的模型。

在 iOS 上使用 CoreML NST 模型

我们创建一个 UIImageView 和 UIButton,添加一个 UIBarButtonItem 来从设备的照片库中选择一个图像并将其设置到 UIImageView 上。

等 UI 完成后,我们就可以把训练的模型导入进来了,只需要将模型文件拖到文件视图中就行了。点击模型文件,可以看到一些模型信息:

这是一个新的 Model 类会被自动创建,查看文件可知我们将要用到的函数是  prediction(inputImage: CVPixelBuffer).

我们选定一张图片后:

  • 首先保存图像的原始大小
  • 把输入 UIImage 转换成一个 720 x 720 的 CVPixelBuffer
  • 输入到模型
  • 把模型的输出 CVPixelBuffer 转换回 UIImage
  • 重新调整图像大小到原始大小

下面是一个 Swift 4 实现:

下面是最终效果:

本文的项目源代码已经放到 Github 上。

PyTorch入门实战教程
0 条回复 A 作者 M 管理员
    所有的伟大,都源于一个勇敢的开始!
欢迎您,新朋友,感谢参与互动!欢迎您 {{author}},您在本站有{{commentsCount}}条评论