GFPGANv1模型.pth转ONNX失败：卷积核形状未知不支持导出

转换GFPGANv1的.pth模型到ONNX时，出现报错：

Unsupported: ONNX export of convolution for kernel of unknown shape

根源是TorchScript图中的Reshape节点导致卷积核形状无法被解析为常量，更换ONNX Opset版本或PyTorch版本（仅2.1.0能显示完整报错）后问题仍存在。

1. 手动固定Reshape节点的静态形状 #

ONNX导出卷积层要求卷积核形状为确定常量，若模型中存在依赖输入动态计算形状的Reshape/View操作，需将其替换为固定数值：

• 定位模型代码中动态Reshape的位置（比如GFPGAN依赖的StyleGAN2模块里的上采样/下采样层）
• 将动态计算的维度（如x.shape[1]）替换为模型固定的通道数/尺寸数值

示例修改：

# 原动态Reshape代码
x = torch.reshape(x, (-1, x.shape[1] // 2, x.shape[2] * 2, x.shape[3] * 2))

# 修改为静态形状（假设通道数固定为256，因此//2后为128）
x = torch.reshape(x, (-1, 128, x.shape[2] * 2, x.shape[3] * 2))

# 若输入尺寸固定，可进一步固定高宽维度
x = torch.reshape(x, (-1, 128, 1024, 1024))

2. 导出ONNX时固定输入形状并限制动态维度 #

导出时仅保留batch维度为动态，其余维度固定为模型预期的输入尺寸（比如512x512），避免TorchScript追踪动态形状：

import torch
from gfpgan import GFPGANer

# 加载预训练模型
model = GFPGANer(
    model_path='GFPGANv1.pth',
    upscale=2,
    arch='clean',
    channel_multiplier=2
).model
model.eval()

# 构造固定形状的dummy输入
dummy_input = torch.randn(1, 3, 512, 512)

# 导出ONNX
torch.onnx.export(
    model,
    dummy_input,
    "GFPGANv1.onnx",
    opset_version=17,
    input_names=["input"],
    output_names=["output"],
    # 仅允许batch维度动态变化
    dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}},
    verbose=False,
    do_constant_folding=True  # 启用常量折叠，帮助解析静态形状
)

3. 用Torch FX追踪并修改模型（PyTorch 1.10+支持） #

通过Torch FX捕获模型计算图，批量修改所有动态Reshape节点为静态形状：

import torch
import torch.fx
from gfpgan import GFPGANer

model = GFPGANer(model_path='GFPGANv1.pth', upscale=2, arch='clean', channel_multiplier=2).model
model.eval()

# 用FX追踪模型
traced_model = torch.fx.symbolic_trace(model)

# 遍历所有节点，修改Reshape/View的动态参数
for node in traced_model.graph.nodes:
    if node.op == 'call_function' and node.target in [torch.reshape, torch.Tensor.view]:
        new_args = list(node.args)
        # 示例：将shape中的动态通道维度替换为固定值128
        if isinstance(new_args[1], tuple):
            fixed_shape = []
            for dim in new_args[1]:
                if isinstance(dim, torch.fx.node.Node) and dim.target == 'shape' and dim.args[1] == 1:
                    fixed_shape.append(128)
                else:
                    fixed_shape.append(dim)
            new_args[1] = tuple(fixed_shape)
            node.args = tuple(new_args)

# 重新编译模型
traced_model.recompile()

# 导出ONNX
dummy_input = torch.randn(1, 3, 512, 512)
torch.onnx.export(
    traced_model,
    dummy_input,
    "GFPGANv1_fixed.onnx",
    opset_version=17,
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={"input": {0: "batch_size"}},
    do_constant_folding=True
)

4. 验证修改效果 #

导出后可使用ONNX Runtime加载模型验证：

import onnxruntime as ort
import numpy as np

session = ort.InferenceSession("GFPGANv1.onnx")
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name

# 测试输入
test_input = np.random.randn(1, 3, 512, 512).astype(np.float32)
output = session.run([output_name], {input_name: test_input})
print("模型加载并推理成功")

内容的提问来源于stack exchange，提问作者Codezilla