DistilBERTの推論速度がCPUとGPUでどれくらい変わるのか比較してみた

Google社が開発した自然言語処理モデルBERTですが、使い方次第では様々なタスクで高い精度を得られるものの、そのパラメータの多さゆえに推論にかなり時間がかかります。
そのためBERTを実運用しようとすると、処理時間がボトルネックになって頓挫する場合もあるのではと思います。

BERTを蒸留したDistilBERT(軽量版BERT)をさらに量子化することで、CPUでも高いパフォーマンスを得られるという記事もあるものの、ある程度長いテキストをBERTで処理する場合は非常にレスポンスに時間がかかります。

そこで、今回は運用コストはいったん度外視し、CPUではなくGPUだとどれくらいBERTの推論処理が速くなるのかを検証します。

なお、コードはColaboratoryでの実行を想定しています。
また、使用するモデルはバンダイナムコ社が公開しているDistilBERTです。
(単純にモデルサイズが軽量でダウンロードが楽という理由です。BERTを使ってもらっても構いません。)

DistilBERTのロード

こちらからDistilBERTをダウンロードし、任意のディレクトリに配置します。
以下のコードでは、Google Driveにマウントし、マイドライブにモデルを配置しています。

from transformers import (DistilBertTokenizer, DistilBertForMaskedLM, DistilBertConfig)
# DitilBERTロード
config = DistilBertConfig.from_json_file(f'{model_path}/config.json')
model_path = '/content/drive/My Drive/Colab Notebooks/DistilBERT-base-jp'
model = DistilBertForMaskedLM.from_pretrained(f'{model_path}/pytorch_model.bin', config=config)
tokenizer = DistilBertTokenizer(f'{model_path}/vocab.txt', do_lower_case=False, do_basic_tokenize=False)
# 推論モード
model.eval()

サンプルテキスト生成

推論速度を比較するため、ある程度大きなサイズの日本語の単語を生成します。
以下のコードでは505単語のリストを作成し、そのリストが10個あるリスト（２次元リスト）を作成しています。
こんなイメージです。

# 作成するリストのイメージ
[
	['私', 'は', '毎日', '朝', 'に', '[MASK]', 'し', 'ます', 'でも', '明日', 'は', '雨', 'が', '降り', 'そう', 'なので', 'やめ', 'ます'],
	['私', 'は', '毎日', '朝', 'に', '[MASK]', 'し', 'ます', 'でも', '明日', 'は', '雨', 'が', '降り', 'そう', 'なので', 'やめ', 'ます'],
	['私', 'は', '毎日', '朝', 'に', '[MASK]', 'し', 'ます', 'でも', '明日', 'は', '雨', 'が', '降り', 'そう', 'なので', 'やめ', 'ます'],
	...
	['私', 'は', '毎日', '朝', 'に', '[MASK]', 'し', 'ます', 'でも', '明日', 'は', '雨', 'が', '降り', 'そう', 'なので', 'やめ', 'ます'],
]

txt = ['私', 'は', '毎日', '朝', 'に', '[MASK]', 'し', 'ます'] # 先頭
append_txt = ['でも', '明日', 'は', '雨', 'が', '降り', 'そう', 'なので', 'やめ', 'ます'] # 追記テキスト
txt.insert(0, '[CLS]')
txt.append('[SEP]')
# n回テキストを追記
n = 45
for i in range(n):
    txt.extend(append_txt)
    txt.append('[SEP]')
print(f'token size:{len(txt)}')

# バッチ化
txt_list = [txt for i in range(10)]
print(f'text size: {len(txt_list)}')

token size:505
text size: 10

CPUとGPUで推論速度比較

CPUとGPUで処理速度を比較します。
1文単位で推論する場合と、全文まとめてバッチ処理する場合に分けて検証します。

import time

def get_elapsed_time_for_bert(model, tokenizer, tokens, device, by_sentence):
    """
    BERTの推論パフォーマンスを計測する
    """
    def predict(model, tokens_tensor):
        with torch.no_grad():
            outputs = model(tokens_tensor)
            predictions = outputs[0]
        return predictions
    
    start = time.time()
    model.to(device)
    ids = [tokenizer.convert_tokens_to_ids(token) for token in tokens]
    if by_sentence:
        for id in ids:
            tokens_tensor = torch.tensor([id]).to(device)
            predict(model, tokens_tensor)
    else:
        tokens_tensor = torch.tensor(ids).to(device)
        predict(model, tokens_tensor)
    elapsed_time = time.time() - start
    print(f'on {device}: {elapsed_time}')

for by_sentence in [True, False]:
    print(f'by_sentence: {by_sentence}')
    for device in ['cpu', 'cuda:0']:
        get_elapsed_time_for_bert(
            model=distilbert_model_fine,
            tokenizer=distilbert_tokenizer_fine,
            tokens=txt_list,
            device=device,
            by_sentence=by_sentence
        )

by_sentence: True
on cpu: 18.735729455947876
on cuda:0: 0.3672614097595215
by_sentence: False
on cpu: 18.155377864837646
on cuda:0: 0.08703923225402832

「by_sentence: True」は1文単位の処理、「by_sentence: False」は全文まとめたバッチ処理を意味しています。
処理単位にかかわらずGPUの方が圧倒的に処理速度が向上していることがわかります。
また、CPUでは1文単位の処理であろうとバッチ処理であろうと推論速度にほとんど差はありませんが、GPUではバッチ処理の方が格段に速くなっています。
これは、GPUがバッチ処理に秀でているということを考えれば当然の結果ではありますが。。

最後に

GPUとCPUを比較した結果、やはりGPUを使った方がBERTの推論処理が圧倒的に速くなることがわかりました。
ただ、だからと言って本番で運用する際に何でもかんでもGPUにすればいいというわけではありません。
GPUはCPUに比べて運用コストが高くなるためです。
実運用していくなかで、推論で入力するデータ量、許容できるレスポンスの時間などを考慮すればCPUでも十分なパフォーマンスを担保できることもあるでしょう。