mindietorch.export_engine
函数功能
将原始TorchScript模型或ExportedProgram进行编译优化生成可在昇腾处理上加速推理的离线OM模型。
函数原型
def export_engine(module: Any, method_name: str="forward", ir="default", inputs=None, **kwargs)
约束说明
- 只有TorchScript路线和torch.export路线模型，在支持整图编译且不设置强制fallback算子的情况下，才可以导出离线OM模型。
- 由于PyTorch的dynamo存在自动图拆分机制，export_engine接口不支持在使用torch.compile的场景下使用。
- mindietorch.export_engine接口由于存在参数校验，在输入非法数据时，可能会抛出异常。故用户必须在try/except语句块内进行调用以及异常处理，防止在使用的过程中出现异常抛出导致程序退出的情况。
参数说明
参数名称 |
参数类型 |
参数说明 |
是否必选 |
|---|---|---|---|
module |
torch.jit.ScriptModule或torch.nn.Module或torch.export.ExportedProgram |
编译优化前的PyTorch模型。在ir设置成不同值时，module传入的类型限制不同，具体请参见下面ir字段的参数说明。 默认值：无 |
是 |
inputs |
List[torch.Tensor]或List[List[torch.Tensor]]或List[mindietorch.Input]或List[List[mindietorch.Input]] |
模型输入。
说明：
当ir参数设置为dynamo时，inputs不支持传入，表示动态分档的List[List[mindietorch.Input]]类型数据。 默认值：None |
否 |
precision_policy |
mindietorch.PrecisionPolicy |
设置模型的推理精度策略，支持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 |
切分子图的最少节点数量，取值范围[0,1024]。 默认值：1 |
否 |
default_buffer_size_vec |
List[Int] |
模型输出在编译阶段无法确定shape时的默认分配内存大小，一般用于动态模型，取值范围(0,36864]，单位为MB。 支持List的长度等于1或者输出的个数，若等于1则所有输出的默认内存大小均为该值，若等于输出个数则为每个输出单独设置默认内存大小。 默认值：[500, ] |
否 |
torch_executed_ops |
List[str] 或 List[torch._ops.OpOverload] |
强制fallback执行的算子。 以add算子为例：
默认值：无 |
否 |
soc_version |
str |
芯片型号。 默认值：Ascend310xxx |
否 |
optimization_level |
int |
模型优化等级，取值如下：
默认值：0 |
否 |
method_name |
str |
方法名，用于指定将传入module对象的某一方法所执行操作编译导出成OM模型。 默认值：“forward“ |
否 |
ir |
str |
模型编译方式，取值如下：
默认值：default 根据module的类型进行匹配，优先选择编译TorchScript模型，其次选择编译torch.export路线模型。 |
否 |
注：参数相关约束请参见export_engine。 |
|||
返回值说明
返回编译优化后以bytearray类型表示的engine，可用于保存om文件。