NPU上如何使能pytorch图模式

news2025/4/21 10:02:44

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 TorchAir

2.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基础上又新增了图相关能力，如离线场景导图、模型编译时间优化、集合通信算子入图等功能，详细介绍参见功能介绍。
在这里插入图片描述
具体可以参考：
简介-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)