onnxruntime (C++/CUDA) 编译安装及部署

前几天使用了LibTorch对模型进行C++转换和测试，发现速度比原始Python的Pytorch模型提升了将近2倍。现在尝试以下另一种跨平台的模型转换方式——Onnx，可实现跨X86/ARM架构的迁移应用。

本文主要介绍C++版本的onnxruntime使用，Python的操作较容易就不再提及了。

一、克隆及编译

git clone --recursive https://github.com/Microsoft/onnxruntime
cd onnxruntime/
git checkout v1.8.0

这里建议checkout到旧tag，否则容易因为版本过新而编译失败，比如Cmake版本要求过高、CUDA版本不匹配等问题。若跟随网上其他教程，大概率会因为版本过新而导致后续编译失败。

接下来编译：

./build.sh --skip_tests --use_cuda --config Release --build_shared_lib --parallel --cuda_home /usr/local/cuda-11.3 --cudnn_home /usr/local/cuda-11.3

其中的use_cuda表示你要使用CUDA的onnxruntime，cuda_home和cudnn_home均指向你的CUDA安装目录即可。

最后就编译成功了：

[100%] Linking CXX executable onnxruntime_test_all
[100%] Built target onnxruntime_test_all
[100%] Linking CUDA shared module libonnxruntime_providers_cuda.so
[100%] Built target onnxruntime_providers_cuda
2022-03-15 13:49:03,260 util.run [DEBUG] - Subprocess completed. Return code: 0
2022-03-15 13:49:03,260 build [INFO] - Build complete

二、onnxruntime使用实例

2.1 onnx模型保存

重点关注example的设置（输入）和如何export出onnx模型。

import numpy as np
from modules.feature_extracter_without_delta_layer import featureExtracter


checkpoint = torch.load("./amodel.pth.tar")
amodel = featureExtracter(channels=1)
amodel.load_state_dict(checkpoint['state_dict']) 
amodel.cuda()
amodel.eval()

example = torch.rand(1, 1, 32, 900)
example = example.cuda()
torch.onnx.export(amodel,
                 (example),
                 'overlapTransformer.onnx',
                 input_names = ['input'],
                 output_names = ['output'],
                 opset_version=11,
                 verbose = True)

2.2 onnx模型的C++调用

其中的cuda_options与AppendExecutionProvider_CUDA的联合使用保证了CUDA加速。

#include <iostream>
#include <vector>
#include <chrono>
#include <string>
#include <vector>
#include <onnxruntime_cxx_api.h>

using namespace std;


int main()
{
    Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "test");
    Ort::SessionOptions session_options;

    OrtCUDAProviderOptions cuda_options{
          0,
          OrtCudnnConvAlgoSearch::EXHAUSTIVE,
          std::numeric_limits<size_t>::max(),
          0,
          true
      };

    session_options.AppendExecutionProvider_CUDA(cuda_options);
    const char* model_path = "./overlapTransformer.onnx";

    int width = 900;
    int height = 32;
    int len_arr = width*height;
    float virtual_image[len_arr];
    for (int i=0; i<height; i++)
        for (int j=0; j<width; j++)
        {
            virtual_image[int(i*width+j)] = 1; // range
        }


    Ort::Session session(env, model_path, session_options);
    // print model input layer (node names, types, shape etc.)
    Ort::AllocatorWithDefaultOptions allocator;

    // print number of model input nodes
    size_t num_input_nodes = session.GetInputCount();
    std::vector<const char*> input_node_names = {"input"};
    std::vector<const char*> output_node_names = {"output"};

    std::vector<int64_t> input_node_dims = {1, 1, 32, 900};
    size_t input_tensor_size = 32 * 900;
    std::vector<float> input_tensor_values(input_tensor_size);
    for (unsigned int i = 0; i < input_tensor_size; i++)
        input_tensor_values[i] = float(virtual_image[i]);
    // create input tensor object from data values ！！！！！！！！！！
    auto memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);

    Ort::Value input_tensor = Ort::Value::CreateTensor<float>(memory_info, input_tensor_values.data(),
                                                            input_tensor_size, input_node_dims.data(), 4);

    std::vector<Ort::Value> ort_inputs;
    ort_inputs.push_back(std::move(input_tensor));

    auto output_tensors = session.Run(Ort::RunOptions{nullptr}, input_node_names.data(), ort_inputs.data(),
                                    ort_inputs.size(), output_node_names.data(), 1);

    float* floatarr = output_tensors[0].GetTensorMutableData<float>();

    for (int i=0; i<256; i++)
    {
        std::cout<<floatarr[i]<<std::endl;
    }

    return 0;
}

相关教程，可供参考：
https://github.com/microsoft/onnxruntime
https://blog.csdn.net/znsoft/article/details/114583048
https://ask.csdn.net/questions/7619412
https://blog.csdn.net/XCCCCZ/article/details/110356437