快速上手

示例说明

TorchAir提供的昇腾编译后端，能够作为参数传入PyTorch的torch.compile接口中。TorchAir的使能示例如下：

# 导入torchair框架
import torch
import torch_npu
import torchair as tng

# 定义模型Model
class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
    def forward(self, x, y):
        return torch.add(x, y)

# 实例化模型model
model = Model()

# 从torchair框架获取npu提供的 默认backend
npu_backend = tng.get_npu_backend()

# 使用torchair的backend去调用compile接口编译模型
model = torch.compile(model, backend=npu_backend, dynamic=False)

# 使用编译后的model去执行
x = torch.randn(2, 2)
y = torch.randn(2, 2)
model(x, y)

接口说明