export_engine

函数功能

将原始TorchScript模型进行编译优化生成可在昇腾处理上加速推理的离线OM模型。

函数原型

std::string export_engine(const torch::jit::Module& module, const std::string& method_name, const CompileSpec& info);

约束说明

参数说明

参数名

输入/输出

说明

module

输入

原始TorchScript模型。

method_name

输入

TorchScript模型中需要进行编译优化的方法名称,需要为TorchScript模型中合法有效的方法名称,一般为“forward”。

info

输入

模型编译优化配置信息,详情请参考CompileSpec

返回值说明

返回编译优化后的engine,可用于保存om文件。