cat set env ONNX_PROVIDER

FlagAI-Open · Jun 3, 2023 · b83e558 · b83e558
1 parent d6f2f5e
commit b83e558
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/onnx/test/onnx/onnx_load.py b/onnx/test/onnx/onnx_load.py
@@ -3,6 +3,7 @@
 import onnxruntime
 from misc.config import ONNX_FP
 from os.path import join
+import os
 from onnxruntime import get_all_providers
 
 session = onnxruntime.SessionOptions()
@@ -12,10 +13,12 @@
 
 def onnx_load(kind):
   fp = join(ONNX_FP, f'{kind}.onnx')
-
+  # 可以在 FlagAI/onnx/.env 中设置环境变量 ONNX_PROVIDER 避免 UserWarning Specified provider xxx is not in available
+  provider = os.getenv('ONNX_PROVIDER')
+  providers = [provider] if provider else get_all_providers()
   sess = onnxruntime.InferenceSession(fp,
                                       sess_options=session,
-                                      providers=get_all_providers())
+                                      providers=providers)
   return sess