2017年1月24日

コサイン類似度に基づくソート処理の実装方法とその性能比較

文書の類似度を計算する方法に「コサイン類似度」を用いる方法があります。

これは、出現する単語を出現回数などで数値化して、空間ベクトルに変換した上でベクトル同士の類似度を計算する、という手法です。
最近、このコサイン類似度を使って、似ているデータを検索するWebアプリを試しに作っていたのですが、ふと、

「このコサイン類似度を使ったソート処理をPostgreSQLでどのように実装するともっとも高速な実装になるのだろうか。また、現実的なパフォーマンスを考えた時にデータ量や次元のサイズはどこまで増やせるのだろうか」

ということが気になりました。

PostgreSQLは、その拡張性の高さがウリの一つですが、そのため「UDFを作る」ということを考えても、実装方法にもいろいろあります。

今回は、PostgreSQL内部でデータ処理を実装するに当たって、どのような実装方法があるのか、それぞれどのように性能が違うのか、そしてその時にデータサイズがどのように影響するのかを見てみます。

■前提条件


今回は以下の前提条件で実装と性能比較を行います。
  • ソート処理するデータはPostgreSQLに蓄積されているものを対象とする
  • 空間ベクトルを表すデータは、PostgreSQL の float8 の配列で1カラムに保持する
  • コサイン類似度による類似度を計算し、もっとも類似度の高いレコードをN件取得する
これらを前提条件として、コサイン類似度の計算を
  • (1) CのUDFで実装した場合
  • (2) PythonのUDFで実装した場合(PL/Python)
  • (3) scikit-learnのcosine_similarity関数を使ってPythonのUDFで実装した場合(PL/Python + scikit-learn)
  • (4) MADlibcosine_similarity関数を使った場合
  • (5) scikit-learnを使ってクライアント側に取得してクライアント側で計算する場合
に、パフォーマンスがそれぞれどう異なるかを確認してみます。

また、レコード数および空間ベクトルの次元数によって実行時間がどのように変わるかも確認してみます。

■処理対象のデータ


処理対象のデータは、 data_vec テーブルに主キーとして整数型のカラム id を持ち、空間ベクトルのデータとして float8[] 型の カラム vec を持つテーブルです。
CREATE TABLE data_vec (
  id INTEGER PRIMARY KEY,
  vec FLOAT8[] NOT NULL
);

このテーブルの中に、以下のようにレコードが入っています。(以下は100次元の空間ベクトルのデータが5件入っている状態)
test_cos_sim=> select * from data_vec ;
 id |                                                                                                    vec

----+-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
  0 | {1,1,0,0,1,0,1,0,0,1,1,1,0,1,1,0,1,1,1,1,0,1,1,1,0,0,0,1,1,1,0,1,0,1,1,0,1,0,1,1,0,0,1,0,0,1,0,1,0,1,1,0,0,0,1,1,0,1,1,1,1,1,1,1,1,0,1,0,1,0,0,0,1,1,0,0,1,1,1,1,1,1,1,0,1,0,1,1,1,1,1,1,0,1,1,1,1,0,1,0}
  1 | {1,1,0,1,0,0,1,0,1,0,0,1,0,1,1,1,1,0,1,1,1,0,0,1,0,0,1,0,0,0,1,1,0,0,0,0,0,1,0,0,1,0,0,1,0,0,1,1,1,0,1,1,0,1,1,0,0,1,0,0,0,0,1,1,0,0,1,0,1,0,1,0,0,0,0,1,0,0,0,1,1,1,1,0,0,1,0,0,0,0,1,0,1,1,0,0,1,0,1,0}
  2 | {0,0,0,1,0,1,1,1,0,0,0,1,1,1,0,1,1,0,1,1,0,1,0,0,1,0,1,1,0,0,1,1,0,1,0,1,0,0,0,0,1,1,0,1,1,0,1,0,0,0,1,1,1,1,1,0,1,1,1,1,1,1,0,1,0,1,0,0,0,1,1,0,1,0,0,1,0,0,1,0,0,0,0,0,1,0,0,0,0,0,1,0,1,1,0,0,1,0,0,1}
  3 | {1,1,1,0,1,1,0,0,1,0,0,0,1,0,0,1,1,0,0,1,0,1,0,0,0,1,1,1,0,1,0,0,0,0,1,0,0,0,0,1,0,1,1,0,0,1,1,0,1,0,0,0,1,0,1,0,0,1,1,1,0,0,0,1,1,1,1,1,0,0,1,0,0,0,0,1,0,0,0,1,1,1,1,1,1,1,1,1,1,1,1,0,1,1,1,0,1,0,1,1}
  4 | {1,1,0,0,1,0,0,1,1,0,0,1,0,1,0,0,1,0,0,0,0,1,1,1,0,1,0,1,0,0,0,0,0,1,1,0,1,0,0,1,0,1,1,0,1,0,0,1,1,1,0,1,1,1,1,1,1,1,1,0,1,0,1,0,1,0,0,1,0,1,1,1,1,0,0,1,1,1,1,0,0,0,0,1,1,1,0,0,1,1,0,0,1,0,1,1,1,0,0,0}
(5 rows)

test_cos_sim=>

■コサイン類似度計算UDFの実装


コサイン類似度の計算をするUDFのそれぞれの実装は以下のようになります。

CによるUDFの実装 (1) は以下のようになります。
Datum
cosine_similarity(PG_FUNCTION_ARGS)
{
  ArrayType  *vec_a = PG_GETARG_ARRAYTYPE_P(0);
  ArrayType  *vec_b = PG_GETARG_ARRAYTYPE_P(1);
  float8 ab = 0;
  float8 aa = 0;
  float8 bb = 0;
  int i;

  int len_a = (ARR_SIZE(vec_a) - ARR_DATA_OFFSET(vec_a)) / sizeof(float8);
  int len_b = (ARR_SIZE(vec_b) - ARR_DATA_OFFSET(vec_b)) / sizeof(float8);

  float8 *a = (float8 *)ARR_DATA_PTR(vec_a);
  float8 *b = (float8 *)ARR_DATA_PTR(vec_b);

  for (i = 0; i < len_a; i++)
  {
    aa += a[i] * a[i];
    bb += b[i] * b[i];
    ab += a[i] * b[i];
  }

  PG_RETURN_FLOAT8(ab/(sqrt(aa)*sqrt(bb)));
}

PL/PythonのUDFとして実装 (2) すると以下のようになります。
CREATE OR REPLACE FUNCTION cosine_similarity_plpy(vec_a float8[], vec_b float8[])
  RETURNS float8
AS $$
    from math import sqrt
    aa = 0
    bb = 0
    ab = 0
    for a,b in zip(vec_a, vec_b):
        aa += a * a
        bb += b * b
        ab += a * b
    return ab/(sqrt(aa)*sqrt(bb))
$$
LANGUAGE 'plpython2u';

PL/Pythonでも scikit-learn を使った実装 (3) は以下のようになります。
CREATE OR REPLACE FUNCTION cosine_similarity_sk(vec_a float8[], vec_b float8[])
  RETURNS float8
AS $$
    from sklearn.metrics.pairwise import cosine_similarity
    return cosine_similarity([vec_a], [vec_b])[0][0]
$$
LANGUAGE 'plpython2u';

■検証用クエリ


今回の評価では、上記のデータとUDFを使って、「id=0」の空間ベクトルを問い合わせクエリとして、「その他(id=0以外)」のレコードの中から最も似ているレコードを10件取得してみます。クエリは、それぞれ以下のようになります。

CによるUDFを使う場合 (1):
SELECT
  a.id as "id1",
  b.id as "id2",
  cosine_similarity(a.vec, b.vec)
FROM
  data_vec a,
  data_vec b
WHERE
  a.id = 0
AND
  b.id <> 0
ORDER BY
  3 DESC
LIMIT 10;

PL/PythonによるUDFを使う場合 (2):
SELECT
  a.id as "id1",
  b.id as "id2",
  cosine_similarity_plpy(a.vec, b.vec)
FROM
  data_vec a,
  data_vec b
WHERE
  a.id = 0
AND
  b.id <> 0
ORDER BY
  3 DESC
LIMIT 10;

PL/Python + scikit-learn によるUDFを使う場合 (3):
SELECT
  a.id as "id1",
  b.id as "id2",
  cosine_similarity_sk(a.vec, b.vec)
FROM
  data_vec a,
  data_vec b
WHERE
  a.id = 0
AND
  b.id <> 0
ORDER BY
  3 DESC
LIMIT 10;

MADlib の cosine_similarity 関数を使う場合 (4):
SELECT
  a.id as "id1",
  b.id as "id2",
  madlib.cosine_similarity(a.vec, b.vec)
FROM
  data_vec a,
  data_vec b
WHERE
  a.id = 0
AND
  b.id <> 0
ORDER BY
  3 DESC
LIMIT 10;

上記は関数名が違うだけで、クエリは基本的には同じです。

なお、クライアント側に取得して scikit-learn で計算する実装 (5) は以下のようになります。
import psycopg2
from sklearn.metrics.pairwise import cosine_similarity

conn = psycopg2.connect("dbname=test_cos_sim")
cur = conn.cursor()
q = """
SELECT
  a.id as "id1",
  a.vec as "vec1",
  b.id as "id2",
  b.vec as "vec2"
FROM
  data_vec a,
  data_vec b
WHERE
  a.id = 0
AND
  b.id <> 0
"""

cur.execute(q)
d = []
for r in cur.fetchall():
    d.append((r[2], cosine_similarity([r[1]], [r[3]])[0][0]))

print "sk_cli"
for r in sorted(d, key=lambda s: s[1], reverse=True)[:10]:
    print(r)

conn.close()

■検証環境


今回検証に用いた環境は以下の通りです。Windows上でのVirtualBoxのVM環境です。
  • Core i7-4785T 2.20GHz (Quad core)
  • CentOS 6.6 (x86_64)
  • Python 2.7.9
  • PostgreSQL 9.5.2
  • VM(VirtualBox)に4コアを100%割り当て
  • VM上で4GB RAM

■実装方式による性能比較


以下は、実装方式による実行時間の比較です。

ここでは、まずはベースラインとして500次元の空間ベクトルを10,000レコード作成して、クエリの実行時間(5回実行した平均値)を取得しています。

縦軸は実行時間(短い方が高速)、横軸はそれぞれ
  • (1) CによるUDF (func_native)
  • (2) PL/PythonによるUDF (func_plpy)
  • (3) PL/Python+scikit-learnによるUDF (func_plpy_sk)
  • (4) MADlibのcosine_similarity関数 (func_madlib)
  • (5) クライアント側に取得してscikit-learnで処理 (cli_sk)
を示しています。


この結果を見ると、
  • CによるUDFとMADlibの関数が圧倒的に高速(それぞれ102msと129ms)
  • PL/Python系の実装が一桁遅い(1906msと2854ms)
  • クライアント側にデータを持ってきてソート処理をするのがもっとも遅い(7381ms)
という結果になっています。

■次元数の違いによる性能比較


次に、空間ベクトルの次元数を500から1,000および2,000に増やして実行時間がどのように変化するかを確認します。レコード数はすべて10,000件としています。


結果を見ると、次元数に応じて実行時間は長くなっています。

CによるUDFとMADlibが圧倒的に高速なのは変わらないのですが、PL/Python系の実装について見てみると、500次元の時には(scikit-learnを使わない)素のPL/Pythonの実装(func_plpy)の方が高速であったのが、次元数が2,000になると、scikit-learnを使った実装(func_plpy_sk)の方が高速になっています。

これは、おそらく素の Python で大きな配列を扱うよりは scikit-learn の方が大量の数値データの扱いに長けているためでしょう。(今回は実施しませんでしたが、numpyを使うと素のPythonで実装するよりは高速になるかもしれません)

■レコード数の違いによる性能比較


最後に、空間ベクトルの次元数は2,000のままにして、レコード数を10,000件から20,000件、40,000件と増やしてみて、実行時間がどのように変化するかを確認します。


レコードの増加に伴って、レコード数と同じ程度に実行時間が延びていることが分かります。

ここでも、レコード数が増えると素の PL/Python による実装よりも、PL/Python の UDF 内で scikit-learn を使った方が高速になる傾向が出ています。

なお、クライアント側に取ってきて scikit-learn で処理する方式(cli_sk)は、40,000レコードの時に Out of Memory エラーで実行できなかったので、結果がありません。

以下に、計測した数値を一覧で示します。(数値はミリ秒)


なお、今回利用したコード類は以下に置いておきましたので、興味のある方は合わせてご利用ください。

■まとめ


今回は、「データベース内のデータに対してコサイン類似度を計算して、より似ているレコードを取得する」というケースを想定して、どのような実装がより高速なのかを検証してみました。

その結果として、今回の前提条件であれば、
  • Cで作成したUDFとMADlibはパフォーマンス的にほとんど変わらない。
  • それに比べると、PL/Pythonでの実装は(scikit-learnを使うかどうかに関わらず)1ケタ以上遅い。
  • 処理するデータ量(次元またはレコード数)が多くなると、scikit-learnを使う実装のメリットが出てくる。
  • オンラインシステムでの利用を想定すると、数百ミリ秒で返ってきたMADlibはそのまま利用できる可能性がある。
  • データをクライアントに転送して処理するのは時間がかかる。
ということが分かりました。

興味のある方は、自分のデータを使って、あるいは別のアルゴリズムについてもデータベース内での処理を試してみていただければと思います。

では。

0 件のコメント:

コメントを投稿