|
前几天使用了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加速。
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
|