ONNXRuntime (Python) GPU 部署配置记录

0. 前言

最近偶尔捣鼓了一下onnxruntime-gpu(python版本)的服务端部署，于是打算简单记录一下一些关键步骤，免得以后忘了。确实，有些时候我们并不全是需要把模型转成MNN/ncnn/TNN后走移动端部署那套，服务端的部署也是个很重要的场景。比较常用的服务端部署方案包括tensorrt、onnxruntime-gpu等等。onnxruntime-gpu版本可以说是一个非常简单易用的框架，因为通常用pytorch训练的模型，在部署时，会首先转换成onnx，而onnxruntime和onnx又是有着同一个爸爸，无疑，在op的支持上肯定是最好的。采用onnxruntime来部署onnx模型，不需要经过任何二次的模型转换。当然，不同的推理引擎会有不同优势，这里就不做对比了，这篇短文主要记录一下onnxruntime-gpu版本配置的一些主要步骤。

1. 基础镜像选择

这一步很重要，只有选择了正确的基础镜像，你才能顺利地使用onnxruntime-gpu版本。onnxruntime-gpu版本依赖于cuda库，因此你选择的镜像中必须要包含cuda库（动态库），否则就算能顺利安装onnxruntime-gpu版本，也无法真正地使用到GPU。进入docker hub 搜索pytorch的镜像，我们看到有很多选择，比如1.8.0版本的，就有cuda10.2、cuda11.1的devel和runtime版本。看到下面这张图。你需要选择的是带有cuda库的devel版本。

下载合适的镜像（看你驱动版本的高低进行选择）

docker pull pytorch/pytorch:1.8.0-cuda11.1-cudnn8-devel # 如果你的nvidia驱动较高
docker pull pytorch/pytorch:1.8.1-cuda10.2-cudnn7-devel # 如果你的nvidia驱动较低

2. 关于runtime和devel的补充

【更新：2022/01/18】以上第1节是文章关于runtime和devel选择的原文。经【小白】大佬提醒，这段理解存在一定的错误。因此这里补充一下学到的新知识~ 选择devel的原因之一是之前发现用runtime版本出现无法使用的问题，换成devel后就正常的，之二是，查到一些相关的教程里面提到需要export导出CUDA库所在的路径到PATH，那么就需要devel版本里面的cuda库。经【小白】大佬指出，runtime版本有问题，“大概是依赖版本问题，这东西重来不提示版本不对，靠自己把握。devel用来编译，runtime用来发版本。” 官方的docker构建脚本在：

onnxruntime: Dockerfile.cuda​github.com/microsoft/onnxruntime/blob/master/dockerfiles/Dockerfile.cuda

因此，如果你只是使用pip进行安装，那么runtime版本和devel版本的pytorch镜像都是可行的。如果你需要自己从源码进行构建，那么就需要devel版本。

3. 正确启动Docker镜像

注意，首先你需要通过nvidia-docker启动并登录pytorch1.8.0的容器，使用docker启动将无法获取GPU信息。

GPU_ID=0
CONTAINER_NAME=onnxruntime_gpu_test

nvidia-docker run -idt -p ${PORT2}:${PORT1} \  # 指定你想设置的映射端口；idt中的d表示后台运行，去掉d表示不后台运行
  -v ${SERVER_DIR}:${CONTAINER_DIR} \  # 挂载共享目录 如果需要 不需要的可以去掉这句
  --shm-size=16gb --env NVIDIA_VISIBLE_DEVICES="${GPU_ID}" \  # 指定GPU_ID
  --name="${CONTAINER_NAME}" pytorch/pytorch:1.8.0-cuda11.1-cudnn8-devel  # 指定后台启动的镜像

安装nvidia-docker有很多参考，这里就不重复了。比如

Ubuntu16.04安装nvidia-docker24 赞同 · 10 评论文章

需要注意的是，首先你得确保宿主机的显卡驱动等是正常的，否则在容器内无法使用GPU，比如

nvidia-smi  # 看看是否识别到GPU
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 440.118.02   Driver Version: 440.118.02   CUDA Version: 10.2     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|===============================+======================+======================|
|   0  Tesla P40           Off  | 00000000:88:00.0 Off |                    0 |
| N/A   26C    P8    10W / 250W |     10MiB / 22919MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+

关于nvidia驱动的安装，也有很多文章可以参考，也不重复了。比如

Ubuntu 18.04 安装 NVIDIA 显卡驱动517 赞同 · 203 评论文章正在上传…重新上传取消

4. 安装onnxruntime-gpu版本

进入容器后，通过pip安装gpu版本的onnxruntime。如果前面的步骤都没有问题的话，这个步骤应该会很顺利。

pip install onnxruntime-gpu # 安装GPU版本

先确认下onnxruntime是否真的能用到你的GPU，如果能获取 TensorrtExecutionProvider 和 CUDAExecutionProvider，那么恭喜你！一切正常！你可以愉快地进行GPU推理部署了。

root@xxx:/workspace# python
Python 3.8.8 (default, Feb 24 2021, 21:46:12)
[GCC 7.3.0] :: Anaconda, Inc. on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import onnxruntime
>>> onnxruntime.get_device()
'GPU'
>>> onnxruntime.get_available_providers()
['TensorrtExecutionProvider', 'CUDAExecutionProvider', 'CPUExecutionProvider']
>>> exit()

如果你发现无法使用GPU，可以尝试增加以下设置。

export PATH=/usr/local/cuda-11.1/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/cuda-11.1/lib64:$LD_LIBRARY_PATH

5. TensorrtExecutionProvider何时有效

【更新：2022/01/18】以上第4节是文章关于 TensorrtExecutionProvider 和 CUDAExecutionProvider使用的原文。经【小白】大佬提醒，TensorrtExecutionProvider 并不一定会被执行，官方文档有提到，通过pip安装的onnxruntime-gpu，只能用到 CUDAExecutionProvider 进行加速。只有从源码编译的onnxruntime-gpu 才能用TensorrtExecutionProvider进行加速（这个我还没试过，之后有时间再来填源码编译的坑~）。官方文档如下：

Official Python packages on Pypi only support the default CPU (MLAS) 
and default GPU (CUDA) execution providers. For other execution providers,
you need to build from source. The recommended instructions build the 
wheel with debug info in parallel. Tune performance

我们也可以来看看onnxruntime python的一小段源码：

    def _create_inference_session(self, providers, provider_options):
        available_providers = C.get_available_providers()

        # validate providers and provider_options before other initialization
        providers, provider_options = check_and_normalize_provider_args(providers,
                                                                        provider_options,
                                                                        available_providers)

        # Tensorrt can fall back to CUDA. All others fall back to CPU.
        if 'TensorrtExecutionProvider' in available_providers:
            self._fallback_providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
        else:
            self._fallback_providers = ['CPUExecutionProvider']

以及，一段初始化调用_fallback_providers的逻辑。

        try:
            self._create_inference_session(providers, provider_options)
        except RuntimeError:
            if self._enable_fallback:
                print("EP Error using {}".format(self._providers))
                print("Falling back to {} and retrying.".format(self._fallback_providers))
                self._create_inference_session(self._fallback_providers, None)
                # Fallback only once.
                self.disable_fallback()
            else:
                raise

可以看到，TensorrtExecutionProvider可以回退到CUDAExecutionProvider。如果你设置了TensorrtExecutionProvider，那么onnxruntime会首先尝试使用包含TensorrtExecutionProvider的providers构建InferenceSession，如果失败的话，则回退到只使用['CUDAExecutionProvider','CPUExecutionProvider']，并且打印两句log：

print("EP Error using {}".format(self._providers))
print("Falling back to {} and retrying.".format(self._fallback_providers))

emmmm......，不过在我的使用中(pip 安装的 onnxruntime-gpu最新版)，并没有打印这两句log，所以是用到了TensorrtExecutionProvider？？？不太清楚是不是新版本优化了这个问题。

6. 使用举例

用法其实很简单，只要在新建InferenceSession的时候，加入 TensorrtExecutionProvider 和 CUDAExecutionProvider 就可以了。以下这行代码，无论是CPU还是GPU部署，都是通用的。在跑推理的时候，检查一下显存占用有没有起来，有就说明一切正常。

self.session = onnxruntime.InferenceSession(
     "YOUR-ONNX-MODEL-PATH", 
     providers=onnxruntime.get_available_providers()
)

如果 要可变可以控制providers，例如：

if device == 'cpu':
    providers = ['CPUExecutionProvider']
elif device == 'gpu':
    providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']

ONNXRuntime-gpu可能无法使用GPU，与CUDA的关系参考：CUDA - onnxruntime

6. 效率对比

简单罗列一下我使用onnxruntime-gpu推理的性能（只是和cpu简单对比下，不是很严谨，暂时没有和其他推理引擎作对比）

CPU	GPU	次数	提速
2637ms（16thread）	131ms	100	15-20x