mindietorch.compile

module

torch.jit.ScriptModule 或 torch.nn.Module 或 torch.export.ExportedProgram

编译优化前的PyTorch模型。

默认值：无

是

inputs

List[torch.Tensor] 或 List[mindietorch.Input]或 List[[mindietorch.Input]]

模型输入。

默认值：None

否

precision_policy

PrecisionPolicy(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算子为例：

• TorchScript路线：以libtorch风格的字符串指定算子，即torch_executed_ops=["aten::add"]；
• torch.export或者torch.compile路线：以PyTorch中的OpOverload对象指定算子，即torch_executed_ops=[torch.ops.aten.add.default]。

默认值：[]

否

soc_version

String

芯片型号。

默认值：Ascend310P3

否

optimization_level

Int

模型优化等级，取值如下：

• 0：表示不优化；
• 1：表示图优化；
• 2：表示算子优化。

默认值：0

否

ir

String

模型编译方式，取值如下：

• torchscript或者ts：表示编译TorchScript路线模型；此时仅支持module传入torch.jit.ScriptModule；
• dynamo：表示编译torch.export路线模型，此时支持module传入torch.nn.Module或torch.export.ExportedProgram；
• torch_compile：表示编译torch.compile路线模型。

默认值：default

根据module的类型进行匹配，优先编译TorchScript路线模型，其次编译torch.export路线模型，最后才是编译torch.compile路线模型。

否

enable_dynamic_torch_compile

Bool

是否让torch.compile捕获动态的模型，仅在编译torch.compile路线模型时生效。

默认值：False

否

注：参数相关约束请参见compile。