開発日誌スマートフォン用の表示で見る開発日誌ソフトウェアエンジニアリングとか機械学習とか2018-05-21TensorFlow.jsがWebGLで並列計算を実行する仕組みを理解する先日、機械学習モデルの学習・デプロイのためのJavaScriptライブラリ「TensorFlow.js」がGoogleから公開されました。APIや、TensorFlow本体・Kerasで学習したモデルをインポートするツールを備えているとのことです。このライブラリでは、WebGLを介してGPUを利用することで計算を高速化しているのが特徴です。WebGLはJavaScriptでグラフィック描画を高速に行うための機能ですが、これをどのようにして機械学習の計算に利用しているのか調べてみました。そもそもWebGLは何をするためのもの?TensorFlow.jsでの計算の仕組みを理解するには、そもそもWebGLがどのようにグラフィック描画をしているかを理解する必要があります。コンピュータグラフィクスの基本的な処理に「レンダリング」というものがあります。これはプログラム内のデータ表現から、画像や映像を生成する処理です。元になるデータは、物体の形や色や質感、光源の色や位置、視点の位置などがあります。同じ物体でも光の当たり方や視点の位置によって見え方が変わるので、物体を移動・回転させたりすると、その都度、物体の各点がどのように見えるかを計算する必要があります。この計算を高速に行うために利用されるのがGPUです。GPUではこの計算を数千単位のコアで並列して実行することができます。TensorFlow.jsではどうやって利用しているの?グラフィック描画ために開発されたWebGLを、機械学習のような汎用計算に利用してやろうというのがTensorFlow.jsです。GPUでの計算は「シェーダ」という機能が担当しており、特にディスプレイ上の各点に何色を表示させるかを「フラグメントシェーダ」が担当しています。WebGLでは、このフラグメントシェーダに行わせる処理を自分で記述することができます。通常フラグメントシェーダでは、物体や光源の情報を受け取り、ディスプレイの各点に表示すべき色情報を出力するのですが、この入出力機能を「関数」ととらえ、より一般的な計算に使ってしまおうというのが、TensorFlow.jsでやっていることです。TensorFlow.jsでは、入出力テンソルをテクスチャ(色情報を定義する2次元の配列)として表現し、出力テンソルの要素1つ1つに対してシェーダを並列実行することで、目的の計算を実行しています。TensorFlow.jsの構成要素ここからはソースコードを見ながら追っていきます。TensorFlow.jsは大きく2つの部分から構成されています。1つはWebGLを利用して計算を行うコア部分。今回はWebGLの利用方法に興味があるので後者を詳しく見ていきます。https://github.com/tensorflow/tfjs-layershttps://github.com/tensorflow/tfjs-coretfjs-coreの主な構成要素Environemntこのクラスはモジュールのインポート時にインスタンス化され、グローバル名前空間に"ENV"として保持されています。いわゆるシングルトンクラスで、他のクラスからも常に同一のインスタンス"ENV"が使われます。使用できるKernelBackendの一覧を持っていて、それぞれのKernelBackendは優先度が設定されています。デフォルトではMathBackendWebGLのほうが高い優先度が設定されているので、WebGLが使える環境ではMathBackendWebGLを、使えない環境ではMathBackendCPUを使うようになっています。EngineEnvironmentからKernelBackendを渡されてインスタンス化されます。EngineクラスにはrunKernelというメソッドが用意されていて、後述する各種オペレータがこのrunKernelメソッドを実行すると、Engineが持つKernelBackendを通じてオペレータの計算が実行されます。tf.addなどの各種オペレータがスタティックメソッドとして定義されています。ユーザーはこれらのメソッドを使ってモデルを構築します。行列乗算の定義はこちら。}Env.engine.runKernelに順伝播関数、入力テンソル、逆伝播関数を渡しているのが確認できます。KernelBackendKernelBackendはオペレータの計算を実行するクラスのインタフェースで、バックエンドが実装すべきメソッド群の宣言がずらずらと並んでいます。KernelBackendインタフェースを実装したクラスは2つあり、それぞれCPU用・GPU用の実装になっています。中略}MathBackendCPUKernelBackendインタフェースをCPU向けに実装したのがMathBackendCPU。各メソッドをCPU上で実行するためのコードが記述されています。行列乗算の実装はこちら。}transposeというオプションがあるのでちょっと複雑に見えますが、普通に3重のfor文を回して各要素を乗算していることがわかります。MathBackendWebGLKernelBackendインタフェースをGPU向けに実装したのがMathBackendWebGL。各メソッドをGPU上で実行するためのコードが記述されています。行列乗算の実装はこちら。}MatMulProgramというクラスをインスタンス化してコンパイルして実行していますが、これだけだとよくわかりません。ここが今回最も興味のあるところなので、詳しく見ていきましょう。MatMulProgramMatMulProgramでは、行列乗算を計算する際にフラグメントシェーダで実行すべき処理が定義されています。}}入力テンソルの変数名、出力テンソルのサイズ、そして「userCode」という文字列をプロパティとして持っています。このuserCodeがフラグメントシェーダのコードになります。フラグメントシェーダは、シェーディング言語というC言語風の独自の言語で書く必要があります。resRC.y)」で、担当している出力要素の値を計算しています。dotARowBColの中身を見てみると、CPUでの計算と異なりfor文が1重しかないのがわかります。フラグメントシェーダでは、担当する出力1要素について計算すれば良いので、出力全要素を走査する2つのループは必要ありません。ちなみにdotARowBColではスカラ同士を直接乗算せずに、4つのスカラを4次元ベクトルと見立ててベクトル同士の乗算を行っています。GPUにはベクトル演算用のプロセッサが搭載されているので、こうすることでスカラ同士を1つずつ乗算するよりも高速化できるようです。compileAndRunメソッド上で定義したフラグメントシェーダをコンパイルして実行しているのがここ。}「uploadToGPU」で入出力テンソルをテクスチャに変換し、GPUからアクセスできるようにしています。実際のコンパイル処理・実行処理は「gpgpu_math.compileProgram」や「gpgpu_math.runProgram」で行われています。このあたりをもっと深掘っていくとWebGLのAPIが見えてきますが、長くなってしまうのでここまで。まとめ今回はTensorFlow.jsのコードを読み、WebGLを利用した汎用計算の実装について見てきました。TensorFlow.jsのAPIの裏で、CPU用・GPU用のバックエンド処理が実装されているGPUでの計算はフラグメントシェーダとして記述されている入出力テンソルはテクスチャとしてGPUに渡している機械学習関連のライブラリはどんどん開発されて便利になってきていますが、こういった低レイヤーの複雑な処理を隠蔽してくれていることを改めて感じました。今度はTensorFlow本体のコードも読んでみたいと思います。機械学習を始めようと思ったときに必要になるのがGPU。»プロフィールid:tfujiwar418ソフトウェアエンジニアリングについて書きます。月別アーカイブはてなブログをはじめよう!tfujiwar418さんは、はてなブログを使っています。ブログを報告するスターをつけました引用をストックしましたストック一覧を見る閉じる引用するにはまずログインしてくださいログイン閉じる引用をストックできませんでした。再度お試しください閉じる限定公開記事のため引用できません。