NPU 上如何使能 pytorch 图模式

作者：zjun

2025-03-24
浙江
本文字数：1294 字
阅读完需：约 4 分钟

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 torchimport 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再导torchairimport torchimport torch_npuimport torchair
# (可选)若涉及集合通信算子入图，可调用patch方法from torchair import patch_for_hcompatch_for_hcom()
# 定义模型Modelclass Model(torch.nn.Module):    def __init__(self):        super().__init__()    def forward(self, x, y):        return torch.add(x, y)
# 实例化模型modelmodel = Model()
# 从TorchAir框架获取NPU提供的默认backendconfig = 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开发文档-昇腾社区

发布于: 刚刚阅读数: 5

zjun

关注

还未添加个人签名 2020-03-06 加入

还未添加个人简介

发布

暂无评论

创作场景