GFPGAN 转换为昇腾 OM 模型

昇思MindSpore 2025-11-25 16:39:56

将 GFPGAN 模型转换为昇腾（Ascend）平台的 OM（Offline Model）模型，通常需要经过以下几个步骤。以下流程基于华为昇腾 CANN（Compute Architecture for Neural Networks）工具链，适用于 Ascend 310/910 等 AI 处理器。

🧩 前提条件

原始模型格式：GFPGAN 通常是 PyTorch 模型（.pth或 .pt文件）。
目标平台：昇腾 AI 芯片（如 Atlas 300I 推理卡、Atlas 800 训练服务器等）。
所需软件环境：
- CANN Toolkit（含 ATC 工具）
- PyTorch（用于导出 ONNX）
- Python 环境
- 昇腾驱动 & 固件（若在真实硬件上运行）

🔁 转换流程概览

PyTorch (.pth) 
    → 导出为 ONNX 
    → 使用 ATC 工具转换为 OM 模型

步骤 1：准备 GFPGAN 的 PyTorch 模型

确保你有训练好的或官方提供的 GFPGAN 模型文件（例如 GFPGANv1.4.pth），并能加载到 PyTorch 中。

可参考官方仓库：
https://github.com/TencentARC/GFPGAN

步骤 2：将 PyTorch 模型导出为 ONNX

import torch
from gfpgan import GFPGANer  # 或直接加载 GFPGAN 模型类

# 示例：加载模型（根据实际代码调整）
model = ...  # 实例化你的 GFPGAN 模型
model.load_state_dict(torch.load('GFPGANv1.4.pth')['params'])
model.eval()

# 构造 dummy input（注意输入尺寸，通常为 [1, 3, H, W]，如 [1, 3, 512, 512]）
dummy_input = torch.randn(1, 3, 512, 512)

# 导出 ONNX
torch.onnx.export(
    model,
    dummy_input,
    "gfpgan.onnx",
    export_params=True,
    opset_version=11,  # 建议使用 opset 11 或 13（需与 ATC 兼容）
    do_constant_folding=True,
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={  # 可选：支持动态尺寸
        'input': {0: 'batch', 2: 'height', 3: 'width'},
        'output': {0: 'batch', 2: 'height', 3: 'width'}
    }
)

⚠️ 注意：GFPGAN 中可能包含自定义算子（如 StyleGAN2 的部分操作），这些在标准 ONNX 中可能不支持，需检查是否能成功导出。如有问题，可能需要修改模型或使用自定义 ONNX 算子。

步骤 3：使用 ATC 工具将 ONNX 转为 OM

在安装了 CANN 的环境中执行：

atc \
  --model=gfpgan.onnx \
  --framework=5 \                 # 5 表示 ONNX
  --output=gfpgan_om \
  --input_format=NCHW \
  --input_shape="input:1,3,512,512" \
  --log_level=error \
  --soc_version=Ascend310         # 根据你的芯片选择，如 Ascend910

✅ --soc_version必须与目标设备匹配。常见值：

Atlas 300I 推理卡：Ascend310
Atlas 900：Ascend910

如果使用了动态 shape，需额外指定 --dynamic_batch_size或 --dynamic_image_size（具体参考 ATC 文档）。

步骤 4：验证 OM 模型（可选）

可使用昇腾推理 API（如 ACL）编写 C++/Python 推理程序加载 gfpgan_om.om进行测试。

⚠️ 常见问题

ONNX 导出失败：GFPGAN 内部使用了非标准操作（如 FusedLeakyReLU、upfirdn2d）。需替换成 ONNX 支持的等效实现，或使用自定义 ONNX 算子。
ATC 报错“不支持的算子”：某些 ONNX 算子在昇腾中未实现。可通过 CANN 自定义算子开发（较复杂），或修改模型结构。
精度下降：OM 默认使用 FP16 推理。可加 --precision_mode=allow_fp32_to_fp16或强制 --precision_mode=force_fp32（性能下降）。

...全文