はじめに
chainerでforwardの計算をするときに行列の対角成分をとりたい時があったが、これを計算するfunctionsが現行のchainer(v1.19.0)に確認できなかった。
ちなみにTensorflowだとdiag_part()がこれに相当する。
http://devdocs.io/tensorflow~python/math_ops#diag_part
Theanoだとdiagonal()という関数がある。がこれは成分でなく対角成分のみでつくった対角行列を返す?(試してない)
github.com
やりたい処理をnumpyで表現すると、具体的には以下のようになる。
>>> import numpy as np >>> x=np.arange(9).reshape(3,3) >>> x array([[0, 1, 2], [3, 4, 5], [6, 7, 8]]) >>> np.diag(x) array([0, 4, 8])
対処
自作で対応する関数を作ってもいいのだが、暇もないので今あるfunctionsで作ることにした。
作ったのは以下の通り。
github.com
案は2つあって、
1. 単位行列を要素ごとにかけて列ごとの和をとる
2. batch_matmulを使って正方行列(nxn)と単位行列(nxn)をベクトル(nx1)ごとに分割して計算する
実験
どちらが速いのか実際やってみた。
cpu計算とgpu計算で正方行列のサイズは100x100,1000x1000,10000x10000。
GPUは NVIDIA Corporation GM206 [GeForce GTX 960] を用いた。
100x100 matrix #cpu $ python chainer_diag.py 100 -1 time of diag_part_elementwise: 0.0002493858337402344 time of diag_part_batch_matmul: 0.00036978721618652344 #gpu $ python chainer_diag.py 100 0 time of diag_part_elementwise: 0.10309004783630371 time of diag_part_batch_matmul: 0.06094169616699219 1000x1000 matrix #cpu $ python chainer_diag.py 1000 -1 time of diag_part_elementwise: 0.0017409324645996094 time of diag_part_batch_matmul: 0.0019404888153076172 #gpu $ python chainer_diag.py 1000 0 time of diag_part_elementwise: 0.10210251808166504 time of diag_part_batch_matmul: 0.06030750274658203 10000x10000 matrix #cpu $ python chainer_diag.py 10000 -1 time of diag_part_elementwise: 0.11125564575195312 time of diag_part_batch_matmul: 0.0674281120300293 #gpu $ python chainer_diag.py 10000 0 time of diag_part_elementwise: 0.12932491302490234 time of diag_part_batch_matmul: 0.1076195240020752
考察
gpu計算よりもcpu計算のが速い・・・
個別に比較すると、大体の場合要素ごとの計算よりもbatch_matmulを使った方が若干速いらしい。
ということで、batch_matmulを使う方法で実装しようと思う。