Swift for TensorFlowと新たなプログラミングモデル

はじめに

Swift for TensorFlowは2018年3月にTensorFlow Developer Summitで発表されました。
Swift for TensorFlowがどのようなプロジェクトなのかをSwiftを使用しているエンジニアや機械学習エンジニアの方にお伝えしたいと思います。
この記事の内容はDocumentの内容をまとめたものになります。
誤りなどありましたらご指摘していただけると幸いです。

Swift for TensorFlow とは何か？

最初多くの方は、Swift で書かれた単なる TensorFlow API ラッパーだと考えるかもしれませんが、Swift for TensorFlowはSwift にコンパイラと言語拡張を追加することによって、機械学習のデベロッパーや研究者にファーストクラスのユーザー エクスペリエンスを提供するものです。
Swift for TensorFlowは機械学習コードをSwiftの制御構造で記述でき、高いPerformanceのTensorFlowグラフを自動で作成します。
コード実行の前に次元が合わない行列計算や型エラーを知ることもできます。
またSwift for TensorFlowにはPythonを使用してきた方が移行する負担を減らすためのPythonとの互換性もあります。

Swift for TensorFlowは機械学習やデータサイエンスの今のツールを置き換える可能性を秘めたプロジェクトです。

Swift for TensorFlow がなぜ必要なのか？

研究や開発には高いUsabilityとPerformanceを実現することが必要です。

• Usability
  機械学習の発展やイノベーションには良いUsabilityが必要です。また最先端の機械学習モデルはますます複雑になっています。
• Performance
  ディープラーニングやニューラルネットワークには高い計算力が必要であり、スマートフォンやIoTデバイス上で機械学習の処理を実行するにはPerformanceが必要です。

近年様々な企業が機械学習に特化したプロセッサを開発しています。
これらのプロセッサを柔軟にサポートすることも重要です。

Machine LearningやData ScienceのコミュニティではPythonを使うのが定着しています。
しかしPythonには実行速度が遅いという課題があります。この課題を既存のFrameworkはどのように解決しているのでしょうか？
機械学習Frameworkの実装にはいくつかアプローチがあります。
まずは既存のFrameworkが採用しているアプローチの長所・短所を説明します。

Explicit graph building APIs

機械学習では計算の流れをグラフ構造により記述します。
このアプローチは計算グラフを構築した後にデータを流し処理を実行します。

# Build a graph.
a = tf.constant(2.0, dtype=tf.float32)
b = tf.constant(2.0, dtype=tf.float32)

c = a + b

# Tensor("add:0", shape=(), dtype=float32)
print(c)

sess = tf.Session()
# 4.0
print(sess.run(b))

このアプローチのライブラリとしてはTensorFlow session APIなどがあります。

• 長所
  Performanceの最適化が可能。
  GPUsやTPUsのような複数のプロセッササポートがあります。
• 短所
  PerformanceのためにUsabilityが犠牲になります。
  例えばTensorFlowは機械学習に特化した言語(DSL)を内包し、ユーザーはPythonを使用してメタプログラミングのようにDSLのTensorFlowグラフを作成します。
  このようなメタプログラミングでは機械学習モデルを作るのにプログラマーの負担になります。また二つの違う言語の実行が伴うのでデバッグも難しくなります。

複雑な機械学習のモデルを作るのには、ループや再帰のようなプログラミング構造のサポートが必要になってきているため、これらの短所があるとユーザーは少しPerformanceを犠牲にしてもUsabilityが良いアプローチを選択するでしょう。

Define-by-run - Interpreters

このアプローチはモデルを直接実行します。
グラフを作るのではなくPythonのプログラム自体がモデルになるアプローチです。
このアプローチの先駆者は日本のPrefferd Networks社のChainerです。その後Pytorch、TensorFlow Eager Executionがこのアプローチを採用するようになりました。Tensorflow 2.0ではSession APIではなくEager Executionがデフォルトになります。

a = tf.constant([2.0, 3.0], [4.0, 5.0])
b = tf.matmul(a, a)
# => tf.Tensor(
#      [[16. 21.]
#       [28. 37.]], shape=(2, 2), dtype=float32)
print(b)

• 長所
  これはPythonの制御構造が使用でき、デバッグも容易でUsabilityが高い。
• 短所
  このアプローチはPythonの実行が伴うため、遅い。
  インタープリタは計算処理全体がわからないことによって最適化が妨げられる。

Performance、表現力両方にとってベストな方法が模索されています。

Define-by-run - Tracing JITs

このアプローチはDefine by runのモデルを取り、最適化しPerformanceが良く実行するために書き直します。
PyTorchのversion 1.0ではJITコンパイラを導入し、PythonまたはPythonのサブセットで書かれたPyTorchモデルを取りPythonインタープリタに依存しないPerformanceの優れたモデルをエクスポートします。

• 長所
  このアプローチはExplicit graph building APIsのPerformanceの利点とDefine-by-run - InterpreterのUsabilityの利点が両方が得られます。
• 短所
  Python直接ではないコードの実行のためデバッグが困難。
  Pythonのサブセットは書き直したあとのモデルを実行するエンジンの表現能力に依存している。
  モデルに行列計算以外も含めると、traceを分けることになりPerformanceを損なう。
  Interpretersのアプローチ同様に計算処理全体がわからないことが最適化を妨げます。

新しいアプローチが必要

これまでのアプローチでは高いPerformanceとUsabilityを実現できていません。
Swift for TensorFlowは高いPerformanceとUsabilityを実現するプロジェクトです。

Graph Program Extraction: a new define-by-run approach(Swift for TensorFlow)

このアプローチはコンパイラが任意の機械学習コードを自動でグラフに変換します。
define-by-runのように書くことができ、同時にExplicit graph building APIsのような高いPerformanceも実現します。

もちろんGPUsやTPUsのような複数のプロセッサもサポートしています。
さらにグラフを作ると同時に、コードの実行前に行列計算のエラーや型のエラーを検知することができます。

let x = Tensor(shape: [2.4], repeating: 2)
let y = Tensor([1, 1, 1, 1])
let z = matmul(x, y)

// error: 'MatMul' op requires rank 2 inputs, have shapes: [2, 4], [4]
// let z = matmul(x, y)  

Swift for TensorFlowはこのアプローチによって高いUsabilityとPerformanceを実現します。

自動微分

機械学習は基本的に多項式があり、そのパラメータを正解データに近ずくように微調整します。
そのためには微分する必要があり、プログラムによって定義された数学的な関数の導関数をアルゴリズムによって機械的に求める手法が自動微分です。

多くの自動微分ツールはライブラリとして実装されていますが、Swift for TensorFlowはコンパイラに自動微分機能を実装することで以下の利点を得ています。

• ユーザーが定義したデータ構造をサポート
• TensorFlow以外のライブラリにも組み込み可能
• 早期にエラーをキャッチできる

自動微分を使うには#gradientキーワードを使用します。

func square(_ x: Float) -> Float {
  return x * x
}

let ∇square = #gradient(square)

square(5) // 25
∇square(5) // 10 

同じようにTensorでも導関数を求めることができます。

func prediction(for x: Tensor<Float>, using parameters: Parameters) -> Tensor<Float> {
  return x ・ parameters.w + parameters.b
} 

let prediction = prediction(for: image)
let prediction = #gradient(prediction)
∇prediction(image, parameters) // => (∂f/∂x, ∂f/∂parameters)

また微分ができない関数を微分しようとした場合、微分ができない関数を表示します。

プログラミング言語

Graph Program Extractionも自動微分にもコンパイラが必要なのでPythonではこのアプローチを実現することができません。
このアプローチにはSwiftがよく合っています。
Swiftは使うのも学ぶのも簡単で、GenericsやFunctional Programmingもあります。
またSwiftはPerformanceに優れています。Performanceを改善するための巨大なコミュニティもあり、将来的にはC++を置き換える可能性もあります。

Python Interoperability

Machine LearningやData ScienceのコミュニティではPythonを使うのが定着しています。
さらにこの巨大なコミュニティのサポートがあるNumpyやPandas、SciPyなどの科学計算ライブラリもあります。
Swift for TensorFlowではPythonのライブラリをあたかもSwiftのライブラリのように使うことができます。

// Python:
//    import numpy as np
//    a = np.arange(15).reshape(3, 5)
//    b = np.array([6, 7, 8])
let np = Python.import("numpy")
let a = np.arange(15).reshape(3, 5)
let b = np.array([6, 7, 8])

Programming Style

Swiftはインタープリターモード、スクリプトモード、Notebookモードをサポートしていて、Pythonを書くようにSwiftのコードをインタラクティブに書くことができます。

• インタープリターモード
  

• Notebookモード
  

Conclusion

Swift for TensorFlowが機械学習とデータサイエンスの分野でSwiftがメインで使われるようになる可能性を秘めているプロジェクトだということが少しでも伝わっていたら嬉しいです。
興味を持っていただいた方は以下のDocumentも読んでみてください。
Swift for TensorFlowはまだearly stageのため、Swift for TensorFlow を使ってディープ ラーニング モデルを書き直すにはまだ少し早すぎますが、機械学習 や言語、コンパイラに興味がある方は、さまざまな方法でこのプロジェクトに参加して貢献することができます。

Document

概要

• Swift for TensorFlow Design Overview
• Graph Program Extraction
• [Swift for TensorFlow - TFiwS (TensorFlow Dev Summit 2018)] (https://www.youtube.com/watch?v=Yze693W4MaU&frags=pl%2Cwn)
• Swift for TensorFlow (TensorFlow @ O’Reilly AI Conference, San Francisco '18)
• Why data scientists should start learning Swift
• Rebuildfm 206: Make Ruby Differentiable
  55:45辺りからSwift for TensorFlowについてわかりやすく説明されています。

プログラミング言語

• Why Swift for TensorFlow?
• Python Interoperability
• On Machine Learning and Programming Languages

自動微分

• Automatic Differentiation in Swift
• First-Class Automatic Differentiation in Swift: A Manifesto
• [Automatic Differentiation in Machine Learning: a Survey] (https://misreading.chat/2018/05/17/episode-9-automatic-differentiation-in-machine-learning-a-survey/)
  Pytorchがどのように自動微分を実装しているか、自動微分のアプローチについてわかりやすく説明されています。