NPU上如何使能pytorch图模式

1 Pytorch的compile技术PyTorch 的 torch.compile 是一个强大的功能，用于优化 PyTorch 模型的性能。它通过将 PyTorch 的动态图转换为静态图，并利用 Just-In-Time（JIT）编译技术，显著提高模型的推理速度和训练效率。1.1 PyTorch torch.compile 的作用torch.compile 是 PyTorch 的一个实验性功能，旨在通过编译优化提升模型的性能。它利用了 PyTorch 的 torch.jit 模块，将动态图转换为静态图，并通过后端编译器（如 LLVM 或 OpenVINO）进一步优化代码。1.2. 如何使用 torch.compile以下是一个简单的示例，展示如何使用 torch.compile 来优化模型：import torch
import torch.nn as nn

定义一个简单的模型

class SimpleModel(nn.Module):

def __init__(self):
    super(SimpleModel, self).__init__()
    self.fc = nn.Linear(10, 1)

def forward(self, x):
    return self.fc(x)

实例化模型

model = SimpleModel()

编译模型

compiled_model = torch.compile(model)

输入数据

input_data = torch.randn(1, 10)

使用编译后的模型进行推理

output = compiled_model(input_data)
print(output)2 TorchAir2.1 Torch Air是什么？TorchAir（Torch Ascend Intermediate Representation）是昇腾为Ascend Extension for PyTorch（torch_npu）提供的图模式能力扩展库，支持PyTorch网络在昇腾设备上进行图模式的训练和推理。TorchAir提供了昇腾设备的图模式编译后端，对接PyTorch的Dynamo特性，将PyTorch的FX（Functionalization）计算图转换为Ascend IR计算图，并通过GE（Graph Engine，图引擎）进行图编译、图执行、图优化等操作，并下发到昇腾硬件执行。TorchAir继承了Dynamo的大部分特性，如动态shape图功能，在PyTorch基础上又新增了图相关能力，如离线场景导图、模型编译时间优化、集合通信算子入图等功能，详细介绍参见功能介绍。

添加图片注释，不超过 140 字（可选）具体可以参考：简介-PyTorch 图模式使用（TorchAir）-套件与三方库-Ascend Extension for PyTorch6.0.0开发文档-昇腾社区2.2 Torch Air使用TorchAir图模式相关的功能配置示例如下，此处代码仅供参考，请根据实际情况开启对应功能项。# 必须先导torch_npu再导torchair
import torch
import torch_npu
import torchair

(可选)若涉及集合通信算子入图，可调用patch方法

from torchair import patch_for_hcom
patch_for_hcom()

定义模型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

config = torchair.CompilerConfig()
npu_backend = torchair.get_npu_backend(compiler_config=config)

使用TorchAir的backend去调用compile接口编译模型

opt_model = torch.compile(model, backend=npu_backend)

使用编译后的model去执行

x = torch.randn(2, 2)
y = torch.randn(2, 2)
opt_model(x, y)

具体可参考：快速上手-PyTorch 图模式使用（TorchAir）-套件与三方库-Ascend Extension for PyTorch6.0.0开发文档-昇腾社区