np.choiceとchainer.utils.WalkerAliasの速度対決

chainerのマニュアルに「chainer.utils.WalkerAliasはnp.choiceよりも速い」と書いてあったので，速度を比較してみた．ただしCPUオンリー．

• chainer.utils.WalkerAlias
• Alias method on Wikipedia

コード

乱数生成

import numpy as np
import chainer
from time import time
import matplotlib.pyplot as plt
%matplotlib inline

loops = 10  # for taking average
verc_size = 10 ** np.arange(1, 7)  # i
n_samples = 10 ** np.arange(1, 7)  # j

t_choice = np.zeros((len(verc_size), len(n_samples)))
t_alias  = np.zeros((len(verc_size), len(n_samples)))
t_aliasv = np.zeros((len(verc_size), len(n_samples)))



for i in range(len(verc_size)):

    p = np.random.randint(100, size=verc_size[i])
    p = p / p.sum()
    # p = p.astype(np.float32)  # for gpu?
    c = np.random.randint(100, size=verc_size[i])

    start_time = time()
    for loop in range(loops):
        sampler = chainer.utils.WalkerAlias(p)
    end_time = time()
    e_time = end_time - start_time
    t_aliasv[i, :] = e_time / loops


    for j in range(len(n_samples)):

        # numpy choice
        start_time = time()
        for loop in range(loops):
            res = np.random.choice(c, size=n_samples[j], p=p)
        end_time = time()
        e_time = end_time - start_time
        print('choice', e_time, n_samples[j], verc_size[i])
        t_choice[i, j] = e_time / loops

        # chainer alisas
        start_time = time()
        for loop in range(loops):
            res = sampler.sample(n_samples[j])
            res = c[res]
        end_time = time()
        e_time = end_time - start_time
        print('alias ', e_time, n_samples[j], verc_size[i])
        t_alias[i, j] = e_time / loopsimport numpy as np

プロット

プロット

fig, ax = plt.subplots(nrows=1, ncols=3, figsize=(15, 5))

def myax(ax, title):
    ax.set_yscale('log')
    ax.set_xscale('log')
    ax.set_ylim(1e-4, 1e2)
    ax.set_title(title)
    ax.set_xlabel('# samples')
    ax.set_ylabel('time [s]')

for i in range(len(verc_size)):
    ax[0].plot(n_samples, t_choice[i], label=verc_size[i])
myax(ax[0], 'np.choice')
ax[0].legend(loc='best')

ax[1].plot(n_samples, t_alias)
myax(ax[1], 'chainer.alias sampling only')

ax[2].plot(n_samples, t_alias + t_aliasv)
myax(ax[2], 'chainer.alias init + sampling')

結果

• alias methodは確かにO(1)ということが分かる（sampling onlyの場合．計算時間はサンプル数に依存しない）．
• alias methodは初期化に膨大な時間がかかるので，初期化の時間を含めると確実にnp.choiceのほうが速い．
• 初期化を除くと，「1万個の中から1つ選ぶ」を1万回繰り返すときには，alias methodのほうが速そう．
• GPUでalias methodがどの程度速いのかは未検証（chainer.utils.WalkerAliasにはto_gpu()がある）．