如何将LightGCN模型的.pth文件转换为.pb文件？转换报错咨询

先搞懂dummy_input的作用 #

dummy_input是模拟模型真实输入的占位张量，ONNX需要通过它追踪模型的前向计算流程，确定每一层的输入输出维度、数据类型，才能生成结构正确的ONNX模型。你之前用的(1,1,28,28)是图像模型的输入形状，和LightGCN的推荐场景输入完全不匹配，这是错误根源之一。

报错原因分析 #

TypeError: LightGCN.forward() missing 1 required positional argument: 'items'说明你的LightGCN模型forward方法需要两个输入参数（通常是users和items，对应用户ID和物品ID），但你只传了一个dummy_input，导致参数缺失。

修正后的完整转换流程 #

1. 准备符合要求的dummy输入 #

根据LightGCN的输入格式，创建匹配的用户ID和物品ID张量（假设你的模型输入是一维的ID张量，batch大小设为32，用户总数user_num，物品总数item_num）：

import torch
import onnx
from onnx_tf.backend import prepare

# 加载模型（注意要实例化模型）
trained_model = Recmodel()
trained_model.load_state_dict(torch.load(weight_file))
trained_model.eval()  # 必须切换到eval模式，避免Dropout、BatchNorm等层影响

# 生成符合模型输入要求的dummy输入
user_num = 1000  # 替换成你的实际用户总数
item_num = 2000  # 替换成你的实际物品总数
dummy_users = torch.randint(0, user_num, (32,), dtype=torch.long)  # 32个随机用户ID
dummy_items = torch.randint(0, item_num, (32,), dtype=torch.long)  # 32个随机物品ID

2. 导出为ONNX模型 #

因为模型有多个输入，要把dummy输入打包成元组传给torch.onnx.export：

# 导出ONNX模型
torch.onnx.export(
    trained_model,
    (dummy_users, dummy_items),  # 多输入用元组传入
    "output/gnn.onnx",
    opset_version=13,  # 指定ONNX算子版本，建议选10以上
    input_names=["users", "items"],  # 给输入命名，方便后续查看
    output_names=["scores"]  # 给输出命名，对应模型的预测得分
)

3. 将ONNX转换为PB文件 #

用onnx-tf库完成转换：

# 加载ONNX模型
onnx_model = onnx.load("output/gnn.onnx")
tf_rep = prepare(onnx_model)

# 导出为PB文件
tf_rep.export_graph("output/gnn.pb")

注意事项 #

• 确保模型已实例化：原代码里trained_model = Recmodel少了括号，应该是trained_model = Recmodel()，否则只是引用类，不是实例化对象。
• 必须切换到eval模式：导出前调用trained_model.eval()，避免训练特有的层影响模型结构。
• 匹配输入的dtype：LightGCN的ID输入一般是torch.long类型，不要用float张量。

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