ゆるいDeep Learning

Hugging FaceのBERTモデルの推論をTorch TensorRTで高速化

Deep Learning

Hugging Faceのモデルで高速に推論したいです。

Torch TensorRTが対応しているぞ

事前準備

Torch TensorRTに関しては下記記事をご覧ください

TRTorchを使用してPyTorchの推論を高速化する

Torch TensorRTの環境はdockerを使用すると簡単に用意できます。

https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch

docker run --gpus all -it --rm -p 8887:8887 --name tensorrt nvcr.io/nvidia/pytorch:22.07-py3

下記コードを参考に実行します。

https://github.com/pytorch/TensorRT/blob/master/notebooks/Hugging-Face-BERT.ipynb

Docker内にすでにコードがあるので、そのコードを使用します。

Jupyter-labを起動してコードにアクセスします。

jupyter-lab --port=8887

/workspace/examples/torch_tensorrt/notebooksディレクトリにHugging-Face-BERT.ipynbのNotebookがあるのでこちらで検証します。

ライブラリの導入

Hugging Faceのライブラリをインストールします。

!pip install transformers

ライブラリをインポートします。

from transformers import BertTokenizer, BertForMaskedLM
import torch
import timeit
import numpy as np
import torch_tensorrt
import torch.backends.cudnn as cudnn

BERTモデル

BERTのモデルはTransformerのモデルをベースに作成されています。

PyTorchで自然言語処理でよく使用されるTransformerを少ないコードで使用してみる

Hugging Faceを使用することで簡単にBERTのモデルを利用できます。

下記コードで事前学習済みモデルを取得します。

enc = BertTokenizer.from_pretrained('bert-base-uncased')

BERTのモデルに与えるダミーのインプットを作成します。

batch_size = 4

batched_indexed_tokens = [[101, 64]*64]*batch_size
batched_segment_ids = [[0, 1]*64]*batch_size
batched_attention_masks = [[1, 1]*64]*batch_size

tokens_tensor = torch.tensor(batched_indexed_tokens)
segments_tensor = torch.tensor(batched_segment_ids)
attention_masks_tensor = torch.tensor(batched_attention_masks)

高速に推論できるようにTorch Script化を行います。

BERTで推論を高速化できるTorchScriptを試してみる
mlm_model_ts = BertForMaskedLM.from_pretrained('bert-base-uncased', torchscript=True)
traced_mlm_model = torch.jit.trace(mlm_model_ts, [tokens_tensor, segments_tensor, attention_masks_tensor])

推論用のデータを準備します。[MASK]部分を推論します。

masked_sentences = ['Paris is the [MASK] of France.', 
                    'The primary [MASK] of the United States is English.', 
                    'A baseball game consists of at least nine [MASK].', 
                    'Topology is a branch of [MASK] concerned with the properties of geometric objects that remain unchanged under continuous transformations.']
pos_masks = [4, 3, 9, 6]

下記コードで推論用のデータを推論します。

自然言語処理ではデータのサイズが異なるケースが多いので最大長を指定して、それに合わせてPadding処理で穴埋めしています。

最大長は128にしています。

encoded_inputs = enc(masked_sentences, return_tensors='pt', padding='max_length', max_length=128)
outputs = mlm_model_ts(**encoded_inputs)
most_likely_token_ids = [torch.argmax(outputs[0][i, pos, :]) for i, pos in enumerate(pos_masks)]
unmasked_tokens = enc.decode(most_likely_token_ids).split(' ')
unmasked_sentences = [masked_sentences[i].replace('[MASK]', token) for i, token in enumerate(unmasked_tokens)]
for sentence in unmasked_sentences:
    print(sentence)

下記のような結果が取得できます。

[MASK]部分が正しく推論できているように見えます。

Paris is the capital of France.
The primary language of the United States is English.
A baseball game consists of at least nine innings.
Topology is a branch of mathematics concerned with the properties of geometric objects that remain unchanged under continuous transformations.

Dummyの入力でTraceしたモデルで推論します。データの与え方は'input_ids', 'token_type_ids', 'attention_mask'の3つのキーで別々で与える形にしています。

encoded_inputs = enc(masked_sentences, return_tensors='pt', padding='max_length', max_length=128)
outputs = traced_mlm_model(encoded_inputs['input_ids'], encoded_inputs['token_type_ids'], encoded_inputs['attention_mask'])
most_likely_token_ids = [torch.argmax(outputs[0][i, pos, :]) for i, pos in enumerate(pos_masks)]
unmasked_tokens = enc.decode(most_likely_token_ids).split(' ')
unmasked_sentences = [masked_sentences[i].replace('[MASK]', token) for i, token in enumerate(unmasked_tokens)]
for sentence in unmasked_sentences:
    print(sentence)

通常のモデルとは異なる結果を出していますが、こちらも正しい推論結果のように見えます。

Paris is the all of France.
The primary all of the United States is English.
A baseball game consists of at least nine each.
Topology is a branch of each concerned with the properties of geometric objects that remain unchanged under continuous transformations.

Torch TensorRT形式に変換

ログレベルを変更します。

new_level = torch_tensorrt.logging.Level.Error
torch_tensorrt.logging.set_reportable_log_level(new_level)

Torch TensorRT形式に変換します。

trt_model = torch_tensorrt.compile(traced_mlm_model, 
    inputs= [torch_tensorrt.Input(shape=[batch_size, 128], dtype=torch.int32),  # input_ids
             torch_tensorrt.Input(shape=[batch_size, 128], dtype=torch.int32),  # token_type_ids
             torch_tensorrt.Input(shape=[batch_size, 128], dtype=torch.int32)], # attention_mask
    enabled_precisions= {torch.float32}, # Run with 32-bit precision
    workspace_size=2000000000,
    truncate_long_and_double=True
)

同様に推論処理をします。

enc_inputs = enc(masked_sentences, return_tensors='pt', padding='max_length', max_length=128)
enc_inputs = {k: v.type(torch.int32).cuda() for k, v in enc_inputs.items()}
output_trt = trt_model(enc_inputs['input_ids'], enc_inputs['token_type_ids'], enc_inputs['attention_mask'])
most_likely_token_ids_trt = [torch.argmax(output_trt[i, pos, :]) for i, pos in enumerate(pos_masks)] 
unmasked_tokens_trt = enc.decode(most_likely_token_ids_trt).split(' ')
unmasked_sentences_trt = [masked_sentences[i].replace('[MASK]', token) for i, token in enumerate(unmasked_tokens_trt)]
for sentence in unmasked_sentences_trt:
    print(sentence)

推論結果は下記のようになります。

Paris is the all of France.
The primary all of the United States is English.
A baseball game consists of at least nine each.
Topology is a branch of each concerned with the properties of geometric objects that remain unchanged under continuous transformations.

数値精度をFP16に変更して実行します。

trt_model_fp16 = torch_tensorrt.compile(traced_mlm_model, 
    inputs= [torch_tensorrt.Input(shape=[batch_size, 128], dtype=torch.int32),  # input_ids
             torch_tensorrt.Input(shape=[batch_size, 128], dtype=torch.int32),  # token_type_ids
             torch_tensorrt.Input(shape=[batch_size, 128], dtype=torch.int32)], # attention_mask
    enabled_precisions= {torch.half}, # Run with 16-bit precision
    workspace_size=2000000000,
    truncate_long_and_double=True
)

モデルの推論性能の比較

ベンチマークを行うための関数を作成します。Warm Up処理後に50回ほど推論処理をします。

def timeGraph(model, input_tensor1, input_tensor2, input_tensor3, num_loops=50):
    print("Warm up ...")
    with torch.no_grad():
        for _ in range(20):
            features = model(input_tensor1, input_tensor2, input_tensor3)

    torch.cuda.synchronize()

    print("Start timing ...")
    timings = []
    with torch.no_grad():
        for i in range(num_loops):
            start_time = timeit.default_timer()
            features = model(input_tensor1, input_tensor2, input_tensor3)
            torch.cuda.synchronize()
            end_time = timeit.default_timer()
            timings.append(end_time - start_time)
            # print("Iteration {}: {:.6f} s".format(i, end_time - start_time))

    return timings

取得したベンチマークの値のレイテンシー、スループットの中央値、平均値、標準偏差を取得します。

比較して見やすいようにコードを若干変更しています。

def printStats(graphName, timings, batch_size):
    times = np.array(timings)
    steps = len(times)
    speeds = batch_size / times
    time_mean = np.mean(times)
    time_med = np.median(times)
    time_99th = np.percentile(times, 99)
    time_std = np.std(times, ddof=0)
    speed_mean = np.mean(speeds)
    speed_med = np.median(speeds)

    msg = ("\n%s =================================\n"
            "batch size=%d, num iterations=%d\n"
            "  Median text batches/second: %.1f, mean: %.1f\n"
            "  Median latency: %.6f, mean: %.6f, 99th_p: %.6f, std_dev: %.6f\n"
            ) % (graphName,
                batch_size, steps,
                speed_med, speed_mean,
                time_med, time_mean, time_99th, time_std)
    print(msg)
    return speed_med, speed_mean, time_med, time_mean
cudnn.benchmark = True

Torch Script化したモデルの性能を計測します。

timings = timeGraph(mlm_model_ts.cuda(), enc_inputs['input_ids'], enc_inputs['token_type_ids'], enc_inputs['attention_mask'])

speed_med, speed_mean, time_med, time_mean = printStats("BERT", timings, batch_size)

dta_dict = {"speed_med": [speed_med], 
            "speed_mean": [speed_mean], 
            "time_med": [time_med], 
            "time_mean": [time_mean]}

著者の環境では下記のような性能になりました。

Warm up ...
Start timing ...

BERT =================================
batch size=4, num iterations=50
  Median text batches/second: 599.1, mean: 597.6
  Median latency: 0.006677, mean: 0.006693, 99th_p: 0.006943, std_dev: 0.000059

Traceしたモデルで性能を計測します。

timings = timeGraph(traced_mlm_model.cuda(), enc_inputs['input_ids'], enc_inputs['token_type_ids'], enc_inputs['attention_mask'])

speed_med, speed_mean, time_med, time_mean = printStats("BERT", timings, batch_size)

dta_dict["speed_med"].append(speed_med)
dta_dict["speed_mean"].append(speed_mean)
dta_dict["time_med"].append(time_med)
dta_dict["time_mean"].append(time_mean)

著者の環境では下記のような性能になりました。

Warm up ...
Start timing ...

BERT =================================
batch size=4, num iterations=50
  Median text batches/second: 951.2, mean: 951.0
  Median latency: 0.004205, mean: 0.004206, 99th_p: 0.004256, std_dev: 0.000015

Torch TensorRT形式に変換したモデルで計測します。

timings = timeGraph(trt_model, enc_inputs['input_ids'], enc_inputs['token_type_ids'], enc_inputs['attention_mask'])

speed_med, speed_mean, time_med, time_mean = printStats("BERT", timings, batch_size)

dta_dict["speed_med"].append(speed_med)
dta_dict["speed_mean"].append(speed_mean)
dta_dict["time_med"].append(time_med)
dta_dict["time_mean"].append(time_mean)

著者の環境では下記のような性能になりました。

Warm up ...
Start timing ...

BERT =================================
batch size=4, num iterations=50
  Median text batches/second: 1216.9, mean: 1216.4
  Median latency: 0.003287, mean: 0.003289, 99th_p: 0.003317, std_dev: 0.000007

Torch TensorRT形式に変換＋数値精度FP16のモデルで計測します。

timings = timeGraph(trt_model_fp16, enc_inputs['input_ids'], enc_inputs['token_type_ids'], enc_inputs['attention_mask'])

speed_med, speed_mean, time_med, time_mean = printStats("BERT", timings, batch_size)

dta_dict["speed_med"].append(speed_med)
dta_dict["speed_mean"].append(speed_mean)
dta_dict["time_med"].append(time_med)
dta_dict["time_mean"].append(time_mean)

著者の環境では下記のような性能になりました。

Warm up ...
Start timing ...

BERT =================================
batch size=4, num iterations=50
  Median text batches/second: 1776.7, mean: 1771.1
  Median latency: 0.002251, mean: 0.002259, 99th_p: 0.002305, std_dev: 0.000015

全結果をまとめて確認します。

import pandas as pd

pd.DataFrame.from_dict(dta_dict, orient='index', columns=["BERT_ts", "BERT_trace", "BERT_Torch_TensorRT", "BERT_Torch_TensorRT_FP16"]).T

Torch Scriptのモデルに比べて中央値のレイテンシーは約3倍ほど高速化できています。

Torch TensorRTがHugging Faceのモデルにも適用できました。

これでいろんな自然言語処理モデルの推論性能を高速化できるな！

Close Bitnami banner
Bitnami