I have a module, whose purpose is to create word features from characters, by passing characters embeddings through an RNN and taking the last hidden state for each word. I am trying to export the module to ONNX. Here is the forward pass (notice that the RNN is initialized with batch_first=False):
def forward(self, inputs: torch.Tensor, lengths: torch.Tensor) -> torch.Tensor:
""""
Args:
inputs (torch.Tensor): Character feature tensor of shape [batch_size x seq_len x max_word_len]
lengths (torch.Tensor): The lengths of each word [batch_size x seq_len]
Returns:
torch.Tensor: The last hidden states of all words.
[batch_size x seq_len x (layers * directions * hidden_size)]
"""
inputs_emb = self.embedding(inputs) # [batch_size x seq_len x max_word_len x embedding_size]
inputs_emb = self.dropout(inputs_emb) # [batch_size x seq_len x max_word_len x embedding_size]
inputs_emb_dims = inputs_emb.size()
inputs_emb_flat = inputs_emb.view(
-1, inputs_emb_dims[2], inputs_emb_dims[3]
) # [(batch_size * seq_len) x max_word_len x embedding_size]
lengths_flat = lengths.view(-1) # [(batch_size * seq_len)]
seq_len_flat = inputs_emb_flat.size()[0] # [(batch_size * seq_len)]
# Move seq dimensions to first prior to RNN
inputs_emb_flat = inputs_emb_flat.permute(1, 0, 2) # [max_word_len, (batch_size * seq_len), embedding_size]
lengths_flat_sorted, lengths_flat_idx_sorted = lengths_flat.sort(dim=0, descending=True)
inputs_emb_flat_sorted = inputs_emb_flat[:, lengths_flat_idx_sorted]
packed = pack_padded_sequence(
inputs_emb_flat_sorted, lengths_flat_sorted.to("cpu"), batch_first=False, enforce_sorted=True
)
if self.rnn.mode == "GRU":
_, h_n = self.rnn(packed) # [(layers * directions) x seq_len_flat x hidden_size]
elif self.rnn.mode == "LSTM":
_, (h_n, _) = self.rnn(packed) # [(layers * directions) x seq_len_flat x hidden_size]
# Move sequence dimension to the first dimension and then
# concatenate the forward/backward hidden states.
h_n = (
h_n.permute(1, 0, 2).contiguous().view(seq_len_flat, -1)
) # [seq_len x (layers * directions * hidden_size)]
# Unsort the hidden states.
_, lengths_flat_idx_unsorted_idx = lengths_flat_idx_sorted.sort(dim=0)
h_n_unsorted = h_n[lengths_flat_idx_unsorted_idx] # [seq_len x (layers * directions * hidden_size)]
# Transform back to original input size
output = h_n_unsorted.view(
inputs_emb_dims[0], inputs_emb_dims[1], -1
) # [batch_size x seq_len x (layers * directions * hidden_size)]
return output
I am using tracing and exporting my model as such:
torch.onnx.export(
model,
input_example,
filepath,
verbose=True,
opset_version=12,
input_names=["char_seqs", "char_seqs_lens"],
output_names=["embeddings"],
dynamic_axes={
"char_seqs": {1: "seq_len"},
"char_seqs_lens": {1: "seq_len"},
"embeddings": {1: "seq_len"},
},
)
When i export the model, I do not get any errors or warnings. However when i try to pass a sequence which has a different length, than the input example, i get the following error:
onnxruntime.capi.onnxruntime_pybind11_state.RuntimeException: [ONNXRuntimeError] : 6 : RUNTIME_EXCEPTION : Non-zero status code returned while running Expand node. Name:'Expand_39' Status Message: /onnxruntime_src/onnxruntime/core/providers/cpu/math/element_wise_ops.h:483 void onnxruntime::BroadcastIterator::Append(int64_t, int64_t) axis == 1 || axis == largest was false. Attempting to broadcast an axis by a dimension other than 1. 500 by 1951
I do not understand why this is not working, since I believe I am taking all the necessary steps to be able to have a dynamic sequence length.
Edit 1
I believe the reason why this is not working, is because I have both a dynamic sequence length (the length of the word in characters), and the length of the word sequence (which is passed as batch_size x seq_len). Any ideas on how to circumvent this problem, would be most welcome.