【环境搭建】onnxruntime

1，介绍     onnxruntime是一个用于onnx模型推理的引擎。

2，安装 2.1 cuda,cudnn 2.2 cmake，版本>=3.13.0

sudo apt-get install libssl-dev sudo apt-get autoremove cmake # 卸载 wget https://cmake.org/files/v3.17/cmake-3.17.3.tar.gz tar -xf cmake-3.17.3 cd cmake-3.17.3 ./bootstrap make -j 8 sudo make install cmake -version cmake version 3.17.3 CMake suite maintained and supported by Kitware (kitware.com/cmake).

2.3 tensorrt 2.4 onnxruntime

conda activate py36 # 切换虚拟环境 git clone https://github.com/microsoft/onnxruntime.git cd onnxruntime git submodule sync git submodule update --init --recursive ./build.sh \ --use_cuda \ --cuda_version=11.0 \ --cuda_home=/usr/local/cuda \ --cudnn_home=/usr/local/cuda \ --use_tensorrt --tensorrt_home=$TENSORRT_ROOT \ --build_shared_lib --enable_pybind \ --build_wheel --update --build pip build/Linux/Debug/dist/onnxruntime_gpu_tensorrt-1.6.0-cp36-cp36m-linux_x86_64.whl # 获取指定版本 git clone -b v1.6.0 https://github.com/microsoft/onnxruntime.git cd onnxruntime git checkout -b v1.6.0 git submodule sync git submodule update --init --recursive

3,pip 安装

pip install onnxruntime-gpu==1.6.0 onnx==1.9.0 onnxconverter_common==1.6.0 # cuda 10.2 pip install onnxruntime-gpu==1.8.1 onnx==1.9.0 onnxconverter_common==1.8.1 # cuda 11.0 import onnxruntime onnxruntime.get_available_providers() sess = onnxruntime.InferenceSession(onnx_path) sess.set_providers(['CUDAExecutionProvider'], [ {'device_id': 0}]) result = sess.run([output_name], {input_name:x}) # x is input

3，onnxruntime and pytorch inference

# -*- coding: utf-8 -*- import torch from torchvision import models import onnxruntime import numpy as np model = models.resnet18(pretrained=True) model.eval().cuda() ## pytorch x = torch.rand(2,3,224,224).cuda() out_pt = model(x) print(out_pt.size()) # onnx onnx_path = "resnet18.onnx" dynamic_axes = {'input': {0: 'batch_size'}} torch.onnx.export(model, x, onnx_path, export_params=True, opset_version=11, do_constant_folding=True, input_names=['input'], dynamic_axes=dynamic_axes) sess = onnxruntime.InferenceSession(onnx_path) sess.set_providers(['CUDAExecutionProvider'], [ {'device_id': 0}]) input_name = sess.get_inputs()[0].name output_name = sess.get_outputs()[0].name # onnxruntime inference xx = x.cpu().numpy().astype(np.float32) result = sess.run([output_name], {input_name:xx}) # x is input print(result[0].shape) # MSE mse = np.mean((out_pt.data.cpu().numpy() - result[0]) ** 2) print(mse)

输出结果

torch.Size([2, 1000]) (2, 1000) 9.275762e-13