torch.compile路线分图特性
示例代码如下所示:
import torch import torch.nn as nn import torch_npu import mindietorch mindietorch.set_device(device_id) class Test(nn.Module): 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") backend_kwargs = { "torch_executed_ops": [torch.ops.aten.tanh.default], "min_block_size": 1, } opt_model = torch.compile(model, backend="mindie", options=backend_kwargs) device_id = 0 npu_input = input.to("npu") infer_ret = compiled_model(npu_input)[0].to("cpu")
父主题: MindIE Torch分图特性