cat set env ONNX_PROVIDER , can print all/can use onnx providers

Author: xxaier

onnx/README.md

@@ -54,8 +54,8 @@ ONNX(Open Neural Network Exchange)，开放神经网络交换，用于在各种
 
 ### 目录
 
-* model/ 存放下载的模型
-* onnx/ 存放导出的 onnx，下载的 onnx 也请解压到这里
+* `model/` 存放下载的模型
+* `onnx/` 存放导出的 onnx，下载的 onnx 也请解压到这里
 
 ### 测试
 
@@ -80,9 +80,36 @@ onnxruntime 有很多版本可以选择，见[onnxruntime](https://onnxruntime.a
 
 * [./test/onnx/onnx_img.py](./test/onnx/onnx_img.py)  生成图片向量 (norm 代表归一化的向量，可用于向量搜索)
 * [./test/onnx/onnx_txt.py](./test/onnx/onnx_txt.py)  生成文本向量
-* [./test/onnx/onnx_test.py](./test/onnx/onnx_test.py) 匹配图片向量和文本向量，进行零样本分类
+* [./test/onnx/onnx_test.py](./test/onnx/onnx_test.py)
 
-可借助向量数据库，提升零样本分类的准确性，参见[ECCV 2022 | 无需下游训练，Tip-Adapter 大幅提升 CLIP 图像分类准确率](https://cloud.tencent.com/developer/article/2126102)。
+  匹配图片向量和文本向量，进行零样本分类
+
+  可借助向量数据库，提升零样本分类的准确性，参见[ECCV 2022 | 无需下游训练，Tip-Adapter 大幅提升 CLIP 图像分类准确率](https://cloud.tencent.com/developer/article/2126102)。
+* [./test/onnx/onnx_load.py](./test/onnx/onnx_load.py)
+
+  onnx 模型的加载代码，运行它可以看到当前机器可用的 onnx provider。
+
+  比如苹果 M2 芯片的笔记本上运行如下：
+
+  ```
+  ❯ ./onnx_load.py 2>/dev/null
+  all providers :
+  ['TensorrtExecutionProvider', 'CUDAExecutionProvider', 'MIGraphXExecutionProvider', 'ROCMExecutionProvider', 'OpenVINOExecutionProvider', 'DnnlExecutionProvider', 'TvmExecutionProvider', 'VitisAIExecutionProvider', 'NnapiExecutionProvider', 'CoreMLExecutionProvider', 'ArmNNExecutionProvider', 'ACLExecutionProvider', 'DmlExecutionProvider', 'RknpuExecutionProvider', 'XnnpackExecutionProvider', 'CANNExecutionProvider', 'CPUExecutionProvider']
+
+  now can use providers :
+  ['CoreMLExecutionProvider', 'CPUExecutionProvider']
+  ```
+
+  可以创建 FlagAI/onnx/.env ，设置环境变量 `ONNX_PROVIDER`，配置当前环境的 Onnx Execution Provider，方便测试对比性能。
+
+  设置的示例如下：
+
+  ```
+  ❯ cat FlagAI/onnx/.env
+  ONNX_PROVIDER=CoreMLExecutionProvider
+  ```
+
+  设置成功后，需要在 `FlagAI/onnx` 目录下运行 `direnv allow` 或者手工 `source .envrc` 让其在当前命令行中生效。
 
 #### pytorch 模型
 

onnx/test/onnx/onnx_img.py

@@ -16,7 +16,11 @@ def img2vec(img):
   from os.path import join
   img = Image.open(join(IMG_DIR, 'cat.jpg'))
 
+  img_data = transform(img)
+  import torch
+  print('img data size', torch.tensor(img_data).size())
+
   vec = img2vec(img)
   print('vec', vec)
   IMG_NORM = onnx_load('ImgNorm')
-  print('norm', IMG_NORM.run(None, {'input': transform(img)})[0])
+  print('norm', IMG_NORM.run(None, {'input': img_data})[0])

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
 
 session = onnxruntime.SessionOptions()
 option = onnxruntime.RunOptions()
@@ -11,13 +12,17 @@
 
 def onnx_load(kind):
   fp = join(ONNX_FP, f'{kind}.onnx')
-
+  provider = os.getenv('ONNX_PROVIDER')
+  providers = [provider] if provider else None
   sess = onnxruntime.InferenceSession(fp,
                                       sess_options=session,
-                                      providers=[
-                                          'TensorrtExecutionProvider',
-                                          'CUDAExecutionProvider',
-                                          'CoreMLExecutionProvider',
-                                          'CPUExecutionProvider'
-                                      ])
+                                      providers=providers)
   return sess
+
+
+if __name__ == '__main__':
+  from onnxruntime import get_all_providers
+  print('all providers :\n%s\n' % get_all_providers())
+  sess = onnx_load('Txt')
+  providers = sess.get_providers()
+  print('now can use providers :\n%s\n' % providers)