rubyのQueueを使ったプログラムのプロファイルをとってみた

rubyのQueueを使ったプログラムののプロファイルをとってみました。

rubyでプロファイルをとるには、以下のように -r profile を指定します。

ruby -r profile {プログラム名}

stderr にプロファイル結果が出力されます。

プロファイル取得対象のプログラムは以下の通りです。
stdin から1行読み込み、キューに push し、各スレッドは キューからpopしたデータを出力します。

require 'thread'

num_threads = 3
q = Queue.new()

ts = []
0.upto(num_threads) do |i|
  ts[i] = Thread.new() do |t|
    while true do
      str = q.pop()
      break if ! str

      #item = itemstr.split(/\t/, -1)
      printf("thread %i [%s]\n", i, str)
    end
  end
end

STDIN.each do |line|
  line.chomp!
  q.push(line)
end

上記のプログラム test1.rb を以下のように実行すると、profile.txt にます。

ruby -e '0.upto(10000) { |i| printf("%d\n", i) }' \
| ruby -r profile test1.rb \
> /dev/null \
2> profile.txt

profile.txt のプロファイルは以下のようになりました。（前半のみです）

  %   cumulative   self              self     total
 time   seconds   seconds    calls  ms/call  ms/call  name
 33.05     0.39      0.39    18415     0.02     0.09  Mutex#synchronize
 18.64     0.61      0.22     8414     0.03     0.27  Queue#pop
 10.17     0.73      0.12    10001     0.01     0.10  Queue#push
  6.78     0.81      0.08     8414     0.01     0.01  Kernel.printf
  6.78     0.89      0.08    18419     0.00     0.01  Mutex#lock
  6.78     0.97      0.08    18415     0.00     0.00  Mutex#unlock
  4.24     1.02      0.05    10689     0.00     0.00  Array#empty?
  4.24     1.07      0.05    18415     0.00     0.00  Array#shift
  1.69     1.09      0.02     8414     0.00     0.00  BasicObject#!
  1.69     1.11      0.02    10001     0.00     0.00  String#chomp!
  1.69     1.13      0.02    12276     0.00     0.00  Array#push
  1.69     1.15      0.02     8414     0.00     0.00  IO#write
  0.85     1.16      0.01     2275     0.00     0.00  Thread#current
  0.85     1.17      0.01     2271     0.00     0.02  Mutex#sleep
  0.85     1.18      0.01        1    10.00    50.00  IO#each

Mutex、Queueで時間がかかっていることがわかります。

nmslib のパラメータ

ベクトル検索の nmslib のパラメータについてのメモ。

ベクトル間の cos 類似度が高い順に検索するには、以下のように space='cosinesimil' を指定して初期化します。

index = nmslib.init(space='cosinesimil') 

■サンプルプログラム（cos 類似度版）

import numpy as np
import nmslib

index = nmslib.init(method='hnsw', space='cosinesimil')

vs = [
    [1.0, 2.0],
    [2.0, 2.0],
    [3.0, 3.0],
]

# インデックス生成
vs = np.array(vs, dtype=np.float32)
index.addDataPointBatch(vs)
index.createIndex({}, print_progress=False)

# 検索
v = np.array([1.0, 1.0], dtype=np.float32)
ids, dists = index.knnQuery(v, 10)

print('ids: %s' % (ids))
print('dists: %s' % (sims))

実行結果

ids: [2 1 0]
dists: [0.0000000e+00 5.9604645e-08 5.1316738e-02]

※cos類似度=1の場合に距離=0

一方、ベクトル間の距離が短い順に検索するには、space='l2' を指定して初期化します。

index = nmslib.init(space='l2')

■サンプルプログラム（距離版）

import numpy as np
import nmslib

index = nmslib.init(method='hnsw', space='l2')

vs = [
    [1.0, 2.0],
    [2.0, 2.0],
    [3.0, 3.0],
]

# インデックス生成
vs = np.array(vs, dtype=np.float32)
index.addDataPointBatch(vs)
index.createIndex({}, print_progress=False)

# 検索
v = np.array([1.0, 1.0], dtype=np.float32)
ids, dists = index.knnQuery(v, 10)

print('ids: %s' % (ids))
print('dists: %s' % (dists))

実行結果

ids: [0 1 2]
dists: [1. 2. 8.]

※距離は2乗和となっている

その他に、method='hnsw' の場合、createIndex() メソッドの第1引数のパラメータに
'post' を指定すると後処理の制御ができます。
0: 後処理なし、1, 2: 後処理あり（2 の方が強力）