TensorFlow で知っていると役に立つ（かもしれない）演算系関数たち

この記事は、TensorFlow Advent Calendar 2016 の13日目です。

TensorFlow で処理をスクラッチする際に知っておくと便利な関数をご紹介したいと思います。

以降の説明は、TensorFlow v0.11.0 の動作に基づいて説明しています。挙動や名称がバージョンによって変化する場合もありますので、ご注意ください。

基本的な演算子とブロードキャスティング

まずは、基本的な演算処理である四則演算の挙動と matmul についてご紹介します。

TensorFlow でも基本的な演算子といえば、+, -, *, / などのことですが、これらの演算子は、こちらに記載のように、基本的にそれぞれの要素ごとに該当する演算を行った（* の場合は、要素ごとに * を適用した）結果を、出力するテンソルの要素とします。

a = tf.Variable([[1, 2],
                 [3, 4]])
y = a * [[1, 2], [3, 4]]

[[ 1  4]
 [ 9 16]]

一方で、これらの演算子には、numpy 等でも採用されているブロードキャスティングと呼ばれる動作も適用されます。一番簡単な例は、下記のようにテンソルとスカラーを演算する場合です。この場合、テンソルの各要素にスカラーとの演算を適用した結果が、出力されるテンソルの要素となります。

a = tf.Variable([[1, 2],
                 [3, 4]])
y = a * 2

[[2 4]
 [6 8]]

ブロードキャスティングは、ベクトルなどの元のテンソルと比べて、あるランクの次元が1になっているものと演算を行った際にも適用されます。例えば、2x2 行列に対しては、1x2 の行列や、2x1 の行列との演算によりブロードキャスティングが適用されます。どのような結果になるかは、下記を見ていただくのが早いかと思います。

2x2 行列 * 1x2 行列

a = tf.Variable([[1, 2],
                 [3, 4]])
y = a * [2, 4]

[[ 2  8]
 [ 6 16]]

2x2 行列 * 2x1 行列

a = tf.Variable([[1, 2],
                 [3, 4]])
y = a * [[2], 
         [4]]

[[ 2  4]
 [12 16]]

行列の掛け算（matmul）

基本的な演算子の次によく使われるものといえば、一般的な行列の掛け算である matmul 演算でしょう。matmul 演算は非常に簡単で、LxN 行列を第一引数に、MxN 行列を第二引数に渡すと、LxM 行列を返してくれます。具体的には、下記の結果を見ていただくのが早いでしょう。

2x3 行列と3x1 行列の掛け算

a = tf.Variable([[1, 2, 3],
                 [3, 4, 5]])
y = tf.matmul(a, [[1],
                  [2],
                  [3]])

[[14]
 [26]]

2x3 行列と3x2 行列の掛け算

a = tf.Variable([[1, 2, 3],
                 [3, 4, 5]])
y = tf.matmul(a, [[1, 2],
                  [2, 3],
                  [3, 4]])

[[14 20]
 [26 38]]

パディングを入れる（tf.pad）

さて以降は、少し変わり種の関数をご紹介していきたいと思います。まずトップバッターは、tf.pad です。

データ長を揃えたりなど、パディングを入れたいという場面は、ままあると思いますが、tf.pad を用いることで簡単にその処理を行うことができます。パディングの指定は、Nx2 行列で行います。階ごとに [前に挿入する分、後に挿入する分] となるようにスカラー値を指定します。

a = tf.Variable([1, 2])
y = tf.pad(a, [[1, 2]])

[0 1 2 0 0]

バッチごとの行列演算（tf.batch_matmul）

行列演算を効率よく行うために用意されている処理が、tf.batch_matmul です。この関数は、名前の通りバッチごとに matmul 演算を行うことができます。正確には、3階以上のテンソルを行列の集合と考え、行列毎に matmul 演算を行いその結果を集約した結果を返します。

a = tf.Variable([[[1, 2],
                  [3, 4]],

                 [[1, 2],
                  [3, 4]]])
b = tf.Variable([[[1, 1],
                  [1, 1]],

                 [[2, 2],
                  [2, 2]]])
y = tf.batch_matmul(a, b)

[[[ 3  3]
  [ 7  7]]

 [[ 6  6]
  [14 14]]]

セグメント毎に値を集約する（Segmentation)

TensorFlow には、Segmentation を行う関数群が用意されています。この関数は、指定したIDに沿ってインプットしたテンソルを集約することができます。IDは0から始まり自由に指定することができますが、同じIDは連続して指定する必要があるという制約があります。

下記には、tf.segmentation_sum の例を上げていますが、乗算や平均の計算も行うことができます。基本的には、名前の通りセグメント毎にデータを集約するものですが、unsorted_segment_sum のように unsorted な実装が増えれば、ラベルごとの平均を求めたい時にも使えそうな関数です。

a = tf.Variable([[1, 2],
                 [2, 3],
                 [3, 4],
                 [4, 5]])
b = tf.Variable([0, 0, 1, 2])
y = tf.segment_sum(a, b,)

[[3 5]
 [3 4]
 [4 5]]

累積演算（Scan）

Scan は累積的に値を変化させる事ができる演算で、tf.cumsum と tf.cumprod が用意されています。i番目の出力は、元のテンソルのi番目までの要素の総和もしくは総乗になります。等差数列や等比数列を作りたい時や、重みを逓増させたり逓減させたい時に使えそうな機能となっています。

a = tf.Variable([1, 2, 3, 4])
y = tf.cumsum(a)

[ 1  3  6 10]

いかがだったでしょうか。TensorFlow には、この他にもこれ何に使うの？といった処理が、公式に用意されているケースが結構あります。自力実装しなければならない場合もありますが、計算速度のことも用意されている関数でなんとかする方が基本的には良いです。バージョンアップ時に API 一覧を眺めてみるだけでも色々と発見があって面白いので、ぜひ皆さんも自分だけの便利関数を探してみてください。

終末 A.I.

データいじりや機械学習するエンジニアのブログ