将.ckpt模型转换为.om格式，运行转换脚本后报错Failed to allocate memory for tensor of shape [1, 64, 160, 160]

[chengxiaoli (Cheng_li)]

1 系统环境

• 硬件环境(Ascend/GPU/CPU): Ascend310
• MindSpore版本: mindspore=2.0.0
• 执行模式（PyNative/ Graph）: 不限
• Python版本: Python=3.9
• 操作系统平台: Linux

2 报错信息

2.1 问题描述

在使用昇腾310芯片部署YOLOv5模型时，尝试将训练好的.ckpt模型文件转换为.om格式，运行转换脚本后报错，错误发生在调用mindspore.ascend的export接口阶段，模型输入尺寸为640×640，批量大小为1。

2.2 脚本信息

import mindspore as ms  
from mindspore import export, load_checkpoint  
net = YOLOv5Backbone()  # 自定义YOLOv5网络  
load_checkpoint("yolov5.ckpt", net)  
input_tensor = ms.Tensor(np.ones([1, 3, 640, 640]), ms.float32)  
export(net, input_tensor, file_name="yolov5", file_format='MINDIR')

2.3 报错信息

TBE compiler error: [Error Code] RT_DEVICE_MEMORY_ALLOC_FAILED
[Detail] Failed to allocate memory for tensor of shape [1, 64, 160, 160]

3 根因分析

• 此处由用户填写

4 解决方案

• 此处由用户填写
• 包含文字方案和最终脚本代码 请将正确的脚本打包并上传附件