洋食の日記

「た・である」調ではなく「です・ます」調で書きはじめれば良かったなと後悔してる人のブログです

近似最近傍探索ライブラリAnnoyのRuby bindingを作成した

はじめに

Annoy (Approximate Nearest Neighbors Oh Yeah) は、C++で書かれた近似最近傍探索ライブラリである。近似最近傍探索とは、その名のとおり、クエリに対して厳密ではなく近似的に近傍にあるものを探索する。高速に探索できるので、大規模データの画像の類似検索などに使われる。インデックスの木の構築では、データ数を二分割するような2点を選び、それらでできる超平面で分割することを繰り返す。k-means tree (vocabulary tree) をアレンジしたような手法になっている。Annoyでは、Python bindingが提供されていて、pipでインストールできる。そのため、Pythonライブラリなイメージがあるが、実際は一つのヘッダーファイルからなるシンプルなライブラリである。これを、native extensionsから叩くかたちでgemにした。

annoy-rb | RubyGems.org | your community gem host

使い方

インストールは、普通にgemコマンドでインストールできる。特別な外部ライブラリもPythonも必要ない。

$ gem install annoy-rb

APIは、AnnoyのPython bindingに合わせた。Python版を使ったことがあれば、すぐに使えると思う。

まずは、検索インデックスを作成する。AnnoyIndexをnewする際に、ベクトルの次元数と、距離基準を指定する(angularはコサイン距離で、他にユークリッド距離やマンハッタン距離などがある)。add_itemメソッドで、データの番号とRuby Arrayによるベクトルを登録していく。そして、buildメソッドで検索インデックスを作成する。作成したインデックスはsaveメソッドで保存できる。

require 'annoy'

n_features = 40 # 検索インデックスに追加するベクトルの次元数
t = Annoy::AnnoyIndex.new(n_features: n_features, metric: 'angular')

# ランダムな値からなるベクトルを1000件追加する.
1000.times do |i|
  v = Array.new(n_features) { rand }
  t.add_item(i, v)
end

# インデックスを構築する. この際, 構築する木の数を指定できる.
# 木の数は, 多いほど検索精度は上がるが遅くなる.
n_trees = 10
t.build(n_trees)

# インデックスを保存する.
t.save('foo.ann')

検索は、get_nns_by_itemとget_nns_by_vectorメソッドで行う。get_nns_by_itemは、検索インデックス中のベクトルを番号で指定し、それをクエリとして検索を行う。get_nns_by_vectorは、クエリとして任意のベクトルをRuby Arrayで与えて検索を行う。

require 'annoy'

# 作成した検索インデックスを読み込む.
n_features = 40
u = Annoy::AnnoyIndex.new(n_features: n_features, metric: 'angular')
u.load('foo.ann')

# アイテム番号0番に近い10件を検索する.
# 近傍とされるアイテムの番号がRuby Arrayで返る.
neighbor_ids = t.get_nns_by_item(0, 10)

# 任意のベクトルに近い10件を検索する.
# include_distancesにtrueを与えると, 距離も返す.
q = Array.new(n_features) { rand }
neighbor_ids, neighbor_dists = t.get_nns_by_vector(1, 10, include_distances: true)

おわりに

Annoyをbindingしたgemを作成した。シンプルにnative extensionsで繋いでるだけで(コードも300行程度)、今後大きくコードを改善する箇所はあまりないが、bugfixやAnnoy側のアップデートなどに対応するなど、ちょこちょこ面倒を見ていきたい。

github.com