TensorFlowでループを使う

TensorFlowの計算グラフ内で、for/whileのようなループを使いたいときがごくまれにあります。

• その時に使える、tf.while_loopを紹介します
• 実際使用したコード

今回は、逆強化学習のアルゴリズム(線形計画法)を実装する際に逆行列の演算が個別に必要な計算箇所があり、これは一度にはできないためループで処理することにしました。

Tips

実際使えるのはWhileで、条件式と条件式がTrueの間実行される関数の2つを定義し、使用する変数を入力することで発動させます。

i = tf.constant(0)
c = lambda i: tf.less(i, 10)
b = lambda i: tf.add(i, 1)
r = tf.while_loop(c, b, [i])

• 複数の引数を取ることもできる
• bodyについて、引数の型と変数の型はあっている必要がある。上記の場合だと、iはTensor/int32なので、bodyの帰り値もTensor/int32の必要がある
• このため、配列をスタックしていくような操作ができない。引数が空の配列で、bodyの中でappendすると、引数は「空の配列」なのにreturnの型が「要素が入っている状態」になるのでエラーになる。
• Eagerを使えば？というのは一理ある。body関数内のエラーは捕捉できないため(while内でエラーがあるということだけ教えてくれる)、デバッグが大変