Improve: Disable enable_cpu_mem_arena for CPU inference only to prevent excessive memory consumption during the inference session of the BERT model
This commit is contained in:
tsukumi
@@ -43,6 +43,7 @@ def load_model(
|
||||
onnx_providers: Sequence[Union[str, tuple[str, dict[str, Any]]]] = [("CPUExecutionProvider", {"arena_extend_strategy": "kSameAsRequested"})],
|
||||
cache_dir: Optional[str] = None,
|
||||
revision: str = "main",
|
||||
enable_cpu_mem_arena: bool | None = None,
|
||||
) -> onnxruntime.InferenceSession: # fmt: skip
|
||||
"""
|
||||
指定された言語の ONNX 版 BERT モデルをロードし、ロード済みの ONNX 版 BERT モデルを返す。
|
||||
@@ -61,6 +62,7 @@ def load_model(
|
||||
onnx_providers (list[str]): ONNX 推論で利用する ExecutionProvider (CPUExecutionProvider, CUDAExecutionProvider など)
|
||||
cache_dir (Optional[str]): モデルのキャッシュディレクトリ。指定しない場合はデフォルトのキャッシュディレクトリが利用される (デフォルト: None)
|
||||
revision (str): モデルの Hugging Face 上の Git リビジョン。指定しない場合は最新の main ブランチの内容が利用される (デフォルト: None)
|
||||
enable_cpu_mem_arena (bool | None): CPU 推論時にもメモリアリーナを有効化するかどうか。デフォルトでは GPU 推論時のみ有効化される (デフォルト: None)
|
||||
|
||||
Returns:
|
||||
onnxruntime.InferenceSession: ロード済みの BERT モデル
|
||||
@@ -101,16 +103,40 @@ def load_model(
|
||||
else:
|
||||
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()
|
||||
# ONNX モデルの作成時にすでに onnxsim により最適化されていることから、ロード高速化のため最適化を無効にする
|
||||
## ONNX モデルの作成時にすでに onnxsim により最適化されていることから、ロード高速化のため最適化を無効にする
|
||||
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
|
||||
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 モデルをロードし、辞書に格納して返す
|
||||
start_time = time.time()
|
||||
__loaded_models[language] = onnxruntime.InferenceSession(
|
||||
model_path,
|
||||
sess_options=sess_options,
|
||||
|
||||
@@ -141,6 +141,7 @@ class TTSModel:
|
||||
if not self.is_onnx_model:
|
||||
from style_bert_vits2.models.infer import get_net_g
|
||||
|
||||
# PyTorch モデルをロード
|
||||
self.net_g = get_net_g(
|
||||
model_path=str(self.model_path),
|
||||
version=self.hyper_parameters.version,
|
||||
@@ -197,24 +198,28 @@ class TTSModel:
|
||||
|
||||
# ONNX 推論時
|
||||
else:
|
||||
sess_options = onnxruntime.SessionOptions()
|
||||
# ONNX モデルの作成時にすでに onnxsim により最適化されていることから、ロード高速化のため最適化を無効にする
|
||||
## DmlExecutionProvider が先頭に指定されているときのみ、DirectML 推論の高速化のためすべての最適化を有効にする
|
||||
# 推論時に一番優先される ExecutionProvider の名前を取得
|
||||
assert len(self.onnx_providers) > 0
|
||||
first_provider_name = (
|
||||
self.onnx_providers[0]
|
||||
if type(self.onnx_providers[0]) is str
|
||||
else self.onnx_providers[0][0]
|
||||
)
|
||||
|
||||
# 推論セッションの設定
|
||||
sess_options = onnxruntime.SessionOptions()
|
||||
## ONNX モデルの作成時にすでに onnxsim により最適化されていることから、ロード高速化のため最適化を無効にする
|
||||
## DmlExecutionProvider が先頭に指定されているときのみ、DirectML 推論の高速化のためすべての最適化を有効にする
|
||||
if first_provider_name == "DmlExecutionProvider":
|
||||
sess_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL # fmt: skip
|
||||
else:
|
||||
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
|
||||
onnxruntime.set_default_logger_severity(3)
|
||||
|
||||
# ONNX モデルをロードし、推論セッションを初期化
|
||||
self.onnx_session = onnxruntime.InferenceSession(
|
||||
str(self.model_path),
|
||||
sess_options=sess_options,
|
||||
|
||||
@@ -53,6 +53,9 @@ def get_onnx_device_options(
|
||||
tuple[str, int, onnxruntime.RunOptions]: 入力テンソルの転送に使用するデバイス種別, デバイス ID, 実行オプション
|
||||
"""
|
||||
|
||||
# ONNX セッションに対応する SessionOptions を取得
|
||||
sess_options = onnx_session.get_session_options()
|
||||
|
||||
# 実際に推論に用いられる ExecutionProvider を取得
|
||||
first_provider = onnx_session.get_providers()[0]
|
||||
|
||||
@@ -85,11 +88,12 @@ def get_onnx_device_options(
|
||||
if "device_id" in first_provider_options:
|
||||
device_id = int(first_provider_options["device_id"])
|
||||
|
||||
# 推論後にメモリアリーナを縮小し、メモリを解放する
|
||||
## onnxruntime.SessionOptions の enable_cpu_mem_arena (デフォルト: True) により、デフォルトでは CPU 推論時にメモリアリーナが構築される
|
||||
## メモリアリーナを無効化すると、推論にのみ使用されたメモリは推論後にすべて解放されるが、一方パフォーマンスがかなり落ちる
|
||||
## そこでメモリアリーナを有効化した上で、推論後にメモリアリーナを縮小し、何回も音声合成するほど漸進的にメモリが消費される現象を回避する
|
||||
## この設定は GPU 推論時にもある程度効果があると思われる
|
||||
# 推論後にメモリアリーナを縮小し、確保されていたメモリを解放する
|
||||
## この設定により、多量の推論処理を行った際に漸進的にメモリが消費される現象を回避できる
|
||||
## 特に CPU 推論時に有効な設定で、これによりアイドル時のメモリ消費量が大幅に削減される
|
||||
## 音声合成モデルの CPU 推論時に SessionOptions の enable_cpu_mem_arena を無効化すると 推論速度が 1 秒以上低下してしまうが、
|
||||
## enable_cpu_mem_arena は有効のまま下記設定を行うと、推論速度を維持しながらメモリ消費量を削減できる
|
||||
## SessionOptions の enable_cpu_mem_arena が False の時は実行しない
|
||||
## 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
|
||||
@@ -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/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") # fmt: skip
|
||||
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}") # fmt: skip
|
||||
if sess_options.enable_cpu_mem_arena is True:
|
||||
if first_provider == "CPUExecutionProvider":
|
||||
# CPU 推論時は cpu:0 を指定
|
||||
run_options.add_run_config_entry("memory.enable_memory_arena_shrinkage", "cpu:0") # fmt: skip
|
||||
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}") # fmt: skip
|
||||
|
||||
return device_type, device_id, run_options
|
||||
|
||||
Reference in New Issue
Block a user