Embedの速度比較

単語埋め込みとかで使う、onehotとlook-up tabelの内積計算を比較する。

埋め込み操作(embedding)の式は以下のとおり。

 \boldsymbol{v} = \boldsymbol{x}\boldsymbol{W}

xは[0,1,0,0,0]って感じのonehotベクトル。欲しいインデックスにビットが立ってて、Wと内積とると欲しいベクトルが引き抜かれる感じ。

これをコードに落とす時、くそまじめにインデックスからonehot作って内積とると重すぎて死ぬので、戒めも兼ねて速度比較を行った。
以下のようなコードを用いる。

Embeddingの操作として、

1. Listを使って目当てのベクトルを引き抜く
2. 内積で引き抜く(式通り)
3. chainerのようなライブラリを使う(中身は知らない)

の三つが考えられるので、それぞれlook-up tableを用意する時間、1万回ベクトルを引き抜く操作を繰り返した時の時間を測る。

ただし、2ではindexをonehotに直す操作、3ではnumpyに包む操作が入って不公平(?)なのでそこは計測しない。
chainerのembedに関しては、中でさらに何かやってるんだろうけど、それは考えないことにする。

import chainer
from chainer import links as L
import numpy as np
import time
import random

iterationSize = 10000
vocSize = 10000
embedSize =100

# 1. リストを用いてembedding
def testListEmbed():
    pStart = time.time()
    # 準備
    embed = [np.random.rand(1, embedSize) for i in range(vocSize)]
    pT = time.time() - pStart
    print('preparation:', pT)
    
    mainStart = time.time()
    # 繰り返し
    for i in range(iterationSize):
        indice = random.randint(0, vocSize-1)
        vec = embed[indice]
    mainT = time.time() - mainStart
    print('time:', mainT)

# 2. numpyで内積
def testNumpyEmbed():
    pStart = time.time()
    # 準備
    embed = np.random.rand(vocSize, embedSize)
    pT = time.time() - pStart
    print('preparation:', pT)

    # 優遇
    onehots = []
    for i in range(iterationSize):
        onehot = np.zeros((1, vocSize))
        onehot[0, random.randint(0, vocSize-1)] = 1
        onehots.append(onehot)
    
    mainStart = time.time()
    # 繰り返し
    for onehot in onehots:
        vec = np.dot(onehot, embed)
    mainT = time.time() - mainStart
    print('time:', mainT)

# 3. chainerで用意されたembeddingを使う
def testChainerEmbed():
    pStart = time.time()
    # 準備
    embed = L.EmbedID(vocSize, embedSize)
    pT = time.time() - pStart
    print('preparation:', pT)

    # 優遇
    indices = [np.array([random.randint(0, vocSize-1)], 'i') for i in range(iterationSize)]

    mainStart = time.time()
    # 繰り返し
    for indice in indices:
        vec = embed(indice)
    mainT = time.time() - mainStart
    print('time:', mainT)

print('List')
testListEmbed()

print('\nNumpy-dot')
testNumpyEmbed()

print('\nChainer')
testChainerEmbed()

3回実行した結果

# 1回目
List
preparation: 0.03716588020324707
time: 0.032792091369628906

Numpy-dot
preparation: 0.01922297477722168
time: 6.139436960220337

Chainer
preparation: 0.05529212951660156
time: 1.03914213180542

# 2回目
List
preparation: 0.12082695960998535
time: 0.0773320198059082

Numpy-dot
preparation: 0.03717994689941406
time: 8.129857063293457

Chainer
preparation: 0.05550503730773926
time: 1.3018569946289062

# 3回目
List
preparation: 0.0423738956451416
time: 0.03646492958068848

Numpy-dot
preparation: 0.021611928939819336
time: 7.694821834564209

Chainer
preparation: 0.05309700965881348
time: 1.3062717914581299


まあそんなもんか、という感想。
chainerはVariableで一気にbackwardできる利便性のぶん、処理速度が犠牲になってる感じ。
深いニューラルならchainerの恩恵を甘受するべきだと思う。

浅いニューラルをスクラッチで書くなら、Listを使って引き抜いた方が速い。
更新する時どうなるかまでは測ってない。
もしかしたら学習のイテレーション全体で見たらchainerが圧勝するのかもしれない。

でもまあ、更新する時もインデックス指定して微分値足すだけなのでリストが爆速なんだと思う。


もしもっと効率的な書き方知ってる方いたら教えてください。