Output Difference between exported ONNX PyG Model and Normal pyG Model

[Harikaraja]

🐛 Describe the bug

I have a custom PyG model with 2 GCN layers. Here is the architecture of my model

class Custom_GGNN(nn.Module):
    def __init__(self):
        super(Custom_GGNN, self).__init__()
        self.ggnn = GCN(in_channels=10, hidden_channels=10, out_channels=10, num_layers=2, norm='BatchNorm', dropout=0.2)
    
    def forward(self, input_tensor):
        num_nodes = 5  # Number of nodes
        in_channels = 10  # Number of input features per node
        state_size = num_nodes * in_channels
        num_edge = 40  # Number of edges
     
        input_tensor = input_tensor.to(next(self.parameters()).device)
        input_tensor = input_tensor.view(-1, state_size + num_edge)
        print("input_tensor inside forward function: ",input_tensor)
        ggnn_input_x = input_tensor[:, :state_size]
        ggnn_input_x = ggnn_input_x.reshape(input_tensor.shape[0], num_nodes, in_channels)  # (num_nodes, in_channels)
     
        edge_index_flat = input_tensor[:, state_size : state_size + num_edge]
        edge_index = edge_index_flat.reshape(input_tensor.shape[0], 2, -1).long()  # (2, num_edges)
       
        batch_size = input_tensor.shape[0]
        data_list = [Data(x=ggnn_input_x[i], edge_index=edge_index[i]) for i in range(batch_size)]  
        batch_data = Batch.from_data_list(data_list=data_list)

        node_mat = self.ggnn(batch_data.x, batch_data.edge_index)
        node_mat = node_mat.reshape(batch_size, num_nodes, in_channels)
        
        input_state_list = node_mat 
        print("Output of Normal GGNN model: ",input_state_list)
        return input_state_list,input_state_list,input_state_list

Now I am trying to export this custom PyG model using ONNX version 1.16.3 as below

model = Custom_GGNN()
model.eval()
model = model.to('cuda')
num_nodes = 5

input_shape = (90,)
state = torch.randn(num_nodes,10)
state_flat = state.flatten()
print("state_flat is: ",state_flat)

edges = np.array([(i, j) for i in range(num_nodes) for j in range(num_nodes) if i != j])

edges_tensor = torch.tensor(edges, dtype=torch.long).reshape(-1)

input_tensor = torch.cat((
    state_flat, 
    edges_tensor
), dim=0)

print("input_tensor: ",input_tensor)
torch.onnx.export(model,              
                  (input_tensor,),                       
                  "dump_custom_pyG_model.onnx",
                  export_params=True,       
                  opset_version=16,          
                  do_constant_folding=True,  
                  input_names=['input_tensor'],
                  output_names=['seed_embeddings','seed_embeddings','seed_embeddings'])

I am giving the same input to both the exported PyG model and Normal PyG model. The weights of the exported model and Normal model are same but the outputs of them are completely different.

My Input to the model is:

The outputs are as below:

Output from Normal PyG model

Output from exported ONNX pyG model

Versions

[pip3] torch-geometric==2.3.0
ONNX version: 1.16.3