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

iOS 11からCore MLが追加されましたが、もちろんiOS 10では利用することができません。代わりに、Metal Performance Shaders（MPS）フレームワークを利用してiOSデバイス上でCNN（Convolutional Neural Network）の推論を行うことが可能です。

• iOSの機械学習フレームワークの比較 - Core ML / Metal Performance Shaders (MPSCNN) / BNNS - Qiita

このMPSのCNN関連APIはiOS 10で追加されました。フレームワーク名に"Metal"とあるとおり、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)

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

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

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

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

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

• iOSのMetal Performance Shadersでニューラルネットを実行する際のモデルの渡し方 - Qiita

##重みデータの順序

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

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

MPSCNNConvolution

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

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

というわけで、「重み」の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() でファイルを書き込みモードでオープン
• 重みについては４次元テンソルの順序を変更する
• 全結合層の場合は２次元テンソルを４次元にreshapeしてからtranspose（順序変更）する
• write() および tobytes() でバイナリデータをファイルに書き込む

##まとめ

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

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

##関連

• TensorFlow for iOS のモデルについて - Qiita

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