mindietorch.export_engine
函数功能
将原始TorchScript模型或ExportedProgram进行编译优化生成可在昇腾处理上加速推理的离线OM模型。
函数原型
def export_engine(module: Any, method_name : str = "forward", inputs=[], precision_policy = _enums.PrecisionPolicy.PREF_FP32, truncate_long_and_double = True, require_full_compilation = False, allow_tensor_replace_int = False, min_block_size = 3, torch_executed_ops = [], soc_version = "Ascend310P3", optimization_level = 0)
约束说明
- 只有TorchScript模型支持整图编译且不设置强制fallback算子才可导出离线OM模型。
- mindietorch.export_engine接口由于存在参数校验，在输入非法数据时，可能会抛出异常。故用户必须在try/except语句块内进行调用以及异常处理，防止在使用的过程中出现异常抛出导致程序退出的情况。
参数说明
参数名称 |
参数类型 |
参数说明 |
是否必选 |
|---|---|---|---|
module |
torch.jit.ScriptModule 或 torch.nn.Module 或 torch.export.ExportedProgram |
编译优化前的Pytorch模型。 默认值：无 |
是 |
inputs |
List[torch.Tensor] 或 List[torch_aie.Input]或 List[[torch_aie.Input]] |
模型输入。 默认值：无 |
是 |
precision_policy |
Enum |
设置模型的推理精度策略，支持混合精度PREF_FP32、FP16以及FP32精度。 默认值：PREF_FP32 |
否 |
truncate_long_and_double |
Bool |
是否允许long和double类型转换。 默认值：True |
否 |
require_full_compilation |
Bool |
是否整图编译，仅在编译TorchScript模型时生效。 默认值：False |
否 |
allow_tensor_replace_int |
Bool |
是否允许采用Tensor代替Int，仅在编译TorchScript模型时生效。 默认值：False |
否 |
min_block_size |
Int |
切分子图的最少节点数量，仅在编译TorchScript模型时生效。 默认值：3 |
否 |
torch_executed_ops |
List[String] |
强制在torch上执行的算子，如：["aten::add"]，仅在编译TorchScript模型时生效。 默认值：无 |
否 |
soc_version |
String |
芯片型号。 默认值：Ascend310P3 |
否 |
optimization_level |
Int |
模型优化等级，取值：0表示不优化，1表示图优化，2表示算子优化。 默认值：0 |
否 |
method_name |
String |
方法名。 默认值：forward |
否 |
注：参数相关约束请参见export_engine。 |
|||
返回值说明
返回编译优化后的engine，可用于保存om文件。