1. onnxruntime 安裝

onnx 模型在 CPU 上進(jìn)行推理，在conda環(huán)境中直接使用pip安裝即可

pip install onnxruntime

2. onnxruntime-gpu 安裝

想要 onnx 模型在 GPU 上加速推理，需要安裝 onnxruntime-gpu 。有兩種思路：

• 依賴于 本地主機 上已安裝的 cuda 和 cudnn 版本
• 不依賴于 本地主機 上已安裝的 cuda 和 cudnn 版本

要注意：onnxruntime-gpu, cuda, cudnn三者的版本要對應(yīng)，否則會報錯 或 不能使用GPU推理。
onnxruntime-gpu, cuda, cudnn版本對應(yīng)關(guān)系詳見: 官網(wǎng)

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
  

• 根據(jù) onnxruntime-gpu, cuda, cudnn 三者對應(yīng)關(guān)系，安裝相應(yīng)的 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 環(huán)境中安裝，不依賴于 本地主機 上已安裝的 cuda 和 cudnn 版本，靈活方便。這里，先說一下已經(jīng)測試通過的組合：

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

如果需要其他的版本， 可以根據(jù) onnxruntime-gpu, cuda, cudnn 三者對應(yīng)關(guān)系自行組合測試。

下面，從創(chuàng)建conda環(huán)境，到實現(xiàn)在GPU上加速onnx模型推理進(jìn)行舉例。

2.2.1 舉例：創(chuàng)建onnxruntime-gpu==1.14.1的conda環(huán)境

## 創(chuàng)建conda環(huán)境
conda create -n torch python=3.8

## 激活conda環(huán)境
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 ... (根據(jù)需求，安裝其他的包)

2.2.2 舉例：實例測試

• 打開終端，輸入 watch -n 0.1 nvidia-smi, 實時查看gpu使用情況

  

• 代碼測試，摘取API文章來源地址http://www.zghlxwxcb.cn/news/detail-468755.html

    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]   
  
    # Run the model on device consuming and producing ORTValues
    def run_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]
  
    def main():
        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()   
  

到了這里，關(guān)于【環(huán)境搭建：onnx模型部署】onnxruntime-gpu安裝與測試（python）的文章就介紹完了。如果您還想了解更多內(nèi)容，請在右上角搜索TOY模板網(wǎng)以前的文章或繼續(xù)瀏覽下面的相關(guān)文章，希望大家以后多多支持TOY模板網(wǎng)！