The model output result is incorrect after exporting the code that called the torch_geometric package as onnx mode

[SpadgerBoy]

from typing import Callable, Union
from torch_geometric.typing import OptPairTensor, Adj, OptTensor, Size

import torch
from torch import Tensor
import torch.nn as nn
import torch.nn.functional as F
from torch_sparse import SparseTensor, matmul
from torch_geometric.nn.conv import MessagePassing

from models.common import MultiLayerPerceptron
import torch.onnx as onnx
import numpy as np
import onnxruntime as ort

class GINConv(MessagePassing):

    def __init__(self, nn: Callable, eps: float = 0., train_eps: bool = False,
                 activation="softplus", **kwargs):
        super(GINConv, self).__init__(aggr='add', **kwargs)
        self.nn = nn  # MLP
        self.initial_eps = eps

        if isinstance(activation, str):
            self.activation = getattr(F, activation)
        else:
            self.activation = None

        if train_eps:
            self.eps = torch.nn.Parameter(torch.Tensor([eps]))
        else:
            self.register_buffer('eps', torch.Tensor([eps]))

    def forward(self, x: Union[Tensor, OptPairTensor], edge_index: Adj,
                edge_attr: OptTensor = None, size: Size = None) -> Tensor:
        if isinstance(x, Tensor):
            x: OptPairTensor = (x, x)

        '''if isinstance(edge_index, Tensor):
            assert edge_attr is not None
            print(x[0].size(-1), edge_attr.size(-1))
            assert x[0].size(-1) == edge_attr.size(-1)
        elif isinstance(edge_index, SparseTensor):
            assert x[0].size(-1) == edge_index.size(-1)
        '''

        out = self.propagate(edge_index, x=x, edge_attr=edge_attr, size=size)

        x_r = x[1]

        if x_r is not None:
            out = out + (1 + self.eps) * x_r

        return self.nn(out)  

    def message(self, x_j: Tensor, edge_attr: Tensor) -> Tensor:

        return self.activation(x_j + edge_attr)

    def __repr__(self):
        return '{}(nn={})'.format(self.__class__.__name__, self.nn)


def code_run(conv_input, edge_index, edge_attr):

    gin = GINConv(MultiLayerPerceptron(128, [128, 128],activation='relu'), activation='relu')

    output = gin(conv_input, edge_index, edge_attr)

    print(output)

    gin.eval()
    torch.onnx.export(gin,
                      (conv_input, edge_index, edge_attr),
                      "gin.onnx",
                      input_names=["conv_input", "edge_index", "edge_attr"],
                      output_names=["output"],
                     )


def onnx_run(conv_input, edge_index, edge_attr):
    model_path = "gin.onnx"
    providers = ['CPUExecutionProvider']
    sess = ort.InferenceSession(model_path, providers=providers)

    input_dict = {
        "conv_input": conv_input,
        "edge_index": edge_index,
        "edge_attr": edge_attr
    }

    output = sess.run(None, input_dict)

    print(output)


if __name__ == '__main__':


    conv_input = ...
    edge_index = ...
    edge_attr =...

    code_run(conv_input, edge_index, edge_attr)

    '''
    tensor([[-16.6324,   4.2831,   4.6628,  ...,  -6.4561,  -0.6724,   1.3455],
        [-14.1573,   3.6967,   4.3299,  ...,  -5.9016,  -0.7184,   1.1747],
        [-12.6177,   3.4545,   3.1947,  ...,  -5.2360,  -0.3239,   1.3559],
        ...,
        [-13.8105,   3.7960,   3.6705,  ...,  -5.5765,  -0.3559,   1.2650],
        [-16.6659,   4.4323,   4.7395,  ...,  -6.4695,  -0.6902,   1.3096],
        [-14.0232,   3.8816,   3.9480,  ...,  -5.7777,  -0.6982,   1.1871]],
       grad_fn=<AddmmBackward0>)
       '''



    conv_input0 = conv_input.numpy()
    edge_index0 = edge_index.numpy()
    edge_attr0 = edge_attr.numpy()

    onnx_run(conv_input0, edge_index0, edge_attr0)

    '''
    [array([[-1.9876492 ,  0.33733028,  0.6393909 , ..., -0.8309902 ,
        -0.21244696,  0.29305094],
       [-1.9019635 ,  0.33062714,  0.6423057 , ..., -0.838735  ,
        -0.21545511,  0.3492959 ],
       [-1.803358  ,  0.48377576,  0.5033815 , ..., -0.84157294,
        -0.24030337,  0.26482943],
       ...,
       [-2.2057521 ,  0.40019044,  0.6885939 , ..., -0.7274695 ,
        -0.2510372 ,  0.17902894],
       [-2.012944  ,  0.36524877,  0.72822946, ..., -0.73702306,
        -0.41365877,  0.28992036],
       [-2.0127623 ,  0.38625523,  0.71876484, ..., -0.7359799 ,
        -0.3926706 ,  0.27177608]], dtype=float32)]
        '''

In my model, I used the torch_geometric package. After exporting the model as onnx, the result of the onnx model is inconsistent with the source code, as explicitly stated in the above code. I don’t know which operator is causing the problem, I hope to get everyone’s guidance.

[rusty1s]

This works for me

from typing import Callable, Union

import numpy as np
import onnxruntime as ort
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.onnx as onnx
# from models.common import MultiLayerPerceptron
from torch import Tensor
from torch_sparse import SparseTensor, matmul

from torch_geometric.nn import MLP, MessagePassing
from torch_geometric.typing import Adj, OptPairTensor, OptTensor, Size


class GINConv(MessagePassing):
    def __init__(self, nn: Callable, eps: float = 0., train_eps: bool = False,
                 activation="softplus", **kwargs):
        super(GINConv, self).__init__(aggr='add', **kwargs)
        self.nn = nn  # MLP
        self.initial_eps = eps

        if isinstance(activation, str):
            self.activation = getattr(F, activation)
        else:
            self.activation = None

        if train_eps:
            self.eps = torch.nn.Parameter(torch.Tensor([eps]))
        else:
            self.register_buffer('eps', torch.Tensor([eps]))

    def forward(self, x: Union[Tensor, OptPairTensor], edge_index: Adj,
                edge_attr: OptTensor = None, size: Size = None) -> Tensor:
        if isinstance(x, Tensor):
            x: OptPairTensor = (x, x)
        '''if isinstance(edge_index, Tensor):
            assert edge_attr is not None
            print(x[0].size(-1), edge_attr.size(-1))
            assert x[0].size(-1) == edge_attr.size(-1)
        elif isinstance(edge_index, SparseTensor):
            assert x[0].size(-1) == edge_index.size(-1)
        '''

        out = self.propagate(edge_index, x=x, edge_attr=edge_attr, size=size)

        x_r = x[1]

        if x_r is not None:
            out = out + (1 + self.eps) * x_r

        return self.nn(out)

    def message(self, x_j: Tensor, edge_attr: Tensor) -> Tensor:

        return self.activation(x_j + edge_attr)

    def __repr__(self):
        return '{}(nn={})'.format(self.__class__.__name__, self.nn)


def code_run(conv_input, edge_index, edge_attr):

    gin = GINConv(MLP([5, 5]))

    output = gin(conv_input, edge_index, edge_attr)

    print(output[:5, :5])

    gin.eval()
    torch.onnx.export(
        gin,
        (conv_input, edge_index, edge_attr),
        "gin.onnx",
        opset_version=16,
        input_names=["conv_input", "edge_index", "edge_attr"],
        output_names=["output"],
    )


def onnx_run(conv_input, edge_index, edge_attr):
    model_path = "gin.onnx"
    providers = ['CPUExecutionProvider']
    sess = ort.InferenceSession(model_path, providers=providers)

    input_dict = {
        "conv_input": conv_input,
        "edge_index": edge_index,
        "edge_attr": edge_attr
    }

    output = sess.run(None, input_dict)

    print(output)


if __name__ == '__main__':

    conv_input = torch.randn(4, 5)
    edge_index = torch.tensor([[0, 1, 1, 2, 2, 3], [1, 0, 2, 1, 3, 2]])
    edge_attr = torch.randn(6, 5)

    code_run(conv_input, edge_index, edge_attr)

    conv_input0 = conv_input.numpy()
    edge_index0 = edge_index.numpy()
    edge_attr0 = edge_attr.numpy()

    onnx_run(conv_input0, edge_index0, edge_attr0)

Can you try to set opset_version=16?