Exporting GraphSAGE model into onnx format : train_sampling.py

Hello. I am trying to export GraphSAGE model of train_sample.py in onnx format. I tried using pytorch.onnx.export with no success. GraphSAGE model’s (in train_sample.py) forward function takes 2 inputs (input graph as DGLblock and input features):

def forward(self, blocks, x):
    h = x
    for l, (layer, block) in enumerate(zip(self.layers, blocks)):
        h = layer(block, h)
        if l != len(self.layers) - 1:
            h = self.activation(h)
            h = self.dropout(h)
    return h

I tried to export model in onnx format as

torch.onnx.export(model, (test_g ,test_g.ndata['features']), "graphsage.onnx", verbose=True, input_names=['input features'], output_names=['output labels'])

This results in following error

RuntimeError: Only tuples, lists and Variables supported as JIT inputs/outputs. Dictionaries and strings are also accepted but their usage is not recommended. But got unsupported type DGLHeteroGraph

I also tried to convert DGLHeteroGraph using to_networkx() but that didn’t help either. Problem is with graph as input. torch.onnx.export() expects input to be in list, variable or tuple form. This is NOT a DGL issue but I am wondering if you guys have some support for simplifying it. Any recommendation on how to export this model into onnx format ? @BarclayII: do you have any trick up in your sleeves ? Thanks in advance.

It’s not supported. Why do you want to do this? Custom op used in dgl is not well supported by onnx I think

Thanks for your response @VoVAllen . Appreciate it. We are trying to do some analysis for our project which requires us to import a model in prototxt, onnx or xml format. So is prototxt or xml export possible?


How could those framework support custom operator? DGL has some custom operators on the graph structure, which is our major concerns on how to export them to other libraries.

@VoVAllen: I am sorry. I missed your response. You are right, onnx and other frameworks won’t support DGL ops. For this part of analysis that I am trying to do, I think I will be better off to use a pure pytorch/TF version of GraphSAGE . Thanks for all help. Much appreciated.

1 Like