【环境搭建：onnx模型部署】onnxruntime-gpu安装与测试（python）

ONNX模型部署环境创建

1. onnxruntime 安装

onnx 模型在 CPU 上进行推理，在conda环境中直接使用pip安装即可

pip install onnxruntime

2. onnxruntime-gpu 安装

想要 onnx 模型在 GPU 上加速推理，需要安装 onnxruntime-gpu 。有两种思路：

• 依赖于 本地主机 上已安装的 cuda 和 cudnn 版本
• 不依赖于 本地主机 上已安装的 cuda 和 cudnn 版本

要注意：onnxruntime-gpu, cuda, cudnn三者的版本要对应，否则会报错 或 不能使用GPU推理。
onnxruntime-gpu, cuda, cudnn版本对应关系详见: 官网

2.1 方法一：onnxruntime-gpu依赖于本地主机上cuda和cudnn

• 查看已安装 cuda 和 cudnn 版本# cuda versioncat /usr/local/cuda/version.txt# cudnn versioncat /usr/local/cuda/include/cudnn.h |grep CUDNN_MAJOR -A2
• 根据 onnxruntime-gpu, cuda, cudnn 三者对应关系，安装相应的 onnxruntime-gpu 即可。## cuda==10.2## cudnn==8.0.3## onnxruntime-gpu==1.5.0 or 1.6.0pip install onnxruntime-gpu==1.6.0

2.2 方法二：onnxruntime-gpu不依赖于本地主机上cuda和cudnn

在 conda 环境中安装，不依赖于 本地主机 上已安装的 cuda 和 cudnn 版本，灵活方便。这里，先说一下已经测试通过的组合：

• python3.6, cudatoolkit10.2.89, cudnn7.6.5, onnxruntime-gpu1.4.0
• python3.8, cudatoolkit11.3.1, cudnn8.2.1, onnxruntime-gpu1.14.1

如果需要其他的版本， 可以根据 onnxruntime-gpu, cuda, cudnn 三者对应关系自行组合测试。

下面，从创建conda环境，到实现在GPU上加速onnx模型推理进行举例。

2.2.1 举例：创建onnxruntime-gpu==1.14.1的conda环境

## 创建conda环境
conda create -n torch python=3.8## 激活conda环境source activate torch
conda installpytorch==1.10.0 torchvision==0.11.0 torchaudio==0.10.0 cudatoolkit=11.3-c pytorch -c conda-forge
conda installcudnn==8.2.1
pip install onnxruntime-gpu==1.14.1
## pip install ... (根据需求，安装其他的包)

2.2.2 举例：实例测试

• 打开终端，输入 watch -n 0.1 nvidia-smi, 实时查看gpu使用情况
• 代码测试，摘取APIimport numpy as np import torch import onnxruntime MODEL_FILE ='.model.onnx' DEVICE_NAME ='cuda'if torch.cuda.is_available()else'cpu' DEVICE_INDEX =0 DEVICE=f'{DEVICE_NAME}:{DEVICE_INDEX}'# A simple model to calculate addition of two tensorsdefmodel():classModel(torch.nn.Module):def__init__(self):super(Model, self).__init__()defforward(self, x, y):return x.add(y)return Model()# Create an instance of the model and export it to ONNX graph formatdefcreate_model(type: torch.dtype = torch.float32): sample_x = torch.ones(3, dtype=type) sample_y = torch.zeros(3, dtype=type) torch.onnx.export(model(),(sample_x, sample_y), MODEL_FILE, input_names=["x","y"], output_names=["z"], dynamic_axes={"x":{0:"array_length_x"},"y":{0:"array_length_y"}})# Create an ONNX Runtime session with the provided modeldefcreate_session(model:str)-> onnxruntime.InferenceSession: providers =['CPUExecutionProvider']if torch.cuda.is_available(): providers.insert(0,'CUDAExecutionProvider')return onnxruntime.InferenceSession(model, providers=providers)# Run the model on CPU consuming and producing numpy arrays defrun(x: np.array, y: np.array)-> np.array: session = create_session(MODEL_FILE) z = session.run(["z"],{"x": x,"y": y})return z[0]# Run the model on device consuming and producing ORTValuesdefrun_with_data_on_device(x: np.array, y: np.array)-> onnxruntime.OrtValue: session = create_session(MODEL_FILE) x_ortvalue = onnxruntime.OrtValue.ortvalue_from_numpy(x, DEVICE_NAME, DEVICE_INDEX) y_ortvalue = onnxruntime.OrtValue.ortvalue_from_numpy(y, DEVICE_NAME, DEVICE_INDEX) io_binding = session.io_binding() io_binding.bind_input(name='x', device_type=x_ortvalue.device_name(), device_id=0, element_type=x.dtype, shape=x_ortvalue.shape(), buffer_ptr=x_ortvalue.data_ptr()) io_binding.bind_input(name='y', device_type=y_ortvalue.device_name(), device_id=0, element_type=y.dtype, shape=y_ortvalue.shape(), buffer_ptr=y_ortvalue.data_ptr()) io_binding.bind_output(name='z', device_type=DEVICE_NAME, device_id=DEVICE_INDEX, element_type=x.dtype, shape=x_ortvalue.shape()) session.run_with_iobinding(io_binding) z = io_binding.get_outputs()return z[0]defmain(): create_model()# print(run(x=np.float32([1.0, 2.0, 3.0]),y=np.float32([4.0, 5.0, 6.0]))) t1 = time.time()print(run(x=np.float32([1.0,2.0,3.0]),y=np.float32([4.0,5.0,6.0])))# [array([5., 7., 9.], dtype=float32)]t1 = time.time() t2 = time.time()print(run_with_data_on_device(x=np.float32([1.0,2.0,3.0,4.0,5.0]), y=np.float32([1.0,2.0,3.0,4.0,5.0])).numpy())# [ 2. 4. 6. 8. 10.] t3 = time.time()print(f'Done. ({(1E3*(t2 - t1)):.1f}ms) Inference.')print(f'Done. ({(1E3*(t3 - t2)):.1f}ms) Inference.')if __name__ =="__main__": main()