Deploying LSTM model to iOS with ONNX

username: dnnagy
cooked: <p>Hello,</p>
<p>I have trained a model which contains LSTMs:</p>
<pre><code class="lang-auto">class Net(nn.Module):
    def __init__(self, W, H, C, num_classes, rnn_type='LSTM', rnn_hidden=128, rnn_layers=2, rnn_b=False):
        super(Net, self).__init__()
        self.W = W
        self.H = H
        self.C = C
        self.num_classes = num_classes

        self.conv_stack = nn.Sequential(
            nn.Conv2d(1, 16, kernel_size=3, padding=1, stride=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=3),
            nn.Conv2d(16, 32, kernel_size=3, padding=1, stride=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=3),
            nn.Conv2d(32, 32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=3),
            nn.Dropout(0.3)
        )

        Cout, Hout, Wout = self.get_conv_out_features((self.C, self.H, self.W))
        self.rnn_in = Cout*Hout*Wout
        self.rnn_hidden = rnn_hidden
        self.rnn_layers = rnn_layers
        self.rnn_b = rnn_b
        
        self.rnn = nn.LSTM(input_size=self.rnn_in,
                           hidden_size=self.rnn_hidden,
                           num_layers=self.rnn_layers,
                           batch_first=True,
                           bidirectional=self.rnn_b)

        self.hidden2out = nn.Sequential(
            nn.Dropout(0.25),
            nn.Linear(in_features=self.rnn_hidden, out_features=64),
            nn.ReLU(),
            nn.Linear(in_features=64, out_features=num_classes),
            nn.LogSoftmax()
        )
    
    def get_conv_out_features(self, input_shape):
        """
        Calculates the output shape of the convolution stack given the input shape
        :param input_shape: tuple of ints (C_{in}, H_{in}, W_{in})
        :return: (C_{out}, H_{out}, W_{out})
        """
        (C, H, W) = input_shape
        for layer in self.conv_stack:
            if isinstance(layer, nn.Conv2d) or isinstance(layer, nn.MaxPool2d):
                kernel_size = layer.kernel_size if type(layer.kernel_size)==tuple else (layer.kernel_size, layer.kernel_size)
                stride = layer.stride if type(layer.stride)==tuple else (layer.stride, layer.stride)
                padding = layer.padding if type(layer.padding)==tuple else (layer.padding, layer.padding)
                dilation = layer.dilation if type(layer.dilation)==tuple else (layer.dilation, layer.dilation)
                
                H = int((H + 2*padding[0] - dilation[0]*(kernel_size[0]-1)-1)/stride[0] + 1)
                W = int((W + 2*padding[1] - dilation[1]*(kernel_size[1]-1)-1)/stride[1] + 1)
            if isinstance(layer, nn.Conv2d):
                C = layer.out_channels
        return (max(C,1), max(H,1), max(W,1))

    def forward(self, x, sequence_lengths=None):
        N, T = x.shape[0], x.shape[1]

        x = x.view(-1, self.W, self.H, self.C)
        x = x.transpose(1,3) # conv_stack requires (N, C, H, W)
        x = self.conv_stack(x)
        x = x.view(N, T, -1)
        x, (hn, cn) = self.rnn(x)
        x = self.hidden2out(hn[-1])
        return x
</code></pre>
<p>Now I need to deploy the trained model for inference to iOS via ONNX.<br>
However I get the following error:<br>
<code>RuntimeError: ONNX export failed: Couldn't export operator aten::lstm</code></p>