使用pytorch_lightning和pytorch-geometric导出ONNX模型遇报错，TorchScript正常

核心原因 #

这个警告是ONNX导出器检测到模型内部存在对输入张量的in-place修改（比如x.add_(y)、x[:] = new_val这类操作），而ONNX的静态计算图语义要求输入张量是只读的，这类修改会破坏图的一致性，导致导出失败。PyTorch Geometric（PyG）的模型和PyTorch Lightning（PL）的封装逻辑容易隐含这类操作。

具体解决方案 #

• 排查并替换in-place操作
  遍历模型forward方法及调用的PyG内置模块，找出所有修改输入张量的代码：

  • 将x += y替换为x = x + y
  • 将x.add_(y)替换为x = x.add(y)
  • 如果必须修改输入内容，先对输入张量执行clone()，比如x_clone = x.clone(); x_clone[:] = new_val，再使用克隆后的张量进行后续计算。
  • 注意PyG的部分Conv层可能隐含in-place操作，可手动封装一层，先克隆输入再传入。

• 调整Lightning模型的导出逻辑
  导出ONNX时直接调用模型的forward方法，避免通过PL的Trainer触发额外的钩子或训练逻辑（这些逻辑可能修改输入数据）。确保forward方法仅基于输入张量生成输出，不修改原始输入。

• 优化ONNX导出参数
  调用torch.onnx.export时添加以下参数：

  torch.onnx.export(
      model,
      dummy_input,  # 这里用拆分为独立张量的输入，比如(x, edge_index)而非完整Data对象
      "model.onnx",
      do_constant_folding=False,  # 关闭常量折叠，规避语义冲突
      dynamic_axes={  # 明确动态维度，适配不同输入规模
          "x": {0: "batch_size"},
          "edge_index": {1: "num_edges"}
      },
      opset_version=17  # 使用较新的opset版本，更好支持PyG的操作
  )
  

• 从TorchScript中转导出ONNX
  既然TorchScript导出正常，可先将模型转为TorchScript，再导出ONNX，间接规避语义冲突：

  # 先导出TorchScript
  scripted_model = torch.jit.script(model)
  # 从TorchScript导出ONNX
  torch.onnx.export(
      scripted_model,
      dummy_input,
      "model.onnx",
      opset_version=17
  )
  

• 拆分PyG输入数据
  不要直接传入PyG的Data对象或dict/tuple形式的完整数据，将输入拆分为独立的张量（比如节点特征x、边索引edge_index等）作为模型的输入参数，确保每个张量在forward中仅被读取，不被修改。

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