Torch.onnx.export and pack_sequence/unpack_sequence

So, I spent a few (good) hours this weekend getting a model exported via torch.onnx.export which contains pack_sequence and unpack_sequence. I managed to get it working after a while and only after rewriting these two little functions myself directly using tensors/for loops.

Made a small repo, perhaps someone might be intersted in the future until this is supported ootb: Files · master · Pîrvu Mihai Cristian / torch-pack-sequence-onnx-export · GitLab

TL;DR

  • pack_sequence → manually pad with zeros and make a for loop
  • unpack_sequence → take the last step from the output of the rnn model
  • only works after you jit the model
  • list of tensor inputs are required and proper typing is necessary throughout the code

As a sidenote: onnxruntime works well to pass List[Tensor] as inputs, but onnxruntime-web doesn’t, but that is a different runtime/project, outside of torch: How do I pass a list of tensors in onnxruntime-web? · Discussion #14828 · microsoft/onnxruntime · GitHub