AIModelManager::GetModelIOTensorDim
获取已加载模型的输入输出尺寸。
函数格式
virtual AIStatus AIModelManager::GetModelIOTensorDim(const std::string& model_name, std::vector<TensorDimension>& input_tensor, std::vector<TensorDimension>& output_tensor) ;
参数说明
参数 |
说明 |
取值范围 |
---|---|---|
model_name |
模型名称。 |
- |
input_tensor |
模型输入尺寸列表。 关于TensorDimension数据类型的定义,请参见TensorDimension。 |
- |
output_tensor |
模型输出尺寸列表。 关于TensorDimension数据类型的定义,请参见TensorDimension。 |
- |
返回值
SUCCESS初始化成功/FAILED初始化失败。
示例
例如获取resnet50模型的输入输出Tensor描述如下。
若output_tensor的size不为0,则获取到的output_tensor数据会追加在原output_tensor数据之后。
input_tensor { name = "data" #输入层的name data_type = 0 #预留数据类型,暂时不用 size = 20 #内存大小,单位字节 format = 0 #预留Tensor排布格式,暂时不用 dims = {1,3,224,224} } output_tensor { name = "output_0_prob_0" #输出Tensor的name,格式为:output_{数字}_{输出节点name}_{输出节点输出索引} data_type = 0 #预留数据类型,暂时不用 size = 20 #内存大小,单位字节 format = 0 #预留Tensor排布格式,暂时不用 dims = {1,1000,1,1} }