TensorFlow計算グラフ最適化処理
11 likes4,131 views
本資料は、以下のイベントにて発表した資料です。 第1回 TensorFlow内部構造 勉強会（2019.3.4） https://tensorflowinternal.connpass.com/event/115653/ TensorFlowの内部で行われている計算グラフ最適化処理について、ソースコードから調べてみました。
![❖ テンソルの形状に関する演算を最適化
Shape Optimizer
Grappler
tensorflow/core/grappler/optimizers/shape_optimizer.cc
Shape Prod Size
単純にテンソルの要素数
(tensor.size)を返せばよい
tensor.shape: [32, 3, 64, 64]
↓
32 * 3 * 64 * 64 = テンソルの要素数](https://image.slidesharecdn.com/tensorflowgraphoptimization-190304221603/85/TensorFlow-20-320.jpg)
TensorFlow計算グラフ最適化処理
- 2. 自己紹介 ❖ Nuka（@nuka_137） ➢ TensorFlowの内部構造記事書いてます https://qiita.com/nuka/items/028780aab338d9b1a7c2 ➢ TensorFlowなにもわからない（TensorFlow歴1年） ❖ 仕事 ➢ HPC向け分散ファイルシステム（1年前くらいまで） ➢ 深層学習フレームワーク、周辺ツール（今） （TF、PyTorch、TVM、ONNX、nGraph、Glowあたりを中心に） 最近調べていた、計算グラフの最適化周りについて話します
- 3. この資料は、 TensorFlowのソースコードから 計算グラフの最適化処理を 調べてまとめたものです https://github.com/tensorflow/tensorflow Commit ID: a8e5c41c5bbe684a88b9285e07bd9838c089e83b (Tag: v1.13.0-rc0) 変化が激しいため、 バージョン注意
- 4. TensorFlowの基本 エッジ：テンソルの流れConst Sqrt Conv Add ReLU ノード：演算 ❖ ユーザが計算グラフを定義して、Session.runで実行 計算グラフが評価されて、ユーザに結果が返る
- 6. デバイス割り当て ❖ 計算グラフの各ノードに、実行デバイスを割り当て Const Add Sqrt Conv Const Add Sqrt Conv CPU CPU CPU GPU デバイス割り当て前 デバイス割り当て後 ノードにデバイスを 割り当てる
- 11. import tensorflow as tf from tensorflow.python.framework.function import Defun @Defun(tf.float32, tf.float32) def fn(x, y): return x + y a = tf.constant(1.5) b = tf.constant(2.7) c = a * b d = a + b e = fn(c, d) f = tf.constant(-7.8) out = e + f tf.Session().run(out) 計算グラフの最適化を確認しよう ❖ 最適化を狙った計算グラフを作ってみる TensorBoard サブグラフ 全入力が Constの演算
- 12. 最適化後の計算グラフを見るために ❖ RunMetadata ➢ Session.runの実行情報を取得可能 メモリ使用量/確保時間、演算毎の実行時間・・・ ❖ RunOptionsの設定で、最適化後の計算グラフも取得可能 run_opts = tf.RunOptions() run_opts.output_partition_graphs = True run_md = tf.RunMetadata() ops = ... with tf.Session() as sess: sess.run(ops, run_metadata=run_md, options=run_opts) print(run_md.partition_graphs) 最適化後の計算グラフを 出力する設定 最適化後の計算グラフが 保存されている Session.runに RunMetadataとRunOptionsを渡す
- 15. 計算グラフ最適化処理 ❖ Grappler ➢ グラフ分割前の最適化 ❖ GraphOptimizer ➢ グラフ分割後の最適化 ❖ GraphOptimizationPass ➢ 4フェーズで追加可能な最適化 グラフ構築 デバイス割り当て Grappler グラフ分割 GraphOptimizer GraphOptimizationPass
- 16. Grappler ❖ Model Pruner ❖ Function Optimizer†1 ❖ Constant Folding†1 ❖ Shape Optimizer ❖ Remapper ❖ Layout Optimizer ❖ Memory Optimizer ❖ Arithmetic Optimizer ❖ Auto Parallel ❖ Loop Optimizer ❖ Dependency Optimizer ❖ Debug Stripper ❖ Scoped Allocator Optimizer ❖ Pin to Host Optimizer tensorflow/core/grappler/optimizers/meta_optimizer.cc †1 p11の計算グラフで実際に実行された最適化処理
- 17. Model Pruner tensorflow/core/grappler/optimizers/model_pruner.cc ❖ 実行不要なノードを削除 Grappler Fetch Fetch 演算結果に影響がない ノードを削除
- 20. ❖ テンソルの形状に関する演算を最適化 Shape Optimizer Grappler tensorflow/core/grappler/optimizers/shape_optimizer.cc Shape Prod Size 単純にテンソルの要素数 (tensor.size)を返せばよい tensor.shape: [32, 3, 64, 64] ↓ 32 * 3 * 64 * 64 = テンソルの要素数
- 22. ❖ GPUに最適なテンソルフォーマット(NCHW)へ変換 Layout Optimizer tensorflow/core/grappler/optimizers/layout_optimizer.cc Grappler CPU GPU NHWC NHWC CPU Trans pose NHWC NHWC ↓ NCHW GPU NHWC テンソルフォーマットを変換
- 24. ❖ 勾配計算に必要なデータを再計算 Memory Optimizer② tensorflow/core/grappler/optimizers/memory_optimizer.cc Grappler Op1 Op2 Op Grad Op1 Op2 Op1 Grad Op1 再計算
- 25. ❖ 計算の中間結果を保持することで、メモリ使用量削減 Memory Optimizer③ tensorflow/core/grappler/optimizers/memory_optimizer.cc Grappler AddN AddN 中間結果を保持
- 26. ❖ 算術演算を簡易化 Arithmetic Optimizer tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc Grappler Add Add AddN 例：Add(Add(x, y), z) ⇒ AddN(x, y, z)
- 27. ❖ 計算グラフを複製し、各GPUで並列実行 Auto Parallel tensorflow/core/grappler/optimizers/auto_parallel.cc Grappler GPU #1 GPU #2
- 29. ❖ ループ内不変ノードをループ外へ移動 Loop Optimizer② tensorflow/core/grappler/optimizers/loop_optimizer.cc Grappler While While ループ内不変 ループ外へ移動
- 30. ❖ Control Edge（データ移動がないエッジ）の最適化 Dependency Optimizer tensorflow/core/grappler/optimizers/dependency_optimizer.cc Grappler Identity 1 3 2 1 3 2 Control Edge 4 4 1の出力テンソルを そのまま4に流す 1～3実行後に実行 1～3実行後に実行
- 32. Scoped Allocator Optimizer tensorflow/core/grappler/optimizers/scoped_allocator_optimizer.cc Grappler ❖ メモリをプーリングし、メモリ確保＆解放時間短縮 2 1 1 2 3 Memory Pool プールした メモリから確保 3
- 33. ❖ 演算効率が良くなるように、CPUでの演算に変更 Pin to Host Optimizer tensorflow/core/grappler/optimizers/pin_to_host_optimizer.cc Grappler Op Const Const GPU GPU 入力は CPUメモリを期待 Op Const Const CPU CPU デバイス割り当てを CPUに変更
- 34. GraphOptimizer ❖ Common Subexpression Elimination（CSE） ❖ Constant Folding ➢ tensorflow/core/common_runtime/constant_folding.cc ➢ GrapplerのConstant Foldingと同じ ❖ Function Inlining ➢ tensorflow/core/common_runtime/function.cc ➢ GrapplerのFunction Optimizerと同じ tensorflow/core/common_runtime/graph_optimizer.cc
- 35. ❖ 演算結果が同じになる演算を共通化 Common SubExpression Elimination tensorflow/core/graph/optimizer_cse.cc GraphOptimizer Add Add Add
- 36. GraphOptimizationPass ❖ PRE_PLACEMENT ➢ XLA†1 向けのグラフ変形 ➢ If/While Lowering ➢ Parallel Concat Optimizer ➢ AccumulateN Optimizer ❖ POST_PLACEMENT ➢ NCCL Replace ❖ POST_REWRITE_FOR_EXEC ➢ XLA向けのグラフ変形 ❖ POST_PARTITIONING ➢ Intel MKL向けの最適化 グラフ構築 デバイス割り当て Grappler グラフ分割 GraphOptimizer †1 グラフコンパイラ
- 37. ❖ PrimitiveなControl Flow Opsから、If/Whileへの変換 ➢ 「If/While Lowering」の逆変換 ➢ 後発となるXLAは、If/Whileの方が扱いやすいため ❖ TensorFlowのグラフから、XLAのグラフへ変換 XLA GraphOptimizationPass tensorflow/compiler/tf2xla/functionalize_control_flow_pass_registration.cc tensorflow/compiler/jit/jit_compilation_pass_registration.cc
- 38. ❖ If/Whileを、PrimitiveなControl Flow Opsに変換 ➢ If ⇒ Switch, Merge ➢ While ⇒ Switch, Merge, Enter, Exit, NextIteration, LoopCond ❖ TensorFlowの内部処理に合わせるための変形処理 If/While Lowering tensorflow/core/common_runtime/lower_if_while.cc GraphOptimizationPass
- 39. ❖ ParallelConcat演算を実現するための計算グラフ変形 ➢ ParallelConcat：並列版のConcat処理 Parallel Concat Optimizer tensorflow/core/common_runtime/parallel_concat_optimizer.cc GraphOptimizationPass A B Parallel Concat A B _Parallel Concat Update _Parallel Concat Update _Parallel Concat Start 部分ごとに Concatを 並列実行 Identity
- 40. ❖ AccumulateN演算を実現するための計算グラフ変形 ➢ AccumulateN：AddNと同じ演算だが、入力が揃ったところから演算 AccumulateN Optimizer tensorflow/core/common_runtime/accumulate_n_optimizer.cc GraphOptimizationPass A B Accumu lateN V2 A B Assign Add Assign Add Assign Const (0) 入力が揃った ところから 実行可能
- 41. ❖ ncclによるGPU間集合通信を実現するための計算グラフ変形 ➢ nccl：NVIDIAが提供する、GPU向けの集合通信用ライブラリ Nccl Replace tensorflow/core/nccl/nccl_rewrite.cc GraphOptimizationPass Nccl Reduce _Nccl Broadcast Send _Nccl Broadcast Send _Nccl Broadcast Recv
- 42. ❖ ノードを結合し、Intel MKLで得意な演算に変更 MKL① tensorflow/core/graph/mkl_layout_pass.cc GraphOptimizationPass Conv 2D Bias Add _Mkl Conv2D WithBias ノードを結合して、 Intel MKLに最適な演算に変更
- 43. ❖ テンソルフォーマット変換（Intel MKL ⇒ TensorFlow） MKL② tensorflow/core/graph/mkl_tfconversion_pass.cc GraphOptimizationPass MKL Op Op1 Tensor Flow Intel MKL ↓ Tenosor Flow Intel MKL MKL Op MklTo TF Intel MKL Op1 Tensor Flow Intel MKLの テンソルフォーマットから戻す
- 44. 最適化の有効/無効を制御する ❖ ConfigProtoを使って、最適化の有効/無効を制御可能 ➢ Grappler: ConfigProto.graph_options.rewrite_options ➢ GraphOptimizer: ConfigProto.graph_options.optimizer_options import tensorflow as tf from tensorflow.core.protobuf import config_pb2 from tensorflow.core.protobuf import rewriter_config_pb2 cfg = config_pb2.ConfigProto() cfg.graph_options.rewrite_options.constant_folding = rewriter_config_pb2.RewriterConfig.OFF cfg.graph_options.optimizer_options.opt_level = config_pb2.OptimizerOptions.L0 with tf.Session(config=cfg) as sess: sess.run(...) GraphOptimizer「Constant Folding」 無効化 Grappler「Constant Folding」 無効化 tf.SessionにConfigProtoを渡す
- 45. ❖ Constant Folding：無効 ❖ Function Optimizer：有効 最適化の無効化を確認する サブグラフ 呼び出し処理 サブグラフ呼び出し処理は 消えている 最適化前 （p11で紹介した計算グラフ） 最適化後 全入力がConstの 演算はそのまま 全入力が Constの演算
- 47. Eager ModeとGraph Mode ❖ TensorFlow 2.0では、Eager Modeがデフォルト化 ➢ Eager Modeでは、最適化処理が適用されない ❖ tf.functionを利用することで、Graph Modeで動作可能 ➢ Graph Modeでは、従来の最適化処理が適用される import tensorflow as tf @tf.function def graph_mode(): a = tf.constant(1.5) b = tf.constant(2.7) c = a * b d = a + b return c + d out = graph_mode() GraphModeで動作するため、 最適化処理も適用される
- 48. ❖ Design Documentに変更方針が記載されている ➢ TensorFlow Design Document（tensorflow/community） https://github.com/tensorflow/community ➢ Design Documentから見たTensorFlow 2.0の変更点 https://qiita.com/nuka/items/6966efeddceb96012819 TensorFlow 2.0の内部構造
- 49. ❖ 最適化後の計算グラフは、RunMetadataを使って 確認できる まとめ ❖ 計算グラフの最適化を行う仕組みは、3つある ➢ 計算グラフ分割前の最適化「Grappler」 ➢ 計算グラフ分割後の最適化「GraphOptimizer」 ➢ 各計算グラフ変形の前後に行う最適化「GraphOptimizationPass」 ❖ TensorFlow 2.0でtf.functionを利用することにより、 最適化処理を活用できる
- 50. 参考資料 ❖ QiitaにもTensorFlowの最適化処理をまとめています ➢ TensorFlow内部構造解析 (4.4) 計算グラフ最適化処理1 Grappler https://qiita.com/nuka/items/f1b0fe9c820e4d5f80cc ➢ TensorFlow内部構造解析 (4.6) 計算グラフ最適化処理2 GraphOptimizationPass https://qiita.com/nuka/items/808fbec073f97be8b872 ➢ TensorFlow内部構造解析 (4.7) 計算グラフ最適化処理3 GraphOptimizer https://qiita.com/nuka/items/1936ad509cd4e68d2fad
- 51. 宣伝 ❖ 技術書典6にて、TechBoosterから頒布予定の技術書に、 「TensorFlowの計算グラフ最適化処理」を寄稿します ➢ https://techbookfest.org/event/tbf06 ➢ 日程：2019/4/14（日）11:00-17:00 ➢ 場所：池袋サンシャインシティ2F 展示ホールD（文化会館ビル2F） ➢ 主催：TechBooster/達人出版会
- 52. おわり