Refactor: Use I/O Binding during BERT inference and always release memory after inference
This commit is contained in:
tsukumi
@@ -11,6 +11,7 @@ from style_bert_vits2.nlp import (
|
||||
cleaned_text_to_sequence,
|
||||
extract_bert_feature_onnx,
|
||||
)
|
||||
from style_bert_vits2.utils import get_onnx_device_options
|
||||
|
||||
|
||||
def __intersperse(lst: list[Any], item: Any) -> list[Any]:
|
||||
@@ -187,34 +188,11 @@ def infer_onnx(
|
||||
np.array(noise_scale_w, dtype=np.float32),
|
||||
]
|
||||
|
||||
# 入力テンソルを転送する GPU デバイスを取得
|
||||
## 本来は device_type="dml" もサポートされているはずだが、手元環境だと常に謎の RuntimeError が発生するため当面無効化している
|
||||
first_provider = onnx_session.get_providers()[0]
|
||||
if first_provider == "CUDAExecutionProvider":
|
||||
device_type = "cuda"
|
||||
# elif first_provider == "DmlExecutionProvider":
|
||||
# device_type = "dml"
|
||||
else:
|
||||
device_type = "cpu"
|
||||
# 入力テンソルの転送に使用するデバイス種別, デバイス ID, 実行オプションを取得
|
||||
device_type, device_id, run_options = get_onnx_device_options(onnx_session, onnx_providers) # fmt: skip
|
||||
|
||||
# 入力テンソルを転送する GPU デバイスの ID を取得
|
||||
## ExecutionProvider に指定したオプションの中から device_id を取得し、入力テンソルの転送先として指定する
|
||||
## InferenceSession で利用するデバイス ID と入力テンソルの転送先デバイス ID は一致している必要がある
|
||||
## 本来は ExecutionProvider に指定したオプションは InferenceSession.get_provider_options() で取得できるはずだが、
|
||||
## 手元環境だと DmlExecutionProvider のみ常に空の辞書が返されるため、当面 onnx_providers から直接オプションを取り出している
|
||||
device_id = 0
|
||||
onnx_providers_dict: dict[str, dict[str, Any]] = {}
|
||||
for provider in onnx_providers:
|
||||
if isinstance(provider, tuple):
|
||||
provider_name, options = provider
|
||||
onnx_providers_dict[provider_name] = options
|
||||
else:
|
||||
onnx_providers_dict[provider] = {}
|
||||
first_provider_options = onnx_providers_dict[first_provider]
|
||||
if "device_id" in first_provider_options:
|
||||
device_id = int(first_provider_options["device_id"])
|
||||
|
||||
# GPU メモリに入力テンソルを割り当て
|
||||
# 推論デバイスに入力テンソルを割り当て
|
||||
## GPU 推論の場合、device_type + device_id に対応する GPU デバイスに入力テンソルが割り当てられる
|
||||
io_binding = onnx_session.io_binding()
|
||||
for name, value in zip(input_names, input_tensor):
|
||||
gpu_tensor = onnxruntime.OrtValue.ortvalue_from_numpy(
|
||||
@@ -224,7 +202,7 @@ def infer_onnx(
|
||||
|
||||
# 推論の実行
|
||||
io_binding.bind_output(output_name, device_type)
|
||||
onnx_session.run_with_iobinding(io_binding)
|
||||
onnx_session.run_with_iobinding(io_binding, run_options=run_options)
|
||||
output = io_binding.get_outputs()
|
||||
|
||||
audio = output[0].numpy()[0, 0]
|
||||
|
||||
@@ -3,10 +3,12 @@ from __future__ import annotations
|
||||
from typing import TYPE_CHECKING, Any, Optional, Sequence, Union
|
||||
|
||||
import numpy as np
|
||||
import onnxruntime
|
||||
from numpy.typing import NDArray
|
||||
|
||||
from style_bert_vits2.constants import Languages
|
||||
from style_bert_vits2.nlp import bert_models, onnx_bert_models
|
||||
from style_bert_vits2.utils import get_onnx_device_options
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -97,34 +99,59 @@ def extract_bert_feature_onnx(
|
||||
NDArray[Any]: BERT の特徴量
|
||||
"""
|
||||
|
||||
# トークナイザーとモデルの読み込み
|
||||
tokenizer = onnx_bert_models.load_tokenizer(Languages.ZH)
|
||||
inputs = tokenizer(text, return_tensors="np")
|
||||
|
||||
session = onnx_bert_models.load_model(
|
||||
language=Languages.ZH,
|
||||
onnx_providers=onnx_providers,
|
||||
)
|
||||
input_names = [input.name for input in session.get_inputs()]
|
||||
output_name = session.get_outputs()[0].name
|
||||
res = session.run(
|
||||
[output_name],
|
||||
{
|
||||
"input_ids": inputs["input_ids"].astype(np.int64), # type: ignore
|
||||
"token_type_ids": inputs["token_type_ids"].astype(np.int64), # type: ignore
|
||||
"attention_mask": inputs["attention_mask"].astype(np.int64), # type: ignore
|
||||
},
|
||||
)[0]
|
||||
|
||||
# 入力テンソルの転送に使用するデバイス種別, デバイス ID, 実行オプションを取得
|
||||
device_type, device_id, run_options = get_onnx_device_options(session, onnx_providers) # fmt: skip
|
||||
|
||||
# 入力をテンソルに変換
|
||||
inputs = tokenizer(text, return_tensors="np")
|
||||
input_tensor = [
|
||||
inputs["input_ids"].astype(np.int64), # type: ignore
|
||||
inputs["token_type_ids"].astype(np.int64), # type: ignore
|
||||
inputs["attention_mask"].astype(np.int64), # type: ignore
|
||||
]
|
||||
# 推論デバイスに入力テンソルを割り当て
|
||||
## GPU 推論の場合、device_type + device_id に対応する GPU デバイスに入力テンソルが割り当てられる
|
||||
io_binding = session.io_binding()
|
||||
for name, value in zip(input_names, input_tensor):
|
||||
gpu_tensor = onnxruntime.OrtValue.ortvalue_from_numpy(
|
||||
value, device_type, device_id
|
||||
)
|
||||
io_binding.bind_ortvalue_input(name, gpu_tensor)
|
||||
# text から BERT 特徴量を抽出
|
||||
io_binding.bind_output(output_name, device_type)
|
||||
session.run_with_iobinding(io_binding, run_options=run_options)
|
||||
res = io_binding.get_outputs()[0].numpy()
|
||||
|
||||
style_res_mean = None
|
||||
if assist_text:
|
||||
# 入力をテンソルに変換
|
||||
style_inputs = tokenizer(assist_text, return_tensors="np")
|
||||
style_res = session.run(
|
||||
[output_name],
|
||||
{
|
||||
"input_ids": style_inputs["input_ids"].astype(np.int64), # type: ignore
|
||||
"token_type_ids": style_inputs["token_type_ids"].astype(np.int64), # type: ignore
|
||||
"attention_mask": style_inputs["attention_mask"].astype(np.int64), # type: ignore
|
||||
},
|
||||
)[0]
|
||||
style_input_tensor = [
|
||||
style_inputs["input_ids"].astype(np.int64), # type: ignore
|
||||
style_inputs["token_type_ids"].astype(np.int64), # type: ignore
|
||||
style_inputs["attention_mask"].astype(np.int64), # type: ignore
|
||||
]
|
||||
# 推論デバイスに入力テンソルを割り当て
|
||||
## GPU 推論の場合、device_type + device_id に対応する GPU デバイスに入力テンソルが割り当てられる
|
||||
io_binding = session.io_binding() # IOBinding は作り直す必要がある
|
||||
for name, value in zip(input_names, style_input_tensor):
|
||||
gpu_tensor = onnxruntime.OrtValue.ortvalue_from_numpy(
|
||||
value, device_type, device_id
|
||||
)
|
||||
io_binding.bind_ortvalue_input(name, gpu_tensor)
|
||||
# assist_text から BERT 特徴量を抽出
|
||||
io_binding.bind_output(output_name, device_type)
|
||||
session.run_with_iobinding(io_binding, run_options=run_options)
|
||||
style_res = io_binding.get_outputs()[0].numpy()
|
||||
style_res_mean = np.mean(style_res, axis=0)
|
||||
|
||||
assert len(word2ph) == len(text) + 2
|
||||
|
||||
@@ -3,10 +3,12 @@ from __future__ import annotations
|
||||
from typing import TYPE_CHECKING, Any, Optional, Sequence, Union
|
||||
|
||||
import numpy as np
|
||||
import onnxruntime
|
||||
from numpy.typing import NDArray
|
||||
|
||||
from style_bert_vits2.constants import Languages
|
||||
from style_bert_vits2.nlp import bert_models, onnx_bert_models
|
||||
from style_bert_vits2.utils import get_onnx_device_options
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -97,32 +99,57 @@ def extract_bert_feature_onnx(
|
||||
NDArray[Any]: BERT の特徴量
|
||||
"""
|
||||
|
||||
# トークナイザーとモデルの読み込み
|
||||
tokenizer = onnx_bert_models.load_tokenizer(Languages.EN)
|
||||
inputs = tokenizer(text, return_tensors="np")
|
||||
|
||||
session = onnx_bert_models.load_model(
|
||||
language=Languages.EN,
|
||||
onnx_providers=onnx_providers,
|
||||
)
|
||||
input_names = [input.name for input in session.get_inputs()]
|
||||
output_name = session.get_outputs()[0].name
|
||||
res = session.run(
|
||||
[output_name],
|
||||
{
|
||||
"input_ids": inputs["input_ids"].astype(np.int64), # type: ignore
|
||||
"attention_mask": inputs["attention_mask"].astype(np.int64), # type: ignore
|
||||
},
|
||||
)[0]
|
||||
|
||||
# 入力テンソルの転送に使用するデバイス種別, デバイス ID, 実行オプションを取得
|
||||
device_type, device_id, run_options = get_onnx_device_options(session, onnx_providers) # fmt: skip
|
||||
|
||||
# 入力をテンソルに変換
|
||||
inputs = tokenizer(text, return_tensors="np")
|
||||
input_tensor = [
|
||||
inputs["input_ids"].astype(np.int64), # type: ignore
|
||||
inputs["attention_mask"].astype(np.int64), # type: ignore
|
||||
]
|
||||
# 推論デバイスに入力テンソルを割り当て
|
||||
## GPU 推論の場合、device_type + device_id に対応する GPU デバイスに入力テンソルが割り当てられる
|
||||
io_binding = session.io_binding()
|
||||
for name, value in zip(input_names, input_tensor):
|
||||
gpu_tensor = onnxruntime.OrtValue.ortvalue_from_numpy(
|
||||
value, device_type, device_id
|
||||
)
|
||||
io_binding.bind_ortvalue_input(name, gpu_tensor)
|
||||
# text から BERT 特徴量を抽出
|
||||
io_binding.bind_output(output_name, device_type)
|
||||
session.run_with_iobinding(io_binding, run_options=run_options)
|
||||
res = io_binding.get_outputs()[0].numpy()
|
||||
|
||||
style_res_mean = None
|
||||
if assist_text:
|
||||
# 入力をテンソルに変換
|
||||
style_inputs = tokenizer(assist_text, return_tensors="np")
|
||||
style_res = session.run(
|
||||
[output_name],
|
||||
{
|
||||
"input_ids": style_inputs["input_ids"].astype(np.int64), # type: ignore
|
||||
"attention_mask": style_inputs["attention_mask"].astype(np.int64), # type: ignore
|
||||
},
|
||||
)[0]
|
||||
style_input_tensor = [
|
||||
style_inputs["input_ids"].astype(np.int64), # type: ignore
|
||||
style_inputs["attention_mask"].astype(np.int64), # type: ignore
|
||||
]
|
||||
# 推論デバイスに入力テンソルを割り当て
|
||||
## GPU 推論の場合、device_type + device_id に対応する GPU デバイスに入力テンソルが割り当てられる
|
||||
io_binding = session.io_binding() # IOBinding は作り直す必要がある
|
||||
for name, value in zip(input_names, style_input_tensor):
|
||||
gpu_tensor = onnxruntime.OrtValue.ortvalue_from_numpy(
|
||||
value, device_type, device_id
|
||||
)
|
||||
io_binding.bind_ortvalue_input(name, gpu_tensor)
|
||||
# assist_text から BERT 特徴量を抽出
|
||||
io_binding.bind_output(output_name, device_type)
|
||||
session.run_with_iobinding(io_binding, run_options=run_options)
|
||||
style_res = io_binding.get_outputs()[0].numpy()
|
||||
style_res_mean = np.mean(style_res, axis=0)
|
||||
|
||||
assert len(word2ph) == res.shape[0], (text, res.shape[0], len(word2ph))
|
||||
|
||||
@@ -3,11 +3,13 @@ from __future__ import annotations
|
||||
from typing import TYPE_CHECKING, Any, Optional, Sequence, Union
|
||||
|
||||
import numpy as np
|
||||
import onnxruntime
|
||||
from numpy.typing import NDArray
|
||||
|
||||
from style_bert_vits2.constants import Languages
|
||||
from style_bert_vits2.nlp import bert_models, onnx_bert_models
|
||||
from style_bert_vits2.nlp.japanese.g2p import text_to_sep_kata
|
||||
from style_bert_vits2.utils import get_onnx_device_options
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -110,32 +112,57 @@ def extract_bert_feature_onnx(
|
||||
if assist_text:
|
||||
assist_text = "".join(text_to_sep_kata(assist_text, raise_yomi_error=False)[0])
|
||||
|
||||
# トークナイザーとモデルの読み込み
|
||||
tokenizer = onnx_bert_models.load_tokenizer(Languages.JP)
|
||||
inputs = tokenizer(text, return_tensors="np")
|
||||
|
||||
session = onnx_bert_models.load_model(
|
||||
language=Languages.JP,
|
||||
onnx_providers=onnx_providers,
|
||||
)
|
||||
input_names = [input.name for input in session.get_inputs()]
|
||||
output_name = session.get_outputs()[0].name
|
||||
res = session.run(
|
||||
[output_name],
|
||||
{
|
||||
"input_ids": inputs["input_ids"].astype(np.int64), # type: ignore
|
||||
"attention_mask": inputs["attention_mask"].astype(np.int64), # type: ignore
|
||||
},
|
||||
)[0]
|
||||
|
||||
# 入力テンソルの転送に使用するデバイス種別, デバイス ID, 実行オプションを取得
|
||||
device_type, device_id, run_options = get_onnx_device_options(session, onnx_providers) # fmt: skip
|
||||
|
||||
# 入力をテンソルに変換
|
||||
inputs = tokenizer(text, return_tensors="np")
|
||||
input_tensor = [
|
||||
inputs["input_ids"].astype(np.int64), # type: ignore
|
||||
inputs["attention_mask"].astype(np.int64), # type: ignore
|
||||
]
|
||||
# 推論デバイスに入力テンソルを割り当て
|
||||
## GPU 推論の場合、device_type + device_id に対応する GPU デバイスに入力テンソルが割り当てられる
|
||||
io_binding = session.io_binding()
|
||||
for name, value in zip(input_names, input_tensor):
|
||||
gpu_tensor = onnxruntime.OrtValue.ortvalue_from_numpy(
|
||||
value, device_type, device_id
|
||||
)
|
||||
io_binding.bind_ortvalue_input(name, gpu_tensor)
|
||||
# text から BERT 特徴量を抽出
|
||||
io_binding.bind_output(output_name, device_type)
|
||||
session.run_with_iobinding(io_binding, run_options=run_options)
|
||||
res = io_binding.get_outputs()[0].numpy()
|
||||
|
||||
style_res_mean = None
|
||||
if assist_text:
|
||||
# 入力をテンソルに変換
|
||||
style_inputs = tokenizer(assist_text, return_tensors="np")
|
||||
style_res = session.run(
|
||||
[output_name],
|
||||
{
|
||||
"input_ids": style_inputs["input_ids"].astype(np.int64), # type: ignore
|
||||
"attention_mask": style_inputs["attention_mask"].astype(np.int64), # type: ignore
|
||||
},
|
||||
)[0]
|
||||
style_input_tensor = [
|
||||
style_inputs["input_ids"].astype(np.int64), # type: ignore
|
||||
style_inputs["attention_mask"].astype(np.int64), # type: ignore
|
||||
]
|
||||
# 推論デバイスに入力テンソルを割り当て
|
||||
## GPU 推論の場合、device_type + device_id に対応する GPU デバイスに入力テンソルが割り当てられる
|
||||
io_binding = session.io_binding() # IOBinding は作り直す必要がある
|
||||
for name, value in zip(input_names, style_input_tensor):
|
||||
gpu_tensor = onnxruntime.OrtValue.ortvalue_from_numpy(
|
||||
value, device_type, device_id
|
||||
)
|
||||
io_binding.bind_ortvalue_input(name, gpu_tensor)
|
||||
# assist_text から BERT 特徴量を抽出
|
||||
io_binding.bind_output(output_name, device_type)
|
||||
session.run_with_iobinding(io_binding, run_options=run_options)
|
||||
style_res = io_binding.get_outputs()[0].numpy()
|
||||
style_res_mean = np.mean(style_res, axis=0)
|
||||
|
||||
assert len(word2ph) == len(text) + 2, text
|
||||
|
||||
@@ -1,9 +1,21 @@
|
||||
from typing import Any, Sequence, Union
|
||||
|
||||
import onnxruntime
|
||||
|
||||
|
||||
def torch_device_to_onnx_providers(
|
||||
device: str,
|
||||
) -> Sequence[Union[str, tuple[str, dict[str, Any]]]]:
|
||||
"""
|
||||
PyTorch のデバイス種別を ONNX の ExecutionProvider に変換する
|
||||
|
||||
Args:
|
||||
device (str): PyTorch のデバイス種別
|
||||
|
||||
Returns:
|
||||
Sequence[Union[str, tuple[str, dict[str, Any]]]]: ExecutionProvider のリスト
|
||||
"""
|
||||
|
||||
if device.startswith("cuda"):
|
||||
return [
|
||||
# cudnn_conv_algo_search を DEFAULT にすると推論速度が大幅に向上する
|
||||
@@ -15,3 +27,77 @@ def torch_device_to_onnx_providers(
|
||||
]
|
||||
else:
|
||||
return ["CPUExecutionProvider"]
|
||||
|
||||
|
||||
def get_onnx_device_options(
|
||||
onnx_session: onnxruntime.InferenceSession,
|
||||
onnx_providers: Sequence[Union[str, tuple[str, dict[str, Any]]]],
|
||||
) -> tuple[str, int, onnxruntime.RunOptions]:
|
||||
"""
|
||||
ONNX 推論時のデバイス関連のオプションを取得する
|
||||
|
||||
Args:
|
||||
onnx_session (onnxruntime.InferenceSession): 初期化済みの ONNX セッション
|
||||
onnx_providers (Sequence[Union[str, tuple[str, dict[str, Any]]]]): ExecutionProvider のリスト
|
||||
|
||||
Returns:
|
||||
tuple[str, int, onnxruntime.RunOptions]: 入力テンソルの転送に使用するデバイス種別, デバイス ID, 実行オプション
|
||||
"""
|
||||
|
||||
# 実際に推論に用いられる ExecutionProvider を取得
|
||||
first_provider = onnx_session.get_providers()[0]
|
||||
|
||||
# 入力テンソルを転送する推論デバイスを取得
|
||||
## GPU への I/O Binding は CUDA と DirectML 以外はサポートされていないので、
|
||||
## それ以外の ExecutionProvider が指定された場合は CPU に入力テンソルを転送する
|
||||
if first_provider == "CUDAExecutionProvider":
|
||||
device_type = "cuda"
|
||||
elif first_provider == "DmlExecutionProvider":
|
||||
device_type = "dml"
|
||||
else:
|
||||
device_type = "cpu"
|
||||
|
||||
# 入力テンソルを転送する GPU デバイスの ID を取得
|
||||
## ExecutionProvider に指定したオプションの中から device_id を取得し、入力テンソルの転送先として指定する
|
||||
## InferenceSession で利用するデバイス ID と入力テンソルの転送先デバイス ID は一致している必要がある
|
||||
## 本来は ExecutionProvider に指定したオプションは InferenceSession.get_provider_options() で取得できるはずだが、
|
||||
## 手元環境では DmlExecutionProvider のみ常に空の辞書が返されることがあったため、当面 onnx_providers から直接オプションを取り出している
|
||||
## CPU 推論時の device_id は 0 で固定
|
||||
device_id = 0
|
||||
onnx_providers_dict: dict[str, dict[str, Any]] = {}
|
||||
if device_type != "cpu":
|
||||
for provider in onnx_providers:
|
||||
if isinstance(provider, tuple):
|
||||
provider_name, options = provider
|
||||
onnx_providers_dict[provider_name] = options
|
||||
else:
|
||||
onnx_providers_dict[provider] = {}
|
||||
first_provider_options = onnx_providers_dict[first_provider]
|
||||
if "device_id" in first_provider_options:
|
||||
device_id = int(first_provider_options["device_id"])
|
||||
|
||||
# 推論後にメモリアリーナを縮小し、メモリを解放する
|
||||
## onnxruntime.SessionOptions の enable_cpu_mem_arena (デフォルト: True) により、デフォルトでは CPU 推論時にメモリアリーナが構築される
|
||||
## メモリアリーナを無効化すると、推論にのみ使用されたメモリは推論後にすべて解放されるが、一方パフォーマンスがかなり落ちる
|
||||
## そこでメモリアリーナを有効化した上で、推論後にメモリアリーナを縮小し、何回も音声合成するほど漸進的にメモリが消費される現象を回避する
|
||||
## この設定は GPU 推論時にもある程度効果があると思われる
|
||||
## ref: https://onnxruntime.ai/docs/get-started/with-c.html
|
||||
## ref: https://github.com/microsoft/onnxruntime/issues/9313#issuecomment-2182919186
|
||||
## ref: https://github.com/microsoft/onnxruntime/issues/11627
|
||||
## ref: https://github.com/microsoft/onnxruntime/issues/22297#issuecomment-2438629814
|
||||
## ref: https://github.com/microsoft/onnxruntime/blob/v1.20.1/onnxruntime/test/python/onnxruntime_test_python.py#L1626-L1647
|
||||
## ref: https://github.com/microsoft/onnxruntime/blob/v1.20.1/include/onnxruntime/core/session/onnxruntime_run_options_config_keys.h#L19-L27
|
||||
run_options = onnxruntime.RunOptions()
|
||||
if first_provider == "CPUExecutionProvider":
|
||||
# CPU 推論時は cpu:0 を指定
|
||||
run_options.add_run_config_entry("memory.enable_memory_arena_shrinkage", "cpu:0")
|
||||
elif first_provider == "DmlExecutionProvider":
|
||||
# DirectML 推論時はこのオプションはサポートされていないようなので、何も指定しない
|
||||
# "The registered allocator for device-id combination is not an arena based allocator: gpu:0" のようなエラーが出る…
|
||||
pass
|
||||
elif first_provider == "CUDAExecutionProvider":
|
||||
# CUDA 推論時は cpu:0;gpu:(device_id) を指定
|
||||
## 公式テストコードを読む限り、CUDA だけでなく CPU のメモリも明示的に解放した方がよいらしい
|
||||
run_options.add_run_config_entry("memory.enable_memory_arena_shrinkage", f"cpu:0;gpu:{device_id}")
|
||||
|
||||
return device_type, device_id, run_options
|
||||
|
||||
Reference in New Issue
Block a user