TensorFlow/Kerasでグラム行列（テンソル）を計算する方法
TensorFlowで分散や共分散が絡む演算を定義していると、グラム行列を計算する必要が出てくることがあります。行列はまだよくてもテンソルのグラム行列はどう計算するでしょうか？今回はテンソルの共分散計算に行く前に、その前提のテンソルのグラム行列の計算から見ていきます。
目次
グラム行列とは
名前は仰々しくてもやっていることは簡単です。$X$という行列があったときに、
$$X^T \cdot X $$
を計算することです。ここで$X^T$は行列$X$の転置行列、真ん中のドットは内積を表します。これがなぜ重要かと言うと、分散共分散行列（≒相関係数の行列）の計算に使うからです。いきなり理論的な話をするとこんがらがるので、とりあえずそれで使うんだなというこということを頭の片隅に入れておいてください。
Numpyでグラム行列
線形代数の授業ではないので、数式よりもプログラムで見たほうが早いですよね。まずは簡単な2×2行列のグラム行列を計算しましょう。
import numpy as np
matrix = np.arange(4).reshape(2,2)
gram = np.dot(matrix.T, matrix)
print("M.T : ")
print(matrix.T)
print("M : ")
print(matrix)
print("gram : ")
print(gram)
これは0～3の値からなる2×2行列のグラム行列の計算です。定義どおりに、「転置した行列と、元の行列」の内積で表されます。これなら手計算でも計算できますね。
M.T :
[[0 2]
[1 3]]
M :
[[0 1]
[2 3]]
gram :
[[ 4 6]
[ 6 10]]
せっかくなので手計算でも確認してみましょう。gramの左上の部分は0×0+2×2=「4」、右上の部分は0×1+2×3=「6」、左下の部分は1×0+2×3=「6」、1×1+3×3=「10」。結果が正しいのが確認できましたね。
通常、行列の内積の順序を入れ替えることはできませんが、グラム行列の場合は同じを行列を転置したかどうかで内積を取っているので、文脈やコードによっては$X\cdot X^T$で書いていることもあります。この書き方が間違いというわけではありません。もちろん内積であることには変わりないので、結果の次元は一緒になることは保証できません。例えば元が2×3行列なら、転置を先にかければ結果は3×3行列になりますし、転置を後からかければ2×2行列になります。数学の定義に厳密になるよりも、状況に応じて適切な順序でかけるのが重要だと思います。
脱線してしまいましたが、Numpyでは定義通りにnp.dot()すればグラム行列を計算できることが確認できました。もう手計算をする必要はありません。
TensorFlowのグラム行列
次にTensorFlowでのグラム行列の計算例をみましょう。TensorFlowというと小難しそうな印象がありますが、グラム行列は実はNumpyよりも簡単にかけます。
import numpy as np
import keras.backend as K
import tensorflow as tf
matrix = K.variable(np.arange(4).reshape(2,2))
gram = tf.matmul(matrix, matrix, transpose_a=True)
print(K.get_value(gram))
セッションの定義が面倒くさいのでKerasのバックエンド関数使っています。Qiitaでも書きましたが、変数だけKerasで定義して、実際の演算はTensorFlowの関数を使うなんて書き方もできます。
TensorFlowではtf.matmulという関数でだいたい内積ができるのですが、この関数がウルトラ高性能で、「transpose_a=True」と指定してくれると勝手に転置までやってくれるのですよね。ちなみに逆側を転置したいときはtranspose_b=Trueとします。これだけだったら「ふーん、なんか変な引数ついてるな」って感じなのですが、この関数の真価は3階以上のテンソルになったときです。それは後で確認します。
[[ 4. 6.]
[ 6. 10.]]
結果はこの通り。ちなみに2階（行列）だったら、Kerasのバックエンドのdotでも計算できますね。
matrix = K.variable(np.arange(4).reshape(2,2))
gram = K.dot(K.transpose(matrix), matrix)
print(K.get_value(gram))
ただ、K.dot()の場合は3階以上になったときにtf.matmul()と挙動が変わります。そこだけ注意してください。
混乱してきたかもしれないので一度整理しましょう。今内積の計算に、Numpyの関数「np.dot()」と、TensorFlowの関数「tf.matmul()」、Kerasのバックエンド関数「K.dot()」の3種類が登場しました。これらの違いを少し細くしておきます。
まず、NumpyとTensorFlow/Kerasの関数は明らかにものが違います。Numpy関数はただ値をパックしただけのNumpy配列を計算しているにすぎませんが、TensorFlow/Kerasはどちらもニューラルネットワーク用に特別に定義されたテンソルを計算しています。ここでいうテンソルとは、線形代数で言われるベクトル→行列→テンソルと階数を上げていくものだけではなく、ニューラルネットワークの訓練に必要な微分を自動でやってくれるよう、計算グラフの連結によって定義された特別な変数です。つまり、TensorFlowのテンソルと言われたときは、線形代数のテンソルよりはかなり狭義（特殊）な意味で使われることが多いです。この投稿もそうですが、広義のテンソルとTensorFlowのテンソルを、どちらも特に断りもなくまぜこぜに「テンソル」と読ぶことが多いので、文脈に注意してください。
次はTensorFlowとKerasの関数の違いですが大まかに言ってこれはかなり似ています。まずフレームワークとして見たときにKerasは上位のAPI、もう少し砕けて言えば「Kerasは入れ物で、TensorFlowはその入れ物の中身（バックエンド）」という位置づけです。Kerasは異なるバックエンドを同一のAPIで表現できるのが大きな売りです。
KerasはバックエンドにTensorFlowやTheano、MXNetなどを指定でき、一番よく使われるのがTensorFlowバックエンドでしょうが、TensorFlowバックエンドの場合は、TensorFlowのテンソルとKerasのテンソルは（ほぼ）同一となります。したがって、先程確認したようにKerasの関数で変数を定義し、TensorFlowの関数に代入するようなことができます。
ではTensorFlowとKerasの関数が完全同一かというとこれは違います。少なくとも内積計算は違います。だいたいはKerasのバックエンド関数はTensorFlowの簡単なラップで同一の結果を返すものが多いものの（例えば絶対値のK.abs()ならtf.abs()を返しているだけ）、一部の関数は他のバックエンドとの整合性を取るためか、少し特殊なラップ実装をしています。詳しくはgithubのKerasのソースコードを確認すると良いでしょう。全てオープンになっています。内積計算はTensorFlowとKerasで関数が異なる1つで、3階以上のテンソルになったときにK.dot()とtf.matmul()の挙動は明確に異なります。それを確認していきましょう。
3階以上のグラム行列
いよいよ本題のテンソルのグラム行列です。急にマニアックな話になるので頑張ってください。
さて、2次元のときのグラム行列の次数をおさらいしましょう。次の問題を考えます。
(問) Xが(12, 3)行列のときグラム行列の次数はどのようになるか？
答えは、(3, 3)行列です。わかりましたでしょうか？TensorFlowでもNumpyでもいいですが、一例としてNumpyで確認してみます。
import numpy as np
matrix = np.zeros((12, 3))
gram = np.dot(matrix.T, matrix)
print(gram.shape)
# (3, 3)
KerasでもTensorFlowでもそうですが、テンソル操作をしたいときは次元（Numpyでいうところのshape）を考えるのが理解への近道です。「入力のshapeがどうで、欲しい出力のshapeはどうなのか」を考え、具体的な値はnp.zeros()でもいいので代入してみると、自ずと正解のコードが導けます。テンソル操作は相当ブラックボックス化しがちなので、「あれ？」と思ったらshapeに適当な値入れてこのように簡単なコードを実行してみてください。その細かなチェックが最終的なデバッグのコストを小さくすると自分は思います。
ちなみにテンソルの次元を確認したいときは、Kerasのバックエンド関数のK.int_shape()、Numpyと同様にTensorFlowのテンソルの末尾に「.shape」をつける、TensorFlow関数のtf.shape()といった関数を活用しましょう。tf.shape()は少し特殊で、バッチサイズのように(None, )で渡される値を実際の値で返してくれます。乱数の発生などで実際の値が必要になったときに使います。確認したいだけならだいたいK.int_shape()で事足ります。
では、次の問題。これはちょっと難しいです。
（問）今データが(64, 3, 16)という次元の3階テンソルで流れているものとする。グラム行列が(64, 16, 16)となるように2～3次元目で内積を計算するにはどのように定義したらよいか？
一番簡単なのはforループですが、forループを使うとハードウェアレベルの並列化を活かせなくなってしまうのでここは一気にテンソル（行列）計算でかたをつけたいです。np.einsumを使うという手もありますが、自分が頭悪いのでこれがろくに理解できません。なんとなくそれっぽいのはできましたが解説は適当です。ご容赦ください。
np.einsum()の場合
おそらくこれが一番難しいやり方。いきなり3階テンソルだと厳しいので行列の場合からやりましょう。数十分ぐらい添字を変えて適当に試してたらできました。
import numpy as np
matrix = np.arange(36).reshape(12,3)
# einsumの場合
gram = np.einsum("ij,ik->jk", matrix, matrix)
print(gram)
#[[4554 4752 4950]
# [4752 4962 5172]
# [4950 5172 5394]]
# 確認用
print(np.dot(matrix.T, matrix))
#[[4554 4752 4950]
# [4752 4962 5172]
# [4950 5172 5394]]
慣れればきっと簡単なんでしょうが、自分はこの表記ろくに使ったことないので正直めんどいです。多分、消したい次元の添字は共通で、残したい次元の添字を列挙するのかなという漠然とした理解です。
3次元の場合のeinsumです。なんか釈然としませんがこれで計算できます。
import numpy as np
a = np.arange(64*3*16).reshape(64,3,16)
# einsumの場合
gram = np.einsum("ijk,ijl->ikl", a, a)
print(gram.shape)
# 確認用
kakunin = np.matmul(np.swapaxes(a, 1, 2), a)
# 全て等しいか
print(all(np.ravel(gram == kakunin)))
#(64, 16, 16)
#True
64で動かさない軸はiで固定、あとは2次元のときと同様です。なんとなくわかってきたようなわからないような。
今しれっと使っていますが、このグラム行列の計算はnp.matmul関数を使うとちょっと簡単にできます。matmulとeinsumの結果が等しいことを確かめています。
np.matmul関数
あまりNumpyのことを重点的にやりたくないのですが、せっかく出してしまったのでやります。こちらは先程見たTensorFlowのtf.matmulと若干似ている関数です。若干というのは軸のスワップの仕方がTensorFlowとNumpyでは違います。
a = np.arange(64*3*16).reshape(64,3,16)
gram = np.matmul(np.swapaxes(a, 1, 2), a)
print(gram.shape)
# (64, 16, 16)
ドキュメントに書いていますが、3階以上のテンソルの場合は、最後の2つの次元を切り出して行列積を計算しています。今回の場合はたまたまこの関数が使えた例です。
ちなみにいくらnp.matmulを使えば簡単とはいえ、行列のときと同じように、aの転置をそのまま左側に代入するのはダメです。
# ダメな例
a = np.arange(64*3*16).reshape(64,3,16)
print(a.T.shape)
gram = np.matmul(a.T, a)
#(16, 3, 64)
#Traceback (most recent call last):
# File…（以下略）
# gram = np.matmul(a.T, a)
#ValueError: shapes (16,3,64) and (64,3,16) not aligned: 64 (dim 2) != 3 (dim 1)
Numpyのテンソルの反転は、次元がそのまま反転してしまうので、(64, 3, 16)となります。これとの行列積は計算しても意味がありません。
Tensorflowのtf.matmulを使おう
実はこれが最も簡単な答えです。先程、tf.matmul()は高性能だよと言いましたが、その理由はここにあります。
import numpy as np
import tensorflow as tf
import keras.backend as K
a = np.arange(64*3*16).reshape(64,3,16)
ta = K.variable(a)
tgram = tf.matmul(ta, ta, transpose_a=True)
gram = K.get_value(tgram)
print(gram.shape)
# (64, 16, 16)
「えっ！？行列のときと同じコードでええの？？」、そうです。何次元だろうがこれでグラム行列が出てきます。だから高性能なのです。人間ではもう何がなんだかわからなくなるような6次元テンソルだってグラム行列計算してくれます。
a = np.zeros((64, 32, 16, 8, 4, 2))
ta = K.variable(a)
tgram = tf.matmul(ta, ta, transpose_a=True)
gram = K.get_value(tgram)
print(gram.shape)
# (64, 32, 16, 8, 2, 2)
ね、簡単でしょ？
ちなみにKerasのdot関数は3次元から既にダメです。
import numpy as np
import tensorflow as tf
import keras.backend as K
a = K.variable(np.zeros((64, 3, 16)))
at = tf.transpose(a, perm=[0,2,1])
print("at = ", K.int_shape(at))
tgram = K.dot(at, a)
print(K.int_shape(tgram))
#at = (64, 16, 3)
#(64, 16, 64, 16)
Transposeまではうまくいっているのに、内積計算が(64, 16, 64, 16)とおかしなことになるんですよね。素直にTensorFlowの関数を使うべきでしょう。
計算結果の確認
最後にこれまでの3つの関数、np.einsum(), np.matmul(), tf.matmul()がそれぞれ同じ結果を返すかを確かます。4階テンソルのグラム行列を求めてみましょう。
import numpy as np
import tensorflow as tf
import keras.backend as K
a = (np.random.rand(64, 3, 5, 16) * 100.0).astype(np.int32)
# numpyの場合(ein_sum)
numpy_ein = np.einsum("ijkl,ijkm->ijlm", a, a)
# numpyの場合(matmul)
numpy_matmul = np.matmul(np.swapaxes(a, 2, 3), a)
# TensorFlowの場合
ta = K.variable(a)
tf_gram = tf.matmul(ta, ta, transpose_a=True)
tf_matmul = K.get_value(tf_gram)
def is_all_same(x, y):
return all(np.ravel(x == y))
print("np_ein == np_matmul :", is_all_same(numpy_ein, numpy_matmul))
print("np_matmul == tf_matmul", is_all_same(numpy_matmul, tf_matmul))
元のテンソルを乱数で発生させています。グラム行列の計算結果の比較なので、元の値によらずに常に全てが等しくなるはずです（整数値のテンソルなので誤差は気になるほどではないと思います）。結果は以下の通り。is_all_same()関数は文字通り、全ての要素が一致しないとTrueを返しません。
np_ein == np_matmul : True
np_matmul == tf_matmul True
OKですね。これでテンソルのグラム行列の3通りの計算方法をマスターできました。
まとめ
かなりハードな内容になってしまいましたがまとめます。
- グラム行列は分散や共分散がからむ計算をする際に使うことがある
- グラム行列の計算は行列までなら内積を取ればいいが、3階以上のテンソルになるとNumpyならeinsumやmatmul、TensorFlowならmatmulで計算する必要がある。Kerasのdotは3階以上のグラム行列の計算に向かない。
- TensorFlowのtf.matmul()が、行列のときもテンソルのときも同じコードで書けるからグラム行列の計算には楽
速度は未検証ですが、tf.matmul()がこんなに便利だとは思いませんでした。次の記事ではここでさんざん頑張って求めたグラム行列を実際どう使うのかというのを書きます。
次の記事
[統計学や機械学習で使われる分散共分散行列、相関行列とグラム行列の関係]
(https://blog.shikoan.com/cov-corr-gram-matrix)
Shikoan's ML Blogの中の人が運営しているサークル「じゅ～しぃ～すくりぷと」の本のご案内
技術書コーナー
北海道の駅巡りコーナー