cat set env ONNX_PROVIDER

Author: xxaier

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