tuna2134/sbv2-v2
1
0
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases Wiki Activity

Improve: Disable enable_cpu_mem_arena for CPU inference only to prevent excessive memory consumption during the inference session of the BERT model

Browse Source
This commit is contained in:
tsukumi
2024-12-22 09:17:31 +09:00
parent 3c218c7087
commit 810ca43615
3 changed files with 61 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

34
style_bert_vits2/nlp/onnx_bert_models.py
View File

@@ -43,6 +43,7 @@ def load_model(
onnx_providers: Sequence[Union[str, tuple[str, dict[str, Any]]]] = [("CPUExecutionProvider", {"arena_extend_strategy": "kSameAsRequested"})], onnx_providers: Sequence[Union[str, tuple[str, dict[str, Any]]]] = [("CPUExecutionProvider", {"arena_extend_strategy": "kSameAsRequested"})],
cache_dir: Optional[str] = None, cache_dir: Optional[str] = None,
revision: str = "main", revision: str = "main",
enable_cpu_mem_arena: bool | None = None,
) -> onnxruntime.InferenceSession: # fmt: skip ) -> onnxruntime.InferenceSession: # fmt: skip
""" """
指定された言語の ONNX 版 BERT モデルをロードし、ロード済みの ONNX 版 BERT モデルを返す。 指定された言語の ONNX 版 BERT モデルをロードし、ロード済みの ONNX 版 BERT モデルを返す。
@@ -61,6 +62,7 @@ def load_model(
onnx_providers (list[str]): ONNX 推論で利用する ExecutionProvider (CPUExecutionProvider, CUDAExecutionProvider など) onnx_providers (list[str]): ONNX 推論で利用する ExecutionProvider (CPUExecutionProvider, CUDAExecutionProvider など)
cache_dir (Optional[str]): モデルのキャッシュディレクトリ。指定しない場合はデフォルトのキャッシュディレクトリが利用される (デフォルト: None) cache_dir (Optional[str]): モデルのキャッシュディレクトリ。指定しない場合はデフォルトのキャッシュディレクトリが利用される (デフォルト: None)
revision (str): モデルの Hugging Face 上の Git リビジョン。指定しない場合は最新の main ブランチの内容が利用される (デフォルト: None) revision (str): モデルの Hugging Face 上の Git リビジョン。指定しない場合は最新の main ブランチの内容が利用される (デフォルト: None)
enable_cpu_mem_arena (bool | None): CPU 推論時にもメモリアリーナを有効化するかどうか。デフォルトでは GPU 推論時のみ有効化される (デフォルト: None)
Returns: Returns:
onnxruntime.InferenceSession: ロード済みの BERT モデル onnxruntime.InferenceSession: ロード済みの BERT モデル
@@ -101,16 +103,40 @@ def load_model(
else: else:
model_path = Path(pretrained_model_name_or_path).resolve() / "model_fp16.onnx" model_path = Path(pretrained_model_name_or_path).resolve() / "model_fp16.onnx"
start_time = time.time() # 推論時に一番優先される ExecutionProvider の名前を取得
assert len(onnx_providers) > 0
first_provider_name = (
onnx_providers[0]
if type(onnx_providers[0]) is str
else onnx_providers[0][0]
)
# 推論セッションの設定
sess_options = onnxruntime.SessionOptions() sess_options = onnxruntime.SessionOptions()
# ONNX モデルの作成時にすでに onnxsim により最適化されていることから、ロード高速化のため最適化を無効にする ## ONNX モデルの作成時にすでに onnxsim により最適化されていることから、ロード高速化のため最適化を無効にする
sess_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL # fmt: skip sess_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL # fmt: skip
# エラー以外のログを出力しない ## エラー以外のログを出力しない
# 本来は log_severity_level = 3 だけで効くはずだが、なぜか抑制できないので set_default_logger_severity() も呼び出している ## 本来は log_severity_level = 3 だけで効くはずだが、なぜか CUDA 系のログが抑制できないので set_default_logger_severity() も呼び出している
sess_options.log_severity_level = 3 sess_options.log_severity_level = 3
onnxruntime.set_default_logger_severity(3) onnxruntime.set_default_logger_severity(3)
# CPU 推論時のみ enable_cpu_mem_arena を無効化し、BERT モデルの推論セッションより富豪的なメモリ消費を防止する
## 既に RunOptions の memory.enable_memory_arena_shrinkage や、ProviderOptions の "arena_extend_strategy": "kSameAsRequested" を指定して
## InferenceSession が構築するメモリアリーナを推論後に縮小するよう構成し、メモリアリーナによるメモリ消費量が漸進的に増加することを防いでいる
## しかし、CPU 推論時の BERT モデルに関しては入力長や入力内容次第では依然大量のメモリが確保される傾向にあるため、CPU 推論時のみメモリアリーナ自体を無効化する
## BERT 特徴量の抽出処理が 0.数秒遅くなるトレードオフがあるが、元々 CPU 推論は CUDA 推論よりかなり遅いこと、
## BERT 特徴量の抽出処理自体は音声合成処理よりも遥かに軽量なこと、低メモリ環境での OOM エラー回避の観点から有益だと判断した
## メモリアリーナを無効化することで、若干の速度低下と引き換えに、多量の推論処理を行ってもメモリリークのような挙動が発生しなくなる
## なお CUDA 推論時は独自に VRAM 管理が行われているようで、CPU 推論時のように過剰に VRAM が消費されることはない
## 明示的に enable_cpu_mem_arena が指定されている場合は、指定された値を利用する
if enable_cpu_mem_arena is not None:
sess_options.enable_cpu_mem_arena = enable_cpu_mem_arena
## 明示的に enable_cpu_mem_arena が指定されていない場合は、推論セッションが CPUExecutionProvider の場合のみメモリアリーナを無効化する
elif first_provider_name == "CPUExecutionProvider":
sess_options.enable_cpu_mem_arena = False
# BERT モデルをロードし、辞書に格納して返す # BERT モデルをロードし、辞書に格納して返す
start_time = time.time()
__loaded_models[language] = onnxruntime.InferenceSession( __loaded_models[language] = onnxruntime.InferenceSession(
model_path, model_path,
sess_options=sess_options, sess_options=sess_options,

15
style_bert_vits2/tts_model.py
View File

@@ -141,6 +141,7 @@ class TTSModel:
if not self.is_onnx_model: if not self.is_onnx_model:
from style_bert_vits2.models.infer import get_net_g from style_bert_vits2.models.infer import get_net_g
# PyTorch モデルをロード
self.net_g = get_net_g( self.net_g = get_net_g(
model_path=str(self.model_path), model_path=str(self.model_path),
version=self.hyper_parameters.version, version=self.hyper_parameters.version,
@@ -197,24 +198,28 @@ class TTSModel:
# ONNX 推論時 # ONNX 推論時
else: else:
sess_options = onnxruntime.SessionOptions() # 推論時に一番優先される ExecutionProvider の名前を取得
# ONNX モデルの作成時にすでに onnxsim により最適化されていることから、ロード高速化のため最適化を無効にする
## DmlExecutionProvider が先頭に指定されているときのみ、DirectML 推論の高速化のためすべての最適化を有効にする
assert len(self.onnx_providers) > 0 assert len(self.onnx_providers) > 0
first_provider_name = ( first_provider_name = (
self.onnx_providers[0] self.onnx_providers[0]
if type(self.onnx_providers[0]) is str if type(self.onnx_providers[0]) is str
else self.onnx_providers[0][0] else self.onnx_providers[0][0]
) )
# 推論セッションの設定
sess_options = onnxruntime.SessionOptions()
## ONNX モデルの作成時にすでに onnxsim により最適化されていることから、ロード高速化のため最適化を無効にする
## DmlExecutionProvider が先頭に指定されているときのみ、DirectML 推論の高速化のためすべての最適化を有効にする
if first_provider_name == "DmlExecutionProvider": if first_provider_name == "DmlExecutionProvider":
sess_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL # fmt: skip sess_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL # fmt: skip
else: else:
sess_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL # fmt: skip sess_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL # fmt: skip
# エラー以外のログを出力しない ## エラー以外のログを出力しない
# 本来は log_severity_level = 3 だけで効くはずだが、なぜか抑制できないので set_default_logger_severity() も呼び出している ## 本来は log_severity_level = 3 だけで効くはずだが、なぜか CUDA 系のログが抑制できないので set_default_logger_severity() も呼び出している
sess_options.log_severity_level = 3 sess_options.log_severity_level = 3
onnxruntime.set_default_logger_severity(3) onnxruntime.set_default_logger_severity(3)
# ONNX モデルをロードし、推論セッションを初期化
self.onnx_session = onnxruntime.InferenceSession( self.onnx_session = onnxruntime.InferenceSession(
str(self.model_path), str(self.model_path),
sess_options=sess_options, sess_options=sess_options,

37
style_bert_vits2/utils/__init__.py
View File

@@ -53,6 +53,9 @@ def get_onnx_device_options(
tuple[str, int, onnxruntime.RunOptions]: 入力テンソルの転送に使用するデバイス種別, デバイス ID, 実行オプション tuple[str, int, onnxruntime.RunOptions]: 入力テンソルの転送に使用するデバイス種別, デバイス ID, 実行オプション
""" """
# ONNX セッションに対応する SessionOptions を取得
sess_options = onnx_session.get_session_options()
# 実際に推論に用いられる ExecutionProvider を取得 # 実際に推論に用いられる ExecutionProvider を取得
first_provider = onnx_session.get_providers()[0] first_provider = onnx_session.get_providers()[0]
@@ -85,11 +88,12 @@ def get_onnx_device_options(
if "device_id" in first_provider_options: if "device_id" in first_provider_options:
device_id = int(first_provider_options["device_id"]) device_id = int(first_provider_options["device_id"])
# 推論後にメモリアリーナを縮小し、メモリを解放する # 推論後にメモリアリーナを縮小し、確保されていたメモリを解放する
## onnxruntime.SessionOptions の enable_cpu_mem_arena (デフォルト: True) により、デフォルトでは CPU 推論時にメモリアリーナが構築される ## この設定により、多量の推論処理を行った際に漸進的にメモリが消費される現象を回避できる
## メモリアリーナを無効化すると、推論にのみ使用されたメモリは推論後にすべて解放されるが、一方パフォーマンスがかなり落ちる ## 特に CPU 推論時に有効な設定で、これによりアイドル時のメモリ消費量が大幅に削減される
## そこでメモリアリーナを有効化した上で、推論後にメモリアリーナを縮小し、何回も音声合成するほど漸進的にメモリが消費される現象を回避する ## 音声合成モデルの CPU 推論時に SessionOptions の enable_cpu_mem_arena を無効化すると 推論速度が 1 秒以上低下してしまうが、
## この設定は GPU 推論時にもある程度効果があると思われる ## enable_cpu_mem_arena は有効のまま下記設定を行うと、推論速度を維持しながらメモリ消費量を削減できる
## SessionOptions の enable_cpu_mem_arena が False の時は実行しない
## ref: https://onnxruntime.ai/docs/get-started/with-c.html ## 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/9313#issuecomment-2182919186
## ref: https://github.com/microsoft/onnxruntime/issues/11627 ## ref: https://github.com/microsoft/onnxruntime/issues/11627
@@ -97,16 +101,17 @@ def get_onnx_device_options(
## 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/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 ## 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() run_options = onnxruntime.RunOptions()
if first_provider == "CPUExecutionProvider": if sess_options.enable_cpu_mem_arena is True:
# CPU 推論時は cpu:0 を指定 if first_provider == "CPUExecutionProvider":
run_options.add_run_config_entry("memory.enable_memory_arena_shrinkage", "cpu:0") # fmt: skip # CPU 推論時は cpu:0 を指定
elif first_provider == "DmlExecutionProvider": run_options.add_run_config_entry("memory.enable_memory_arena_shrinkage", "cpu:0") # fmt: skip
# DirectML 推論時はこのオプションはサポートされていないようなので、何も指定しない elif first_provider == "DmlExecutionProvider":
# "The registered allocator for device-id combination is not an arena based allocator: gpu:0" のようなエラーが出る… # DirectML 推論時はこのオプションはサポートされていないようなので、何も指定しない
pass # "The registered allocator for device-id combination is not an arena based allocator: gpu:0" のようなエラーが出る…
elif first_provider == "CUDAExecutionProvider": pass
# CUDA 推論時は cpu:0;gpu:(device_id) を指定 elif first_provider == "CUDAExecutionProvider":
## 公式テストコードを読む限り、CUDA だけでなく CPU のメモリも明示的に解放した方がよいらしい # CUDA 推論時は cpu:0;gpu:(device_id) を指定
run_options.add_run_config_entry("memory.enable_memory_arena_shrinkage", f"cpu:0;gpu:{device_id}") # fmt: skip ## 公式テストコードを読む限り、CUDA だけでなく CPU のメモリも明示的に解放した方がよいらしい
run_options.add_run_config_entry("memory.enable_memory_arena_shrinkage", f"cpu:0;gpu:{device_id}") # fmt: skip
return device_type, device_id, run_options return device_type, device_id, run_options