要求:已知模型输入tensor shape是[1,3,640,640],求输入tensor长度,一般我们用简单乘法1*3*640*640就可以实现,但是我看github大佬们都喜欢用std::accumulate去做,于是记录一下求解长度方法
测试环境
vs2019
onnxruntime==1.12.0
代码部分,这里以yolov8s.onnx为例子:
#include <iostream>
#include<onnxruntime_cxx_api.h>
#include <numeric>
using namespace std;
using namespace Ort;
template <typename T>
T VectorProduct(const std::vector<T>& v)
{
return std::accumulate(v.begin(), v.end(), 1, std::multiplies<T>());
};
int main()
{
const wchar_t* model_path = L"D:\\yolov8s.onnx";//模型路径
Ort::Env env;//创建env
Ort::Session session(nullptr);//创建一个空会话
Ort::SessionOptions sessionOptions{ nullptr };//创建会话配置
session = Ort::Session(env, model_path, sessionOptions);
//获取输入节点数量,名称和shape
size_t inputNodeCount= session.GetInputCount();
std::cout << "输入节点数量:" << inputNodeCount << "\n";
Ort::AllocatorWithDefaultOptions allocator;
std::shared_ptr<char> inputName = std::move(session.GetInputNameAllocated(0, allocator));
std::vector<char*> inputNodeNames;
inputNodeNames.push_back(inputName.get());
std::cout << "输入节点名称:" << inputName << "\n";
Ort::TypeInfo inputTypeInfo = session.GetInputTypeInfo(0);
auto input_tensor_info = inputTypeInfo.GetTensorTypeAndShapeInfo();
ONNXTensorElementDataType inputNodeDataType = input_tensor_info.GetElementType();
std::vector<int64_t> inputTensorShape = input_tensor_info.GetShape();
std::cout << "输入节点shape:";
for (int i = 0; i<inputTensorShape.size(); i++)
{
std::cout << inputTensorShape[i]<<" ";
}
std::cout << "\n";
//获取输出节点数量、名称和shape
size_t outputNodeCount = session.GetInputCount();
std::cout << "输出节点数量:" << outputNodeCount << "\n";
std::shared_ptr<char> outputName = std::move(session.GetOutputNameAllocated(0, allocator));
std::vector<char*> outputNodeNames;
outputNodeNames.push_back(outputName.get());
std::cout << "输出节点名称:" << outputName << "\n";
Ort::TypeInfo type_info_output0(nullptr);
type_info_output0 = session.GetOutputTypeInfo(0); //output0
auto tensor_info_output0 = type_info_output0.GetTensorTypeAndShapeInfo();
ONNXTensorElementDataType outputNodeDataType = tensor_info_output0.GetElementType();
std::vector<int64_t> outputTensorShape = tensor_info_output0.GetShape();
std::cout << "输出节点shape:";
for (int i = 0; i<outputTensorShape.size(); i++)
{
std::cout << outputTensorShape[i]<<" ";
}
std::cout << "\n";
size_t input_tensor_length = VectorProduct(inputTensorShape);
std::cout << "输入tensor的长度:" << input_tensor_length << "\n";
getchar();
}
输出结果: