昇腾 CANN 实战：超分辨率模型从 PyTorch 到 ONNX 再到 OM 的完整转换指南

最近在边缘计算和端侧AI部署的浪潮里，不少开发者开始关注如何将训练好的模型高效地运行在专用硬件上。昇腾（Ascend）处理器及其配套的CANN（Compute Architecture for Neural Networks）软件栈，为这个需求提供了一个颇具吸引力的选择。但当你兴致勃勃地拿到一块昇腾开发板，准备把实验室里用PyTorch精心调校的超分辨率模型搬上去时，可能会发现从熟悉的框架到陌生的硬件之间，横亘着一条名为“模型转换”的沟壑。这篇文章，我就想和你聊聊，怎么一步步填平这条沟，把PyTorch模型顺畅地“翻译”成昇腾设备能理解的OM格式，并跑起来。整个过程，远不止是敲几行命令那么简单，里面有不少细节和“坑”，我会结合自己趟过的路，把关键节点和实用技巧都摊开来讲。

1. 转换前的核心认知与准备工作

在动手之前，我们需要先理清几个基本概念，这能帮你理解后续每一步操作背后的逻辑，而不是机械地复制命令。

昇腾CANN的推理生态，其核心思想是将不同深度学习框架（如PyTorch、TensorFlow）训练的模型，通过一个中间表示（Intermediate Representation, IR）进行统一，最终编译成在昇腾硬件上高效执行的离线模型（Offline Model, OM）。这个中间表示，最常见的就是ONNX（Open Neural Network Exchange）。所以，我们的转换路径很清晰：PyTorch -> ONNX -> OM。

注意：虽然CANN也支持直接转换MindSpore的MindIR格式，但对于PyTorch和TensorFlow用户，ONNX是目前最通用、生态最成熟的桥梁。

为什么需要OM模型？OM模型是经过ATC（Ascend Tensor Compiler）工具编译优化后的二进制文件。它已经针对特定的昇腾芯片（如Ascend 310B）进行了算子融合、内存优化、精度校准等深度优化，因此推理时无需框架开销，能获得极致的性能和能效。这类似于移动端的NCNN、MNN，或者NVIDIA的TensorRT。

准备工作清单：

• 硬件环境：确保你有一台可以访问昇腾设备的开发环境。这可以是Atlas 200I DK A2开发者套件，或者云端带有昇腾芯片的ECS实例。本文的操作主要基于开发板环境。
• 软件环境：在开发环境（通常是x86的Ubuntu服务器或PC）上安装好CANN Toolkit。这是包含ATC转换工具、AscendCL（Ascend Computing Language）运行时库等一整套开发套件。务必确认安装版本与目标昇腾设备的驱动版本匹配。
• 模型与数据：准备好你的PyTorch模型文件（.pth）和一小部分用于验证转换正确性的测试数据（如图片）。对于超分辨率任务，像Set5、Set14这样的标准测试集就很好用。

这里有一个常见的环境配置对照表，帮你快速检查：

组件	开发环境（用于转换）	运行环境（昇腾设备）
操作系统	Ubuntu 18.04/20.04 x86_64	Ubuntu 20.04 aarch64 (开发板)
核心软件	CANN Toolkit (含ATC)	CANN Runtime (含AscendCL)
Python环境	PyTorch, ONNX, onnxruntime等	通常只需Python运行AscendCL脚本
关键工具	atc 模型转换命令行工具	acl 推理库

2. 从PyTorch到ONNX：导出模型的陷阱与技巧

这一步看似简单，一句torch.onnx.export就能搞定，但导出的ONNX模型是否“健康”，直接决定了后续转换的成败。

首先，确保你的模型处于推理模式。这不仅仅是调用model.eval()，还要注意模型中是否存在训练阶段特有的行为，比如Dropout层、BatchNorm层的不同运行模式。一个稳妥的做法是在导出前，遍历所有子模块，将其设置为eval()。

import torch
from your_model_arch import YourSRModel

# 加载模型架构和权重
model = YourSRModel(scale_factor=2)
state_dict = torch.load('your_model.pth', map_location='cpu')
model.load_state_dict(state_dict)

# 切换到推理模式
model.eval()
# 显式设置所有子模块，避免遗漏
for m in model.modules():
    if hasattr(m, 'training'):
        m.training = False

其次，构造正确的输入样例（Dummy Input）。输入张量的形状(batch_size, channels, height, width)和数据类型必须与模型训练时一致。对于超分辨率模型，高度和宽度通常可以动态，但通道数（如RGB图像的3）和批次维度（通常是1）是固定的。数据类型也要留意，很多模型权重是float32，但输入可能是uint8的图片转换而来。

# 创建一个符合模型预期的虚拟输入
# 假设模型接受任意尺寸的输入，但通道数为3，批次为1
dummy_input = torch.randn(1, 3, 256, 256, dtype=torch.float32)

接下来是导出ONNX的核心操作。这里的关键是dynamic_axes参数，它定义了哪些维度是动态的。对于超分辨率模型，我们通常希望输入图片的height和width是动态的，以适应不同尺寸的输入。

# 导出ONNX模型
onnx_model_path = "super_resolution.onnx"
torch.onnx.export(
    model,                      # PyTorch模型
    dummy_input,                # 虚拟输入
    onnx_model_path,            # 输出路径
    input_names=["input"],      # 输入节点名称
    output_names=["output"],    # 输出节点名称
    dynamic_axes={              # 指定动态维度
        "input": {2: "height", 3: "width"},
        "output": {2: "height", 3: "width"}
    },
    opset_version=11,           # 建议使用较新的opset，如11或12，兼容性更好
    do_constant_folding=True    # 常量折叠优化
)

导出后，强烈建议使用Netron（一个开源模型可视化工具）打开生成的.onnx文件。检查以下几点：

1. 模型输入/输出的名称和维度是否符合预期。
2. 模型结构是否完整，有无缺失的节点。
3. 是否存在ONNX不支持的PyTorch算子。如果存在，你需要寻找替代实现或自定义算子。

提示：如果遇到不支持的算子，可以尝试更新PyTorch和ONNX的版本。对于一些复杂操作（如某些插值方式、自定义激活函数），可能需要手动实现其ONNX导出逻辑，或者查阅CANN文档看ATC是否提供了该算子的直接支持。

3. 使用ATC工具将ONNX转换为OM模型

拿到“健康”的ONNX模型后，就可以请出CANN的“编译器”——ATC工具了。这一步是将与框架、硬件无关的中间表示，编译成高度优化的昇腾硬件指令。

转换命令的核心是atc，它的参数繁多，但掌握几个关键的就够了。我们从一个支持动态分辨率的超分辨率模型转换命令开始：

atc --model=./super_resolution.onnx \
    --framework=5 \
    --input_shape="input:1,3,-1,-1" \
    --dynamic_image_size="256,256;512,512;1024,768" \
    --output=./super_resolution_om \
    --soc_version=Ascend310B1 \
    --log=info

我们来拆解这些参数：

• --model: 指定输入的ONNX模型路径。
• --framework: 指定输入框架，5代表ONNX。
• --input_shape: 定义输入张量的形状。-1表示该维度是动态的。这里"input:1,3,-1,-1"意味着批次为1，通道为3，高度和宽度动态。
• --dynamic_image_size: 这是为动态的height和width维度预设的具体尺寸组合。ATC会根据这些组合进行编译优化。至少需要提供两组，例如"256,256;512,512"。提供的组合越多，模型对不同尺寸的适应性越好，但转换时间会变长，生成的OM文件也可能略大。
• --output: 指定输出的OM模型文件名（无需后缀）。
• --soc_version: 至关重要，必须与你的昇腾设备型号严格匹配，如Ascend310B1、Ascend310P1等。填错会导致模型无法在目标设备上运行。
• --log: 设置日志级别，调试时设为debug或info可以看到更多细节。

关于动态形状的深入讨论： --dynamic_image_size参数预设的是(height, width)的组合。如果你需要批次（batch size）也是动态的，需要使用--dynamic_batch_size参数，例如--dynamic_batch_size="1,2,4,8"。甚至可以使用更灵活的--dynamic_dims参数来定义多个维度的动态范围。具体选择哪种，取决于你的应用场景。对于单张图片推理的超分辨率，动态height/width是最常见的需求。

转换过程中，终端会输出大量编译信息。如果成功，最后会看到ATC run success的提示。如果失败，日志信息是排查问题的第一手资料。常见错误包括：不支持的算子、输入形状定义冲突、内存不足等。

4. 使用AscendCL进行推理：从示例到实战

OM模型转换成功后，就来到了最后一步：在昇腾设备上加载并执行推理。这里我们使用昇腾提供的基础推理接口——AscendCL（Ascend Computing Language）。它是一套C语言API，但官方也提供了Python绑定，用起来相对方便。

与其从零开始编写AscendCL代码，更高效的方法是复用官方示例。在CANN的安装目录或开源样本库中，通常会有resnet50_imagenet_classification这样的分类样例。我们的策略是，理解其流程，然后将其“改造”成适合超分辨率模型的程序。

一个典型的AscendCL推理流程包括以下步骤：

1. 初始化：初始化AscendCL运行管理资源。
2. 加载模型：从OM文件加载模型到设备内存。
3. 准备输入：将输入数据（如图片）处理成模型需要的格式（NCHW， FP32等），并拷贝到设备内存。
4. 执行推理：调用模型进行前向计算。
5. 获取输出：从设备内存取回推理结果数据。
6. 后处理与释放：对结果进行后处理（如超分辨率图像的保存），并释放所有申请的资源。

下面，我以改造图片预处理和结果后处理为例，展示关键代码的修改思路。假设我们有一个处理静态输入(1,3,200,200)的OM模型。

首先，图片预处理函数需要将任意输入图片，处理成模型需要的固定尺寸和格式。对于尺寸不匹配的图片，常见的做法是填充（Padding）或缩放。这里采用填充以保持原始内容比例。

import numpy as np
from PIL import Image

def preprocess_for_sr(image_path, target_size=(200, 200)):
    """
    将图片预处理为模型输入张量。
    目标：转换为NCHW，FP32，数值范围[0,1]，并填充到固定尺寸。
    """
    # 1. 使用PIL打开图片，并转换为RGB
    img = Image.open(image_path).convert('RGB')
    original_size = img.size  # (W, H)

    # 2. 计算填充，使图片居中放置在target_size的画布上
    new_img = Image.new('RGB', target_size, (0, 0, 0))  # 创建黑色背景画布
    # 计算粘贴位置（左上角坐标）
    paste_x = (target_size[0] - original_size[0]) // 2
    paste_y = (target_size[1] - original_size[1]) // 2
    new_img.paste(img, (paste_x, paste_y))

    # 3. 转换为numpy数组，并调整数值范围和维度顺序
    img_np = np.array(new_img, dtype=np.float32) / 255.0  # HWC, [0,1]
    img_np = img_np.transpose(2, 0, 1)  # HWC -> CHW
    img_np = np.expand_dims(img_np, axis=0)  # CHW -> NCHW

    # 4. 记录填充信息，用于后处理时裁剪
    padding_info = (paste_x, paste_y, original_size[0], original_size[1])
    
    # AscendCL需要的是数据指针和大小，这里我们返回flatten后的数据和元信息
    input_data = img_np.flatten().astype(np.float32)
    return input_data, padding_info, target_size

然后，在推理结果后处理函数中，我们需要将模型输出的NCHW格式张量，转换回图片，并根据之前记录的填充信息进行裁剪，恢复原始比例。

def postprocess_from_sr(output_data, padding_info, target_size):
    """
    将模型输出张量处理回图片。
    output_data: 模型输出的原始一维数组（已从设备内存拷贝回）。
    padding_info: (paste_x, paste_y, original_w, original_h)
    target_size: (target_w, target_h)
    """
    # 1. 将一维输出数据重塑为NCHW形状
    # 假设模型输出是[1, 3, H, W]，且H,W与输入相同（对于超分，通常是scale倍）
    scale = 2  # 假设是2倍超分
    out_h, out_w = target_size[1]*scale, target_size[0]*scale
    output_np = output_data.reshape(1, 3, out_h, out_w)

    # 2. 取第一个批次，转换维度顺序，并调整数值范围
    sr_img_chw = output_np[0]  # CHW
    sr_img_hwc = sr_img_chw.transpose(1, 2, 0)  # CHW -> HWC
    sr_img_hwc = np.clip(sr_img_hwc, 0, 1) * 255
    sr_img_hwc = sr_img_hwc.astype(np.uint8)

    # 3. 根据填充信息裁剪，还原有效区域
    paste_x, paste_y, orig_w, orig_h = padding_info
    # 注意：超分后，填充区域和原始区域都放大了scale倍
    crop_x1 = paste_x * scale
    crop_y1 = paste_y * scale
    crop_x2 = crop_x1 + (orig_w * scale)
    crop_y2 = crop_y1 + (orig_h * scale)
    
    valid_sr_img = sr_img_hwc[crop_y1:crop_y2, crop_x1:crop_x2, :]

    # 4. 转换为PIL图像并返回或保存
    result_img = Image.fromarray(valid_sr_img, 'RGB')
    return result_img

最后，在主函数中，你需要将官方示例中关于ResNet的分类逻辑（如加载标签文件、计算top-k精度）替换为调用上述预处理、推理、后处理的流程。核心的AscendCL模型加载、内存分配、推理执行等代码框架通常无需大改，只需确保输入输出数据的大小与你的模型匹配。

调试时，可以先在开发环境用ONNX Runtime加载同一个ONNX模型，用相同的预处理流程跑一个CPU/GPU的推理，将结果与你昇腾设备上的推理结果进行对比，这是验证整个转换和预处理流水线是否正确的最有效方法。