将原始TorchScript模型或ExportedProgram进行编译优化生成可在昇腾处理上加速推理的离线OM模型。
def export_engine(module: Any, method_name : str, ir = "default", inputs=None, **kwargs)
参数名称 |
参数类型 |
参数说明 |
是否必选 |
---|---|---|---|
module |
torch.jit.ScriptModule 或 torch.nn.Module 或 torch.export.ExportedProgram |
编译优化前的PyTorch模型。 默认值:无 |
是 |
inputs |
List[torch.Tensor] 或 List[mindietorch.Input]或 List[[mindietorch.Input]] |
模型输入。 默认值:None |
否 |
precision_policy |
Enum |
设置模型的推理精度策略,支持FP16精度、FP32精度以及混合精度PREF_FP32、PREF_FP16。 默认值:PREF_FP32 |
否 |
truncate_long_and_double |
Bool |
是否允许long和double类型转换。 默认值:True |
否 |
require_full_compilation |
Bool |
是否强制要求整图编译,若模型中存在不支持的算子,开启此项时会在编译模型时抛出异常并提示用户无法编译整图。 默认值:False |
否 |
allow_tensor_replace_int |
Bool |
是否允许采用Tensor代替Int。 默认值:False |
否 |
min_block_size |
Int |
切分子图的最少节点数量。 默认值:1 |
否 |
torch_executed_ops |
List[String] |
强制fallback执行的算子。 以add算子为例:
默认值:无 |
否 |
soc_version |
String |
芯片型号。 默认值:Ascend310P3 |
否 |
optimization_level |
Int |
模型优化等级,取值:0表示不优化,1表示图优化,2表示算子优化。 默认值:0 |
否 |
method_name |
String |
方法名。 默认值:forward |
否 |
ir |
Str |
模型编译方式,取值如下:
默认值:default 根据module的类型进行匹配,优先编译TorchScript模型,其次编译torch.export路线模型。 |
否 |
注:参数相关约束请参见export_engine。 |
返回编译优化后的engine,可用于保存om文件。