当前位置：首页 > article >正文

resnet网络迁移到昇腾执行（OM上篇）

article 2025/4/1 11:26:46

总体介绍

pytorch迁移OM模型

原始代码详细介绍

模型加载和初始化

初始化统计变量

数据推理及归一化

统计每个样本的结果

基本概念

Softmax（归一化指数函数）

作用

代码示例

应用场景

argmax取最大值索引

作用

代码示例

两者配合使用

原始代码总结

迁移到昇腾

模型加载

数据推理及归一化

统计每个结果

迁移过程总结

OM模型不需要实例化

OM输入数据要求

后处理的运行位置

迁移后代码执行

need size

input size

总结

本系列文章记录模型到310芯片的迁移，在上一篇中已经介绍了模型的转换，后续介绍实际的迁移过程。

总体介绍

总体而言，总共有torch_npu和om模型两种方式。

1）torch_npu采用原始的pt模型文件，相关接口采用pytorch的兼容实现。

2）om模型方式，即先将pt模型转换到om。同时相关接口采用cann或者cann的高层封装例如ais_bench等实现。

文件类型

格式

使用场景

优点

缺点

.pt

PyTorch原生格式

动态图模式

（训练/推理）

灵活修改模型结构，支持动态计算图

需要运行时解析，性能可能略低

.om

昇腾离线模型格式

静态图模式

（仅推理）

高度优化，NPU执行效率最高

模型结构固定，无法修改

通过上述描述，我们可以发现，在模型调试时采用pt方式。在业务部署时采用om方式。

如果代码需要继续训练或微调模型，必须使用 .pt 文件（.om 仅支持推理）

pytorch迁移OM模型

原始代码详细介绍

模型加载和初始化

cnn_model = Resnet50_2d()
cnn_model.load_state_dict(torch.load('ok.pt'))
cnn_model.cuda()

创建一个ResNet50_2d模型实例
从'ok.pt'文件加载预训练权重
将模型转移到CUDA设备(GPU)上

初始化统计变量

target_correct = list(0. for i in range(num_classes))
target_total = list(0. for i in range(num_classes))
conf = np.zeros([num_classes, num_classes])
confnorm = np.zeros([num_classes, num_classes])

target_correct: 记录每个类别预测正确的样本数
target_total: 记录每个类别的总样本数
conf: 初始化混淆矩阵(实际类别×预测类别)
confnorm: 用于存储归一化的混淆矩阵

数据推理及归一化

for data, _, target in test_loader:
    data, target = Variable(data).float().cuda(), Variable(target).long().cuda()
    [feature, output] = cnn_model(data)
    output = soft_max(output)

从测试数据加载器(test_loader)中获取数据
将数据和标签转换为PyTorch变量并转移到GPU
前向传播获取模型输出(这里模型返回特征和分类结果)
对输出应用Softmax得到概率分布

统计每个样本的结果

for i in range(output.shape[0]):
    test_label = torch.argmax(output[i, :])  # 获取预测类别
    label = target.data[i]  # 获取真实类别
    
    # 统计正确预测数
    target_correct[label] += test_label.eq(label)
    
    # 统计总样本数
    target_total[label] += 1
    
    # 更新混淆矩阵
    conf[label, test_label] += 1

对批次中的每个样本:
- 获取预测类别(概率最大的类别)
- 获取真实类别标签
- 如果预测正确，增加对应类别的正确计数
- 增加对应类别的总样本计数
- 在混淆矩阵中对应位置+1

基本概念

Softmax（归一化指数函数）

作用

将神经网络的原始输出（logits）转换为 概率分布，满足：

所有类别的概率值 ∈ [0, 1]
所有类别的概率之和 = 1

代码示例

output = torch.tensor([3.0, 1.0, 0.2])  # 原始输出（logits）
prob = torch.softmax(output, dim=0)      # 转换为概率
print(prob)  # 输出例如：tensor([0.8360, 0.1131, 0.0509])

应用场景

多分类任务中，模型需要输出每个类别的置信度（概率）。
在交叉熵损失函数（Cross-Entropy Loss）中，Softmax 常与对数运算结合使用

`argmax取最大值索引`

`作用`

用于获取张量中最大值所在的索引位置。在分类任务中，它通常用于从模型输出的概率分布中确定预测的类别。

test_label = torch.argmax(output[i, :])

这里的 output[i, :] 表示第 i 个样本对所有类别的预测分数（通常是 softmax 后的概率分布），argmax 找出这些分数中最大值对应的索引，即模型预测的类别标签

代码示例

tensor([0.1, 0.7, 0.2])

torch.argmax(output) 将返回 1（因为第二个元素 0.7 是最大值）

两者配合使用

原始数据经过归一化处理后，再获取最大值索引即得到分类标签

原始输出 (Logits) → [3.0, 1.0, 0.2]
  ↓ Softmax  
概率分布 → [0.8360, 0.1131, 0.0509]  
  ↓ Argmax  
预测类别 → 0

原始代码总结

原始代码总共分为如下3个部分。

迁移到昇腾

模型加载

本例处采用ais封装的高层接口，不采用直接的cann接口。

device_id = 0 
model_path = "resnet50_2d.om"  # 转换后的OM模型文件
# 创建推理会话
session = InferSession(device_id, model_path)

数据推理及归一化

for data, _, target in test_loader:
    # 将数据转换为numpy数组并符合OM模型输入要求
    input_data = data.numpy().astype(np.float32)
    
    # 执行推理
    outputs = session.infer([input_data])  # 返回list，每个元素对应一个输出
   
    print(outputs)
    
    # 获取输出并计算softmax
    output = torch.from_numpy(outputs[1])
    output = torch.softmax(output, dim=1)
    #output = torch.softmax(output)

统计每个结果


    # 统计结果
    for i in range(output.shape[0]):
        test_label = torch.argmax(output[i, :]).item()
        label = target[i].item()
        target_correct[label] += (test_label == label)
       
        target_total[label] += 1
        conf[label, test_label] += 1

迁移过程总结

OM模型不需要实例化

OM模型是华为昇腾处理器专用的离线模型格式，它已经包含了：

完整的模型架构：已经固化了网络结构
优化后的计算图：针对昇腾芯片进行了深度优化
预编译的算子：所有算子都已编译为昇腾芯片可执行的指令
模型已完全序列化
运行时不依赖Python框架

OM输入数据要求

必须与模型编译时指定的输入数据类型一致
常见类型：float32、float16、int8、int32等

input_data = data.numpy().astype(np.float32)

此处为float32

后处理的运行位置

即在GPU(NPU) 还是CPU上。

output = soft_max(output)
        print(output.device)

通过上述代码可以看出后处理运行的位置。

原始代码运行在GPU上；

而迁移后的代码还在CPU上。

迁移后代码执行

报如下错误

[ERROR] check i:0 name:input in size:8000000 needsize:262144 not match
[ERROR] check input vector failed ret:-1

还是tensor size不匹配导致。

need size

之前模型转换时采用的tensor size： :2,2,2,8192，对应的数据大小为：

计算所有维度的乘积： 2×2×2×8192=65536 个元素2×2×2×8192=65536 个元素

如果每个元素是 float32（4字节），则总内存占用为：65536×4=262144 字节 (≈256 KB)65536×4=262144 字节 (≈256 KB)

也就是错误提示中need size的值。

input size

那么input的值又从何来呢？通过分析读取数据集的接口

单个样本的 Shape: (2，100，100)
。2通道：实部(1)和虚部(Q)。
。100x100:信号片段重塑为2D矩阵。
。 Batch Shape:(batch_size, 2, 100, 100)
·用途：适用于处理复信号(IQ数据）的2D表示(如时频分析、信号图像化等）。
如果需要调整输入尺寸，修改data_dim的值即可(但需同步调整模型结构以适应新的H和w）。

再根据数据集文件的个数，我们得到输入的size为：--input_shape="input:160,2,100,100"

修改后的模型转换命令：

atc  --model=ok100100.onnx     --framework=5     --output=resnet50_2d     --soc_version=Ascend310B1    --input_format=NCHW     --input_shape="input:160,2,100,100"     --log=info

采用此模型后，在NPU上可以正常推理了。