chainerで正方行列の対角成分を抜き出す

はじめに

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。 GPUNVIDIA 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を使う方法で実装しようと思う。