MLMultiArrayとポインタ

はじめに

[CoreML, Pytorch, Swift] PytorchのモデルをCoreMLを使ってiOSで実行
こちらの記事を参考にして、
githubで公開されているPyTorch学習モデルのiOS移植に取り組みました。(画像解析系)
動作確認をしてみると、動作に結構時間がかかることがわかりました。
モデルの推論に時間がかかるのは承知の上で、他にどの処理で時間がかかっているか確認していたところ、MLMuitiArrayへ画素値をコピーする処理に時間を要していました。

そこで、今回はMLMultiArrayへのコピー処理の時間短縮について
備忘録として書き残しておきます。

MLMultiArrayとは

https://developer.apple.com/documentation/coreml/mlmultiarray
（グーグル翻訳ありがとう。。。！）

A multidimensional array used as a feature input or feature output for a model.
モデルの機能入力または機能出力として使用される多次元配列。

そのままですね、coreMLの入力また出力として用いられる配列です。

普通に実装したとき

参考記事のUIImageからグレースケールの値を算出する箇所をコピーして、
RGB値を取得できるようソースを書き換えました。

#uiImage: UIImageの画像
let imagePixel = uiImage.getPixelRgb()

extension UIImage {
    //pixelBUfferに変換(RGB値)
    func getPixelRgb() -> [[[Double]]]
    {
        guard let cgImage = self.cgImage else {
            return []
        }
        let bytesPerRow = cgImage.bytesPerRow
        let width = cgImage.width
        let height = cgImage.height
        let bytesPerPixel = 4
        let pixelData = cgImage.dataProvider!.data! as Data

        var r_buf : [Double] = []
        var g_buf : [Double] = []
        var b_buf : [Double] = []

        var r_arr=[[Double]]()
        var g_arr=[[Double]]()
        var b_arr=[[Double]]()

        var arr=[[[Double]]]()

        for j in 0..<height {
            for i in 0..<width {
                let pixelInfo = bytesPerRow * j + i * bytesPerPixel
                let r = Double(pixelData[pixelInfo])
                let g = Double(pixelData[pixelInfo+1])
                let b = Double(pixelData[pixelInfo+2])
                r_buf.append(Double(r/255.0))
                g_buf.append(Double(g/255.0))
                b_buf.append(Double(b/255.0))
            }
            r_arr.append(r_buf)
            g_arr.append(g_buf)
            b_arr.append(b_buf)
            r_buf = []
            g_buf = []
            b_buf = []
        }

        arr.append(b_arr)
        arr.append(g_arr)
        arr.append(r_arr)

        return arr
    }
}

CGImageの時点では、配列が[横]×[縦]×[RGB]となっているので、
目的に合わせて転置を行います。
私が移植したモデルの入力は[BGR]×[横]×[縦](元のソースではopenCVを用いたため)だったので
ここでは上のような配列に変換しています。

MLMultiArrayへの変換処理は、参考にした記事やApple Developer Forumsでは、下記のように紹介されております。

//(3, 256, 192)のMLMultiArrayを生成（メモリ確保）
let mlArray = try! MLMultiArray(shape: [3,  256, 192], dataType: MLMultiArrayDataType.double)

//(3, 256, 192)のMLMultiArrayに画素データを代入
for i in 0..<3 {
    for j in 0..<256{
        for k in 0..<192{
            mlArray[192*256*i + 192*j + k] = imagePixel[i][j][k] as NSNumber
        }
    }
}

はじめに、画素配列と同じ大きさの空のMLMultiArrayを用意して、このMLMultiArrayにそのまま代入を行います。
とてもシンプルでわかりやすいと思います。わかりやすいのはいいのですが、この処理に0.1秒くらいかかってしまっております。
バックカメラを使ったリアルタイム解析だとリアル解析など言ってられないくらいのストレスを感じると思います。

ポインタを使ってみる

MLMuitiArrayのリファレンスにポインタが使える旨が明記されておりました。
これを用いて画素配列をポインタとして参照させます。

その前に画素配列を一次元にします。

#uiImage: UIImageの画像
let imagePixel = uiImage.getPixelRgb()

extension UIImage {
    //pixelBUfferに変換(RGB値)
    func getPixelRgb() -> [Double]
    {
        guard let cgImage = self.cgImage else {
            return []
        }
        let bytesPerRow = cgImage.bytesPerRow
        let width = cgImage.width
        let height = cgImage.height
        let bytesPerPixel = 4
        let pixelData = cgImage.dataProvider!.data! as Data

        var r_buf : [Double] = []
        var g_buf : [Double] = []
        var b_buf : [Double] = []

        for j in 0..<height {
            for i in 0..<width {
                let pixelInfo = bytesPerRow * j + i * bytesPerPixel
                let r = Double(pixelData[pixelInfo])
                let g = Double(pixelData[pixelInfo+1])
                let b = Double(pixelData[pixelInfo+2])
                r_buf.append(Double(r/255.0))
                g_buf.append(Double(g/255.0))
                b_buf.append(Double(b/255.0))
            }
        }
        return ((b_buf + g_buf) + r_buf)
    }

次にこの配列のポインタを取得します。

#画素配列のポインタ
let imagePointer : UnsafePointer<Double> = UnsafePointer(imagePixel)

下準備は完了です。

メモリに参照させる処理ですが、
上での実装と同じように、空のmlMultiArrayを用意します。
そしてMLMultiArrayの関数"dataPointer"で入力先のポインタを指定し、
"initializeMemory"という関数を用いて画素配列の値を参照させます。

//(3, 256, 192)のMLMultiArrayを生成（メモリ確保）
let mlArray = try! MLMultiArray(shape: [3,  256, 192], dataType: MLMultiArrayDataType.double)

#as: 何型として代入するか
#from: どこのポインタから
#count: 何個分参照するのか
mlArray.dataPointer.initializeMemory(as: Double.self, from: imagePointer, count: imagePixel.count)

まとめ

速度的に言えば1000倍近く早くなっていました。
他の処理でもポインタを用いればもっと早くなる部分はたくさんありそうな一方、
メモリ解放等のタイミング等swiftでのポインタ関係の理解が不十分なため、
少し諸刃の剣な感じがあります。
きちんと使う場合には、それなりに学習が必要がありそうです。

今回はこのような実装を行っておりましたが実装する際、
下調べが足りておらずVisionという画像処理フレームワークの存在を知りませんでした。
もしかすると、上記の実装よりVisionを用いた方が早いかもしれませんw
今後比較してみたいと思います。

問題点等ありましたらコメントでご指摘いただけると嬉しいです。

参考文献

https://qiita.com/Mco7777/items/09976fe31e3a4a6f408c
ポインタを用いた実装に関しては、
実装時参考する記事が少ないと思っていましたが、
この記事の執筆時にたくさん記事を見つけてしまい頭を抱えました。。。
今後はきちんと調査を行います。。。