本記事は、連載記事 TensorFlow内部構造解析の1つで、GraphOptimizerによるTensorFlowの計算グラフ最適化処理について説明した記事になります。

TensorFlowにおける計算グラフの最適化

TensorFlowでは、以下の3つの最適化機能の仕組みを使って、ユーザが定義した計算グラフを最適化した後に実行します。

GraphOptimizationPass
Grappler
GraphOptimizer ★本記事で説明

本記事では、これら3つの最適化の仕組みの中で最後に行う最適化処理である、GraphOptimizerによる計算グラフ最適化処理について説明します。

GraphOptimizerによる最適化項目

GraphOptimizerは、TensorFlowの計算グラフを実行デバイスごとに分割（Partition）した後に行う最適化処理です。
実行デバイスごとに分割したことにより新たに最適化可能になるケースがあるため、本最適化が提供されています。

GraphOptimizerのソースコードは tensorflow/core/common_runtime/graph_optimizer.cc で、適用される最適化は GraphOptimizer::Optimize の処理から確認することができます。
以下に示す3種類の最適化をサポートします。
なお、Common Subexpression EliminationとConstant Foldingは、デフォルトで有効化されています。

Common Subexpression Elimination
Constant Folding
Function Inlining

1. Common Subexpression Elimination

ソースコード：tensorflow/core/graph/optimizer_cse.cc

共通の演算を行うことで演算結果が同じ値になるノードを検索し、1つのノードに置き換えます。
これにより、演算回数を減らすことができます。

2. Constant Folding

ソースコード：tensorflow/core/common_runtime/constant_folding.cc

ノードの入力が全て定数値で構成されるテンソルである場合、変換対象とするノードの演算をCPUで実行し、Constノードで置き換えます。

3. Function Inlining

ソースコード：tensorflow/core/common_runtime/function.cc

ユーザが定義したサブグラフ（tf.Defun を使ってユーザ定義したFunction）をインライン展開することにより、サブグラフの呼び出しコストを削減します。

GraphOptimizerによる最適化の有効・無効化

GraphOptimizerによる計算グラフの最適化は、tensorflow/core/protobuf/config.proto でProtocol Buffers形式により定義された OptimizerOptions を使って、ユーザが有効・無効を制御できます。

GraphOptimizerによる最適化の有効・無効の切り替えは、Grapplerと比較して少し複雑です。
OptimizerOptions には、各最適化の有効・無効を制御するためのフィールドの他に、最適化レベルと呼ばれるものがあるためです。

tensorflow/core/protobuf/config.proto

// Options passed to the graph optimizer
message OptimizerOptions {
  // 各最適化の有効・無効を制御
  bool do_common_subexpression_elimination = 1;
  bool do_constant_folding = 2;
  bool do_function_inlining = 4;

  // 最適化レベル
  enum Level {
    // L1 is the default level.
    // Optimization performed at L1 :
    // 1. Common subexpression elimination
    // 2. Constant folding
    L1 = 0;

    // No optimizations
    L0 = -1;
  }
  Level opt_level = 3;
}

最適化レベルを示すフィールド opt_level が L1 に指定されていると、do_common_subexpression_elimination と do_function_inlining の設定値に関わらず、Constant FoldingとCommon Subexpression Eliminationが有効化されます。

tensorflow/core/common_runtime/graph_optimizer.cc

GraphOptimizer::GraphOptimizer(const OptimizerOptions& opts) : opts_(opts) {
  // opt_levelがL1の場合は、強制的にConstant Foldingと
  // Common Subexpression Eliminationが有効化される
  if (opts_.opt_level() >= OptimizerOptions::L1) {
    opts_.set_do_common_subexpression_elimination(true);
    opts_.set_do_constant_folding(true);
  }
}

一方Function Inliningに関しては、opt_level の値に関わらず do_function_inlining の設定によって、有効・無効が決まります。デフォルトでは、Function Inliningは無効化されています。

例えば以下は、GraphOptimizerによる全ての最適化をPythonプログラムから無効化するためのソースコードです。

import tensorflow as tf
from tensorflow.core.protobuf import config_pb2

# GraphOptimizerによる全ての最適化を無効化
cfg = config_pb2.ConfigProto()
cfg.graph_options.optimizer_options.opt_level = config_pb2.OptimizerOptions.L0

# 計算グラフ構築
# ...

# 計算グラフ実行
with tf.Session(config=cfg) as sess:
  sess.run(...)

デフォルトで無効化されているFunction Inliningを有効化する場合は、以下のように明示的に有効化する必要があります。

import tensorflow as tf
from tensorflow.core.protobuf import config_pb2

# Function Inliningを有効化
cfg = config_pb2.ConfigProto()
cfg.graph_options.optimizer_options.do_function_inlining = True

# 計算グラフ構築
# ...

# 計算グラフ実行
with tf.Session(config=cfg) as sess:
  sess.run(...)

TensorFlow内部構造解析 (4.7) 計算グラフ最適化処理3 GraphOptimizer