Hello,
I have trained a model which contains LSTMs:
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
Now I need to deploy the trained model for inference to iOS via ONNX.
However I get the following error:
RuntimeError: ONNX export failed: Couldn't export operator aten::lstm