Export to ONNX multi-input network


(Alex Lopez) #1

Hello,

I am currently working on a project that relies on a Convolutional Neural Network that has two different inputs(two images).

My idea is to obtain some information about the network and I need to export it through torch.onnx.export (ONNX Export). However, I am having some problems when creating and passing as arguments two dummy inputs for torch.onnx.export.

Does any one know if this function supports CNNs with two inputs?

My code is something as follows:

dummy_input1 = torch.randn(1, 3, 224, 224).cuda()
dummy_input2 = torch.randn(1, 150, 224, 224).cuda()
torch.onnx.export(model, (dummy_input1, dummy_input2))

Thanks in advance!

Alex.


(Benjamin) #2

I have run across a similar issue. You could try to change your network a bit and put both pictures in the same array and break them apart in the forward function of your model. So your dummy_input could be shape (1,153,224,224) where channels 151-153 are the “dummy_input1” and 1-150 are “dummy_input2”. You can then reshape within your forward model. This worked for me. Goodluck!