2

I was trying to convert my pytorch model to onnx but I am facing RuntimeError: Only tuples, lists and Variables are supported as JIT inputs/outputs. Dictionaries and strings are also accepted, but their usage is not recommended. Here, received an input of unsupported type: DGLHeteroGraph Error

import torch.onnx
model.load_state_dict(state,strict=True)
dummy_input = [get_graph_from_smile(Chem.MolToSmiles(Chem.AddHs(Chem.MolFromSmiles('CC(C)(C)Br')))), get_graph_from_smile(Chem.MolToSmiles(Chem.AddHs(Chem.MolFromSmiles('CC(C)(C)Br'))))]
torch.onnx.export(model, dummy_input, "solubility.onnx")

I have kept the code which I am working, don't know where I am doing wrong

1 Answers1

0

I think that your dummy_input is not something torch supports for export. I guess get_graph_from_smile returns DGLHeteroGraph which torch cannot trace/jit to then export it in ONNX.

I don’t know well the data types you use but in general, you should provide the same kind of input (only a batch with N=1 should be enough unless you want batched execs for the exported NN then use the N you want for the dummy batch) during export as you provide during training.

A classic pattern for a NN taking images with NCHW format 1, 1, 28, 28 e.g. is to provide dummy_input = torch.rand(1, 1, 28, 28)

Edit: found a similar issue online: https://discuss.dgl.ai/t/exporting-graphsage-model-into-onnx-format-train-sampling-py/1263

It could be that the framework is not well supported for onnx export

IceTDrinker
  • 151
  • 5