1. 程式人生 > 其它 >【深度學習】圖片風格轉換應用程式:使用CoreML建立Prisma

【深度學習】圖片風格轉換應用程式:使用CoreML建立Prisma

WWDC 2017讓我們瞭解了蘋果公司對機器學習的看法以及它在移動裝置上的應用。CoreML框架使得將ML模型引入iOS應用程式變得非常容易。

大約一年前,我們在iOS和Android上實現了自己的神經網路推理機,現在我們很興奮,因為蘋果已經推出了類似的本機版本。在這篇文章中,我將向你展示如何使用只有開源模型和CoreML的方式實現一個小型的風格轉換應用程式。

最近,我們在GitHub上共享了一個工具,這個工具將Torch7模型轉換為CoreML。

GitHub地址:https://github.com/prisma-ai/torch2coreml

我將用它來轉換Justin Johnson預先訓練的模型,該模型來自於他的作品“實時樣式傳輸和超解析度的知覺損失”。

作品地址: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層。

如果我們看一下TotalVariation層,我們會看到它只是在前進階段的一個無op層。因此,我們可以安全地從模型中刪除這個層。

ShaveImage層:https://github.com/jcjohnson/fast-neural-style/blob/master/ fast_neural_style/ShaveImage.lua SpatialZeroPadding層:https://github.com/torch/nn/blob/master/SpatialZero Padding.lua 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.serializationimport 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.nnimport Module
from torch.autogradimport Variable
from torch.nnimport InstanceNorm3d

class InstanceNormalization(Module):
    def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=False):
        super(Module,self).__init__()
        if momentumis 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

接下來,我們用新實現的類的例項替換模型中的每個未知的例項規範化(instancenalization)物件。從這一點來看,我們有完全工作的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.biasis 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轉換器指令碼。

指令碼地址:https://github.com/prisma-ai/torch2coreml/blob/master/example/fast-neural-style/convert-fast-neural-style.py

它有一些額外的修改,例如,修復空間全卷積層(SpatialFullConvolution)的載入。

.mlmodel檔案可以嵌入到iOS應用程式內部,執行生成的python CoreML模型來測試是否在影象上應用了風格轉換。為此,我建立了一個簡單的指令碼:

import argparse

from PILimport Image
from coremltools.modelsimport 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