TensorFlowで学習したモデルを使ってiOS/Swiftで実行する #Swift

最近 Qoncept では TensorFlow を使った案件が続いており、その中で TensorFlow を iOS 上で使いたいことがありました。

ぱっと浮かんだ選択肢は次の二つでした。

TensorFlow を iOS 用にビルドして C++ の API を Swift から叩く
学習は TensorFlow / Python で行って、テンソルの計算だけを iOS / Swift でシミュレーションする

しかし、前者ついては、~~まだ TensorFlow を iOS 用にビルドできなさそうでしたし~~できるようになりました（コメント参照）、たとえできたとしても C++ の API を Swift から叩くのは辛そうです。

TensorFlow がありがたいのは学習時の自動微分等の機能であって、学習済みのモデルを利用するときはただテンソルの計算をしてるだけです。別に学習を iOS 上でやりたいわけではありませんでしたし、それくらいの計算なら自力でシミュレーションできそうです。

そういうわけで、 TensorFlow で学習させたモデルを使って Swift 上でテンソルを計算してくれるライブラリ TensorSwift を作りました。

冒頭の画像は TensorFlow のチュートリアル Deep MNIST for Experts を再現し、 iOS アプリで数字認識をしている様子です。こちらのアプリは TensorSwift のデモアプリとしてリポジトリに同梱されているので、興味があれば clone して実行してみてください。

以下、 TensorSwift がどんなものなのか、どのように使うのかを説明します。

インストール

CocoaPods

pod 'TensorSwift', '~> 0.1'

Carthage

github "qoncept/TensorSwift" ~> 0.1

`Tensor`

TensorSwift では学習は TensorFlow で行うことを前提に、学習済みのモデルを使ってテンソルを計算する部分しか担当しません。そのため、 TensorFlow で出てくる色々な ~~ややこしい~~ 概念がありません。 placeholder も Variable も Graph も Session もありません。 TensorSwift の Tensor はただの値です。 CGPoint なんかと変わりません。

let a = Tensor(shape: [2, 3], elements: [1, 2, 3, 4, 5, 6])

単純に↑のようにすれば次のようなテンソルが出来上がります。

[[1, 2, 3],
 [4, 5, 6]]

計算もシンプルです。

let a = Tensor(shape: [2, 3], elements: [1, 2, 3, 4, 5, 6])
let b = Tensor(shape: [2, 3], elements: [7, 8, 9, 10, 11, 12])
let sum = a + b

TensorFlow と違って Session を run しなくても即時計算が実行されます。 sum は次のようになります。

[[8, 10, 12],
 [14, 16, 18]]

なお、今のところ Tensor の要素の型は Float しかサポートされていません。

値型としての特徴

Swift の特徴を活かして、 Tensor は値型となるよう struct で実装されています。その結果 Tensor は次の特徴を持ちます。

var か let かでミュータビリティをコントロールできる
値型だけど Copy-on-Write が効いているので代入してもコピーは発生しない

// ミュータビリティのコントロール
var a = Tensor(shape: [2, 3], elements: [1, 2, 3, 4, 5, 6])
a[1, 2] = 111 // OK
let b = a
b[1, 2] = 222 // コンパイルエラー

// Copy-on-write
let zeros = Tensor(shape: [100, 100, 1000], elements: [Float](count: 10000000, repeatedValue: 0.0))
let foo = zeros // 値型なのに代入してもコピーは発生しない

Deep MNIST for Experts を再現

それでは、 Deep MNIST for Experts をオリジナルの Python 版のコードと見比べながら、 TensorSwift でどのように書けるのかを見ていきましょう。

First Convolutional Layer

最初の畳込み＆プーリング層は、 Python / TensorFlow のコードでは次のようになっていました。

# Python
h_conv1 = tf.nn.relu(conv2d(x_image, W_conv1) + b_conv1)
h_pool1 = max_pool_2x2(h_conv1)

なお、 conv2d および max_pool はチュートリアルの中で定義された↓のような関数です。

# Python
def conv2d(x, W):
  return tf.nn.conv2d(x, W, strides=[1, 1, 1, 1], padding='SAME')

def max_pool_2x2(x):
  return tf.nn.max_pool(x, ksize=[1, 2, 2, 1],
                        strides=[1, 2, 2, 1], padding='SAME')

これらの関数は書き下すことにして、 Swift / TensorSwift では次のようになります。

// Swift
let h_conv1 = (x_image.conv2d(filter: W_conv1, strides: [1, 1, 1]) + b_conv1).relu
let h_pool1 = h_conv1.maxPool(kernelSize: [2, 2, 1], strides: [2, 2, 1])

見比べると、 conv2d と maxPool が TensorSwift では Tensor のメソッドになっているのがわかると思います。 relu も関数ではなく Tensor のプロパティです。これは、メソッドやプロパティの方が Swift っぽいという判断です（ Swift では map や filter なんかも Python と違って関数ではなくメソッドとして用意されています）。また、個人的にはメソッドやプロパティの方が、コード上の記述順と実行される計算の順序が一致するので書きやすい／読みやすいと思います。例えば、 h_conv1 の例では、 Swift では conv2d を計算してから b_conv1 を足し、その結果を relu に適用するという計算順がコード上の順番と一致します。しかし、 Python では、実行順が最後である relu が最初に来てしまいます。

そのような表面上の違いに加えて、一番大きな違いは conv2d, maxPool に渡す strides と、 maxPool の kernelSize が 4 階ではなく 3 階のテンソルになっていることです（ kernelSize はテンソルそのものではなく Shape ですが）。学習時にはまとめて複数のデータを渡して計算させることが多いですが、学習済みモデルを利用するときには一つずつデータを渡すケースが多いと思います。そのため、いちいち 4 階のテンソルに reshape しなくても 3 階のテンソルとして渡せるようにしています。

なお、今のところ padding は SAME しかサポートしていません。

Second Convolutional Layer

次の層は最初の層とまったく同じなので説明は省略してコードだけ載せます。

# Python
h_conv2 = tf.nn.relu(conv2d(h_pool1, W_conv2) + b_conv2)
h_pool2 = max_pool_2x2(h_conv2)

// Swift
let h_conv2 = (h_pool1.conv2d(filter: W_conv2, strides: [1, 1, 1]) + b_conv2).relu
let h_pool2 = h_conv2.maxPool(kernelSize: [2, 2, 1], strides: [2, 2, 1])

Densely Connected Layer

次の全結合層を見比べてみましょう。

# Python
h_pool2_flat = tf.reshape(h_pool2, [-1, 7*7*64])
h_fc1 = tf.nn.relu(tf.matmul(h_pool2_flat, W_fc1) + b_fc1)

keep_prob = tf.placeholder(tf.float32)
h_fc1_drop = tf.nn.dropout(h_fc1, keep_prob)

// Swift
let h_pool2_flat = h_pool2.reshape([1, 7 * 7 * 64])
let h_fc1 = (h_pool2_flat.matmul(W_fc1) + b_fc1).relu

TensorSwift では reshape と matmul が Tensor のメソッドになっています。また、 Swift / TensroSwift で dropout がないのは、 dropout は過学習を防ぐ仕組みで学習時のみ必要だからです。

reshape 時の Shape を見比べると、 Python で [-1, 7*7*64] となっているのが [1, 7 * 7 * 64] となっています。 -1 （不定）が 1 になっているのは、学習時（ Python / TensorFlow ）は N 件のデータをまとめて渡すのに対して、今（ Swift / TensorSwift ）は 1 件だけを識別したいためです。

Readout Layer

最後に、 digit ごとの確率に対応した 10 次元ベクトルに変換します。

# Python
y_conv=tf.nn.softmax(tf.matmul(h_fc1_drop, W_fc2) + b_fc2)

// Swift
let y_conv = (h_fc1.matmul(W_fc2) + b_fc2).softmax

ここでは softmax を使いますが、これも TensorSwift では関数ではなくプロパティになっています。

`y_conv` からクラスへの変換

実際にどのクラスとして識別されたかを判断するには、 y_conv の最も大きい要素のインデックスを取得する必要があります。

TensorFlow では次のように argmax を使ってインデックスを取り出せます。

# Python
tf.argmax(y_conv,1)

しかし、今のところ TensorSwift では argmax を提供していません。 argmax がなくても次のようにしてインデックスを取得できます。

// Swift
y_conv.elements.enumerate().maxElement { $0.1 < $1.1 }!.0

モデルの書き出し、読み込み

W_conv1, b_conv1 などは Python / TensorFlow で学習させた値を使います。本当は TensorFlow で書きだした Checkpoint ファイルを読み込めればいいのですが今のところそのような機能はありません。何らかの方法で Python 側でテンソルを書き出し、それを Swift で読み込む必要があります。

Python に関しては大した経験がないのでもっと良いやり方があるかもしれませんが、参考までに僕は次のようにやっています。

# Python ( 書き出し )
import struct

ndarray = W_conv1.eval(session=sess)
list = ndarray.reshape([-1]).tolist()
f = open("W_conv1", "wb")
f.write(struct.pack("%df" % len(list), *list))
f.close()

// Swift ( 読み込み )
let data = NSData(contentsOfFile: "W_conv1")!
let array = [Float](UnsafeBufferPointer(start: UnsafeMutablePointer<Float>(data.bytes), count: data.length / 4))
let W_conv1 = Tensor(shape: [5, 5, 1, 32], elements: array)

その他

TensorSwift はテンソルの計算部分しか含まないので軽量なライブラリになった
TensorSwift は単純にテンソルを計算するだけなので、試してないけど Chainer とでも Caffe とでも組み合わせて使うことができるはず
現状ではまだα版くらいの完成度なので TensorFlow の API のごく一部しか再現できておらず、 Deep MNIST for Experts の畳み込みニューラルネットワークが再現できるくらいの段階
GPU を使った計算の高速化はしていない
行列演算やコンボリューションの計算で BLAS を使った高速化はしている

Kotlin

実は Kotlin 版も作ってます。値型でないところ以外はほとんど同じように使えます。

まとめ

iOS で TensorFlow を使いたくなった
TensorFlow を iOS 用にビルドするのは現状では難しそうだし、たとえできても Swift から C++ の API を叩くのは辛そう
学習は Python でやると割りきって Swift では学習済みのモデルを使って計算だけすることにした
テンソルの計算部分だけをシミュレーションする軽量なライブラリ TensorSwift を作った
TensorSwift を使って TensorFlow のチュートリアル Deep MNIST for Experts の動作を再現した
デモアプリとして、 Deep MNIST for Experts の結果を使って手書き文字認識する iOS アプリを作った

TensorFlowで学習したモデルを使ってiOS/Swiftで実行する