Chainerでのスパース行列の取り扱い

Chainer v5からスパース行列の演算ができるようになるそうです。まだ、行列積(sparse_matmul)だけですが、今後使える関数が増えるとありがたいです。早速使用してGraph Convolutionを実装してみました(Github)。この記事は実装するときのメモに修正を加えたものです。

Chainer v5のインストール

まだテスト版だそうですが、pipで入れることができます。
追記：正式にバージョンアップしていました。

pip install -U chainer

スパース行列について

まず、スパース行列(sparse matrix, 疎行列)について簡単に説明しておきます。スパース行列は行列の要素がほとんどゼロで、非ゼロ要素が少しあるというような行列です。

>>> a = np.array([[0, 2, 0], [-1, 0, 0]], np.float32)
>>> a
array([[ 0.,  2.,  0.],
       [-1.,  0.,  0.]], dtype=float32)

こんなやつですね。要は要素がスカスカな行列なのでsparseと呼ばれるわけです。ほとんどゼロなので、そのままの行列で取り扱うとメモリを圧迫します。なので、別の形式で保存して取り扱うわけです。今回、Chainerに実装されたのは**COOrdinate format (COO format)**と呼ばれる形式です。簡単に言えば、非ゼロ要素のインデックスをx, yで記録しておくというものです。記録されてないやつ以外はゼロで埋めます。

scipy.sparse

Pythonでスパース行列を扱うにはscipy.sparseを使うのが一般的です。ここからはscipy.sparseを用いてみます。COO形式の行列を作成するには、非ゼロ要素の配列、非ゼロ要素のx座標、非ゼロ要素のy座標、配列全体の形状の4つが必要です。

>>> from scipy.sparse import coo_matrix
>>> row  = np.array([0, 3, 1, 0])
>>> col  = np.array([0, 3, 1, 2])
>>> data = np.array([4, 5, 7, 9])
>>> coo_matrix((data, (row, col)), shape=(4, 4)).toarray()
array([[4, 0, 9, 0],
       [0, 7, 0, 0],
       [0, 0, 0, 0],
       [0, 0, 0, 5]])

という感じで作成ができます。
参考：scipy.sparse.coo_matrix — SciPy v1.1.0 Reference Guide

スパース行列に関わるutils, function

参考：Sparse utilities — Chainer 5.0.0rc1 documentation

chainer.utils.to_coo

numpyやcupyの配列をCOO形式に変換できます。

>>> import chainer
>>> import numpy as np
>>> data = np.array([[0, 2, 0], [-1, 0, 0]], np.float32)
>>> x = chainer.utils.to_coo(data)
>>> x.data
variable([ 2., -1.])
>>> x.row
array([0, 1], dtype=int32)
>>> x.col
array([1, 0], dtype=int32)
>>> x.shape
(2, 3)

chainer.utils.CooMatrix

COO形式のスパース行列を作成できます。

>>> data = np.array([[0, 2, 0], [-1, 0, 0]], np.float32)
>>> x = chainer.utils.to_coo(data)
>>> x
<chainer.utils.sparse.CooMatrix object at 0x0000020932F76828>
>>> y = chainer.utils.CooMatrix(x.data.data, x.row, x.col, x.shape)
>>> y
<chainer.utils.sparse.CooMatrix object at 0x0000020932F850F0>
>>> y.data
variable([ 2., -1.])
>>> y.row
array([0, 1])
>>> y.col
array([1, 0])
>>> y.shape
(2, 3)

chainer.functions.sparse_matmul

片方がスパース行列でも行列積が計算できる関数です。

>>> import chainer.functions as F
>>> a = np.array([[0, 2, 0], [-1, 0, 0]], np.float32)
>>> b = np.arange(6.).reshape(3,2)
>>> a_coo = chainer.utils.to_coo(a)
>>> c = F.sparse_matmul(a_coo, b)
>>> c
variable([[ 4.,  6.],
          [ 0., -1.]])
>>> a @ b
array([[ 4.,  6.],
       [ 0., -1.]])

帰ってくる行列は密行列です。

scipy.sparseからchainer.Variableへの変換

chainer.utils.CooMatrixを使って変換ができます。前処理はscipyでして、学習時に変換という使い方ができます。

def sparse_mx_to_chainer_sparse_variable(sparse_mx):
    """Convert a scipy sparse matrix to a chainer sparse variable."""
    sparse_mx = sparse_mx.tocoo().astype(np.float32)
    data = sparse_mx.data
    row = sparse_mx.row
    col = sparse_mx.col
    shape = sparse_mx.shape
    return chainer.utils.CooMatrix(data, row, col, shape)