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

最近 [Qoncept](https://qiita.com/organizations/qoncept) では [TensorFlow](https://www.tensorflow.org/) を使った案件が続いており、その中で TensorFlow を iOS 上で使いたいことがありました。

![deep-mnist.gif](https://raw.githubusercontent.com/qoncept/TensorSwift/master/Resources/DeepMnist.gif)

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

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

しかし、前者ついては、~~まだ TensorFlow を iOS 用にビルド[できなさそう](https://github.com/tensorflow/tensorflow/issues/16)でしたし~~できるようになりました（[コメント参照](http://qiita.com/koher/items/2c0bfca4d6e31cde674b#comment-3003403c289e35757189)）、たとえできたとしても C++ の API を Swift から叩くのは辛そうです。

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

そういうわけで、 __TensorFlow で学習させたモデルを使って Swift 上でテンソルを計算してくれるライブラリ [TensorSwift](https://github.com/qoncept/TensorSwift) を作りました__。

冒頭の画像は TensorFlow のチュートリアル [Deep MNIST for Experts](https://www.tensorflow.org/versions/r0.8/tutorials/mnist/pros/index.html) を再現し、 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` なんかと変わりません。

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

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

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

計算もシンプルです。

```swift
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 が効いているので代入してもコピーは発生しない

```swift
// ミュータビリティのコントロール
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 // コンパイルエラー
```

```swift
// 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](https://www.tensorflow.org/versions/r0.8/tutorials/mnist/pros/index.html) をオリジナルの Python 版のコードと見比べながら、 TensorSwift でどのように書けるのかを見ていきましょう。

### First Convolutional Layer

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

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

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

```python
# 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
// 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
# Python
h_conv2 = tf.nn.relu(conv2d(h_pool1, W_conv2) + b_conv2)
h_pool2 = max_pool_2x2(h_conv2)
```

```swift
// 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
# 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
// 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
# Python
y_conv=tf.nn.softmax(tf.matmul(h_fc1_drop, W_fc2) + b_fc2)
```

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

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

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

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

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

```python
# Python
tf.argmax(y_conv,1)
```

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

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

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

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

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

```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
// 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 版](https://github.com/qoncept/TensorKotlin)も作ってます。値型でないところ以外はほとんど同じように使えます。

## まとめ

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