support LSTM

i use pytorch to build my lstm network, it looks like

    class TestNet(torch.nn.Module):
        def __init__(self):
            super(TestNet, self).__init__()
            self.lstm = nn.LSTM(256,
                                128, 2,
                                batch_first=True, bidirectional=True)
            return
        def forward(self, x):
            self.lstm.flatten_parameters()
            res = self.lstm(x)
            return res

input_tensor = torch.randn(30, 61, 256)

and lstm converter

op = trt.RNNOperation.LSTM
ctx.network.add_rnn_v2(input_tensor._trt, layer_count, hidden_size, max_seq_length, op)

max_seq_length should be input_tensor.shape[0]?

got a error

[TensorRT] ERROR: Parameter check failed at: ../builder/Network.cpp::addRNNCommon::397, condition: input.getDimensions().d[di.seqLen()] == maxSeqLen

and, how could i set a reverse weights?

please any help…

[bfortuner]

If this is your first time converting an LSTM, I found it helpful to use the raw TensorRT API instead to make sure you understand what’s happening underneath the surface. The torch2trt repo is good for basic stuff, but for anything non-trivial, you’ll have to use the TensorRT API directly. I would start there.

https://docs.nvidia.com/deeplearning/sdk/tensorrt-api/python_api/index.html

It will take more time initially, but it will save you tons of time in the future.

[bfortuner]

A small update to forward fixed it for me:

From this

      def forward(self, x):
            self.lstm.flatten_parameters()
            res = self.lstm(x)
            return res

To this:

        def forward(self, x):
            self.lstm.flatten_parameters()
            out, (h0, c0) = self.lstm(x)
            return out