CANN/GE NPU模型装饰器API文档

npu_model【免费下载链接】geGEGraph Engine是面向昇腾的图编译器和执行器提供了计算图优化、多流并行、内存复用和模型下沉等技术手段加速模型执行效率减少模型内存占用。 GE 提供对 PyTorch、TensorFlow 前端的友好接入能力并同时支持 onnx、pb 等主流模型格式的解析与编译。项目地址: https://gitcode.com/cann/ge产品支持情况产品是否支持Atlas A3 训练系列产品/Atlas A3 推理系列产品√Atlas A2 训练系列产品/Atlas A2 推理系列产品√函数功能如果UDF部署在host侧执行时数据需要从device拷贝到本地进行运算。对于PyTorch场景如果计算全在device侧输入输出也是在device侧执行时数据需要从device拷贝到host执行后PyTorch再将数据搬到device侧影响执行性能使用npu_model可以优化为不搬移数据即直接下沉到device执行的方式触发执行。函数原型装饰器npu_model参数说明参数名称数据类型取值说明optimize_levelint- 1PyTorch场景下通过UDF nn引擎完成输入输出数据下沉到device执行默认值为1。- 2把PyTorch模型编译成图直接作为nn模型导出优化为npu模型加载执行需要配合input_descs使用。 说明 该配置项在修饰类的时候起作用修饰函数不能配置。input_descs[TensorDesc]当optimize_level2时用于表达torch导出成图的输入tensor描述示例如下input_descs[TensorDesc(dtype df.DT_INT64, shape [2,1,4]),TensorDesc(dtype df.DT_FLOAT, shape [2,1,4])],当shape中某一维度为负值表示输入是动态的通过npu_model最终会导出成动态图。num_returnsint装饰器装饰函数时用于表示函数的输出个数不设置该参数时默认函数返回一个返回值。该参数与使用type annotations方式标识函数返回个数与类型的方式选择其一即可。resourcesdict用于标识当前func需要的资源信息支持memory、num_cpus和num_npus。memory单位为M; num_npus表示需要使用npu资源数量为预留参数当前仅支持1。例如{memory: 100, num_cpus: 1, num_npus: 1}env_hook_funcfunction此钩子函数用于给用户自行扩展在Python UDF初始化之前必要的Python环境准备或import操作。visible_device_enablebool开启后UDF进程会根据用户配置num_npus资源自动设置ASCEND_RT_VISIBLE_DEVICES调用get_running_device_id接口获取对应的逻辑ID当前num_npus仅支持1因此该场景下get_running_device_id结果为0。返回值正常场景下返回被装饰的函数。异常情况下会抛出DfException异常。可以通过捕捉异常获取DfException中的error_code与message查看具体的错误码及错误信息。详细信息请参考DataFlow错误码。调用示例df.npu_model(optimize_level1) class FakeModel1(nn.Module): def __init__(self): super().__init__() # 模拟模型推理 df.method() def forward(self, input_image): return F.interpolate(input_image, size(256, 256), modebilinear) df.npu_model(optimize_level1, input_descs[df.TensorDesc(dtypedf.DT_FLOAT, shape[1, 3, 768, 768])]) class FakeModel2(nn.Module): def __init__(self): super().__init__() self.mean 0.5 self.std 0.5 # 模拟模型推理 df.method() def forward(self, input_image): return (input_image - self.mean) / self.std df.npu_model() def preprocess(input_image): # 模拟图片裁切 transform transforms.Compose([transforms.CenterCrop(512)]) return transform(input_image) df.npu_model() def postprocess(input_image): mean 0.5 std 0.5 img input_image * std mean return F.interpolate(img, size(512, 512), modebilinear)约束说明需安装对应Python版本的torch_npu包。输入输出必须为npu tensor。一组输入对应一组输出不支持流式输入输出。【免费下载链接】geGEGraph Engine是面向昇腾的图编译器和执行器提供了计算图优化、多流并行、内存复用和模型下沉等技术手段加速模型执行效率减少模型内存占用。 GE 提供对 PyTorch、TensorFlow 前端的友好接入能力并同时支持 onnx、pb 等主流模型格式的解析与编译。项目地址: https://gitcode.com/cann/ge创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考