torch.compile路线分图特性

以下代码将手动把tanh算子fallback到torch_npu执行，必须与torch_npu的配套使用，详情请参见配套Torch_NPU使用安装torch_npu环境。

示例代码如下所示：
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
# 1.导入mindietorch框架
import torch
import torch.nn as nn
import torch_npu
import mindietorch

device_id = 0
mindietorch.set_device(device_id) # 设置使用device 0设备

# 2.模型导出
class Test(nn.Module): # 定义模型Test
    def forward(self, x):
        x = torch.ops.aten.relu.default(x)
        x= torch.ops.aten.tanh.default(x)
        out = torch.ops.aten.sigmoid.default(x)
        return out
shape = (2, 2)
input = torch.randn(shape)
model = Test().to("npu")

# 3.模型编译
backend_kwargs = {
    "torch_executed_ops": [torch.ops.aten.tanh.default],
    "min_block_size": 1,
}
opt_model = torch.compile(model, backend="mindie", options=backend_kwargs)

# 4.模型推理
npu_input = input.to("npu")
infer_ret = opt_model(npu_input)[0].to("cpu")
父主题： MindIE Torch分图特性