iOSのMPSCNNに渡すモデルパラメータのフォーマット / TensorFlowからの書き出し

  • 23
    いいね
  • 0
    コメント

iOS 10でMetal Performance Shadersフレームワークに、CNN(Convolutional Neural Network)演算機能群が追加されました。iOSデバイスのGPUで畳み込みニューラルネットワークの計算をさせることができる、という代物です。iOSデバイスのGPUも日々進歩しているとはいえ、モバイルデバイスなのでさすがに学習までは厳しいところですが、学習済みのネットワークの計算であれば使えるケースはありそうです。

先日書いた記事』では、MPSCNNを用いて手書き文字認識を行うCNNを実装するコードを示しました。

当該サンプル(AppleのMPSCNNHelloWorld)では、学習済みのモデルパラメータ(重み/バイアス)をxxxx.datというファイルでアプリに持たせて、実行時にロードしています。

この.datとはどういうファイルなのでしょうか?何かAppleの独自形式なのか、MPSCNNはその形式しか使えないのか、中身はどういうフォーマットになっているのか等々、本記事では、MPSCNNに渡せるモデルパラメータのフォーマットについて書いてみたいと思います。

ファイルフォーマットについて

Apple の MPSCNNHelloWorld や MetalImageRecognition といったサンプルに入っている xxxx.dat は、Apple独自とかMPSCNN専用とかそういうものではなく、普通のバイナリファイルです。

こんな感じで、MetalやMetalPerformanceShadersフレームワークに一切依存することなくメモリにロードすることができます。

let fd = open( filePath, O_RDONLY, S_IRUSR | S_IWUSR | S_IRGRP | S_IWGRP | S_IROTH | S_IWOTH)
assert(fd != -1, "Error: failed to open output file at \""+filePath+"\"  errno = \(errno)\n")

guard let hdr = mmap(nil, size, PROT_READ, MAP_FILE | MAP_SHARED, fd, 0) else {
    close(fd)
    return nil
}

UnsafeMutableRawPointer から UnsafePointer<Float> にキャストして、

let p = UnsafePointer<Float>(hdr.assumingMemoryBound(to: Float.self))
assert(p != UnsafePointer<Float>(bitPattern: -1), "mmap failed with errno = \(errno)")

このポインタを、畳み込み層を示す MPSCNNConvolution や全結合層を示す MPSCNNFullyConnected に渡して初期化します。

MPSCNNConvolution
public init(device: MTLDevice, convolutionDescriptor: MPSCNNConvolutionDescriptor, kernelWeights: UnsafePointer<Float>, biasTerms: UnsafePointer<Float>?, flags: MPSCNNConvolutionFlags)

MPSCNNConvolutionMPSCNNFullyConnected は初期化の際パラメータを内部にコピーするので、初期化し終わったらメモリにロードしたデータを開放し、ファイルをクローズします。

munmap(hdr, Int(size))
close(fd)

この処理もMetalやMetalPerformanceShadersフレームワークに依存しない、一般的なメモリやファイルに対する処理になります。

つまり、ファイルフォーマット自体はiOSで読み込みさえできれば何でもOKです。

実際、Apple の MPSCNNHelloWorld や MetalImageRecognition はTensorFlowで学習させたパラメータを書き出した、と説明に書いてありますし、また下記記事ではChainerやKelasで学習させたモデルを「HDF5」というファイル形式で保存してiOSで利用 1 したそうです。

重みデータの順序

ファイルフォーマットは何でもいい、と書きましたが、ファイルの中身についてはいくらか決まりがあります。

ここでもう一度、MPSCNNConvolutionの初期化メソッドを見てみます。

MPSCNNConvolution
public init(device: MTLDevice, convolutionDescriptor: MPSCNNConvolutionDescriptor, kernelWeights: UnsafePointer<Float>, biasTerms: UnsafePointer<Float>?, flags: MPSCNNConvolutionFlags)

CNNにおける「重み」パラメータは 入力チャネル数 × カーネル幅 × カーネル高さ × 出力チャネル数4次元テンソルになりますが、ファイルから読み込んだ重みパラメータについては、第4引数 kernelWeightsFloat型のポインタとして渡しているだけです。4次元テンソルの格納順序を指定するようなパラメータはありません。

というわけで、「重み」の4次元テンソルの格納順は、次のように決められています。

weight[ outputChannels ][ kernelHeight ][ kernelWidth ][ inputChannels ]

forループで順番に読んでいくとしたらこんな感じ。

for o in 0..<outputChannels {
    for ky in 0..<kernelHeight {
        for kx in 0..<kernelWidth {
            for i in 0..<inputChannels {
                let index = Int(((o * kernelHeight + ky) * kernelWidth + kx) * inputChannels + i)
                print(String(format: "\(index): %.3f", kernelWeights[index]))
            }
        }
    }
}

参考: MPSCNN Weight Ordering - Stack Overflow

TensorFlowから書き出す

import tensorflow as tf

// 中略

w_conv1 = weight_variable([5, 5, 1, 32])
b_conv1 = bias_variable([32])

// 後略(学習)

こんな感じでTensorFlowで学習した重み/バイアスパラメータが得られたとして、これらをiOSのMPSCNN向けにそれぞれ"weights_xxx.dat", "bias_xxx.dat"という名前でファイルに書き出したい場合は次のようになります。

with open('weights_conv1.dat', 'w') as f:
  w_conv1_p = tf.transpose(w_conv1, perm=[3, 0, 1, 2])
  f.write(session.run(w_conv1_p).tobytes())
with open('bias_conv1.dat', 'w') as f:
  f.write(session.run(b_conv1).tobytes())

全結合層の重みについてはこんな感じ。

with open('weights_fc1.dat', 'w') as f:
    w_fc1_shp = tf.reshape(w_fc1, [7, 7, 64, 1024])
    w_fc1_p = tf.transpose(w_fc1_shp, perm=[3, 0, 1, 2])
    f.write(session.run(w_fc1_p).tobytes())

ポイントは、

  • open() でファイルを書き込みモードでオープン
  • 重みについては4次元テンソルの順序を変更する
  • 全結合層の場合は2次元テンソルを4次元にreshapeしてからtranspose(順序変更)する
  • write() および tobytes() でバイナリデータをファイルに書き込む

まとめ

  • モデルのファイル形式はiOSで読み込めるなら何でもいい
  • モデルの中身のフォーマットには決まりがある
    • 「重み」の4次元テンソルの格納順は weight[ outputChannels ][ kernelHeight ][ kernelWidth ][ inputChannels ]
  • 上記を満たせば、学習するために使うツールはTensorFlowでもChainerでも何でもいい

※とはいえ、現行のMPSCNNでは実装できないCNNもあるでしょうし、メモリの都合でiOSデバイス上では実行できないCNNもあると思います。そのへんの制約については調査しつつまた書きたいと思います。

関連


  1. SwiftでHDF5を読み取るHDF5Kitというオープンソースがあるとのこと