欢迎来到尧图网

客户服务 关于我们

您的位置:首页 > 文旅 > 文化 > NPU上如何使能pytorch图模式

NPU上如何使能pytorch图模式

2025/3/25 23:04:42 来源:https://blog.csdn.net/xyz3120/article/details/146477448  浏览:    关键词: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 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)

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

版权声明:

本网仅为发布的内容提供存储空间,不对发表、转载的内容提供任何形式的保证。凡本网注明“来源:XXX网络”的作品,均转载自其它媒体,著作权归作者所有,商业转载请联系作者获得授权,非商业转载请注明出处。

我们尊重并感谢每一位作者,均已注明文章来源和作者。如因作品内容、版权或其它问题,请及时与我们联系,联系邮箱:809451989@qq.com,投稿邮箱:809451989@qq.com

热搜词