Exporting GraphSAGE model into onnx format : train_sampling.py

Models & Apps
#1

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.

DGL Serving Best Practice
#2

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

#3

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?

#4

Hi,

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.

#5

@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