【环境搭建:onnx模型部署】onnxruntime-gpu安装与测试(python)

ONNX模型部署环境创建

  • 1. onnxruntime 安装
  • 2. onnxruntime-gpu 安装
    • 2.1 方法一:onnxruntime-gpu依赖于本地主机上cuda和cudnn
    • 2.2 方法二:onnxruntime-gpu不依赖于本地主机上cuda和cudnn
      • 2.2.1 举例:创建onnxruntime-gpu==1.14.1的conda环境
      • 2.2.2 举例:实例测试

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 version
    cat /usr/local/cuda/version.txt
    
    # cudnn version
    cat /usr/local/cuda/include/cudnn.h | grep CUDNN_MAJOR -A 2
    
  • 根据 onnxruntime-gpu, cuda, cudnn 三者对应关系,安装相应的 onnxruntime-gpu 即可。

    ## cuda==10.2
    ## cudnn==8.0.3
    ## onnxruntime-gpu==1.5.0 or 1.6.0
    pip 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 install pytorch==1.10.0 torchvision==0.11.0 torchaudio==0.10.0 cudatoolkit=11.3 -c pytorch -c conda-forge
conda install cudnn==8.2.1
pip install onnxruntime-gpu==1.14.1
## pip install ... (根据需求,安装其他的包)

2.2.2 举例:实例测试

  • 打开终端,输入 watch -n 0.1 nvidia-smi, 实时查看gpu使用情况

    【环境搭建:onnx模型部署】onnxruntime-gpu安装与测试(python)_第1张图片

  • 代码测试,摘取API

      import 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 tensors
      def model():
          class Model(torch.nn.Module):
              def __init__(self):
                  super(Model, self).__init__()
      
              def forward(self, x, y):
                  return x.add(y)
      
          return Model()
      
      # Create an instance of the model and export it to ONNX graph format
      def create_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 model
      def create_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 
      def run(x: np.array, y: np.array) -> np.array:
          session = create_session(MODEL_FILE)
          z = session.run(["z"], {"x": x, "y": y})
          return z[0]   
    
      def main():
          create_model()
          print(run(x=np.float32([1.0, 2.0, 3.0]),y=np.float32([4.0, 5.0, 6.0])))
      
      if __name__ == "__main__":
          main()   
    

你可能感兴趣的:(python,深度学习,pytorch)