I am having trouble getting a model with several LSTMs to export to ONNX properly. The main issue is that I intend to use the model in an online fashion, i.e. feeding in one frame of data at a time.
My LSTM code is similar to the following:
class MyLSTM(torch.nn.Module):
def __init__(self, dim_in, dim_out):
super().__init__()
self.deployed = False
self.hidden = torch.zeros(1, 1, dim_out)
self.cell = torch.zeros(1, 1, dim_out)
self.lstm = torch.nn.LSTM(input_size=dim_in,
hidden_size=dim_out,
batch_first=True,
bidirectional=False)
def deploy(self):
self.deployed = True
def forward(self, x):
if self.deployed:
out, (self.hidden, self.cell) = self.lstm(x, (self.hidden, self.cell))
else:
out, _ = self.lstm(x)
return out
Obviously, before I export the model to ONNX, I call deploy(). The PyTorch model works as expected, and I even tried saving it as a ScriptModule with torch.jit. However, when I load and attempt to use the exported model using onnxruntime, it’s behavior suggests that it never updates the hidden/cell state.
Is it possible to do this? If so, how?