Search code examples
pytorchonnxseq2seq

ONNX export of Seq2Seq model - issue with decoder input length


I have a simple Seq2Seq model trained according to "Attention is all you need" and implemented using PyTorch. The model works fine. I decided to export it to ONNX. I exported the encoder and decoder separately. When using the ONNX model, the encoder works fine. However, the decoder only works for the same length of input sequence for which it was exported. For all other lengths, it ends with an error: The input tensor cannot be reshaped to the requested shape. Input shape:{2,1,300}, requested shape:{1,20,15} The embedding size is 300. I don't think this is a problem with dynamic axes, as I set them correctly after the first failure. I tried to solve the problem by using a constant input length for the decoder and applying a mask, but this resulted in nonsensical output. Thank you in advance for any tips.


Solution

  • ok. That's a problem in Pytorch 2.2.x onnx exporter. There is no such a problem with Pytorch 1.X.