将GFPGAN模型转换为ONNX格式时出现AttributeError错误

错误原因 #

你直接用torch.load()加载的是模型的权重字典（dict类型），而不是实例化后的GFPGAN模型对象。torch.onnx.export()需要传入的是PyTorch模型实例，不是权重字典，所以会触发'dict' object has no attribute 'modules'错误。

正确转换流程 #

1. 克隆GFPGAN仓库并安装依赖 #

必须先获取GFPGAN的模型定义代码，因为权重文件本身不包含模型结构：

!git clone https://github.com/TencentARC/GFPGAN.git
%cd GFPGAN
!pip install -r requirements.txt
!pip install onnx torch torchvision

2. 下载预训练权重 #

import urllib.request
url = "https://github.com/TencentARC/GFPGAN/releases/download/v0.2.0/GFPGANCleanv1-NoCE-C2.pth"
urllib.request.urlretrieve(url, "GFPGANCleanv1-NoCE-C2.pth")

3. 实例化GFPGAN模型并加载权重 #

从仓库的封装类中获取实际模型实例，再加载权重：

import torch
from gfpgan import GFPGANer

# 实例化GFPGAN封装类，自动加载权重
gfpgan = GFPGANer(
    model_path='GFPGANCleanv1-NoCE-C2.pth',
    upscale=2,
    arch='clean',
    channel_multiplier=2,
    bg_upsampler=None
)

# 提取核心模型对象
model = gfpgan.gfpgan
model.eval()  # 切换到评估模式，必须执行，否则影响ONNX导出

4. 导出为ONNX格式 #

注意传入实际张量输入，而非输入形状，同时可指定动态轴适配Android端的尺寸调整：

import onnx

# 采用GFPGAN官方默认输入尺寸512x512，避免尺寸不兼容
input_shape = (1, 3, 512, 512)
example_input = torch.randn(input_shape)

output_path = "GFPGAN.onnx"
torch.onnx.export(
    model,
    example_input,
    output_path,
    opset_version=11,
    do_constant_folding=True,
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={'input': {0: 'batch_size', 2: 'height', 3: 'width'},
                  'output': {0: 'batch_size', 2: 'height', 3: 'width'}}
)

# 验证ONNX模型有效性
onnx_model = onnx.load(output_path)
onnx.checker.check_model(onnx_model)
print("ONNX模型导出成功并验证通过")

关键注意事项 #

• 必须先实例化模型结构，再加载权重文件，不能直接将权重字典传入导出函数。
• 导出前务必调用model.eval()，关闭训练相关的层（如BatchNorm、Dropout），保证导出结果稳定。
• 优先使用官方默认的输入尺寸，避免后续Android端推理出现形状不匹配问题。

内容的提问来源于stack exchange，提问作者Hammad Ali Shah