apple / apple/coremltools

Support LSTMCell layer from PyTorch

Open
#1,344 2 comments 0 reactions 0 assignees View on GitHub
LSTM/RNN missing layer type PyTorch (traced) triaged
Dominant language
Python
Stars
5.4k
Forks
850
Avg merge
4d 5h
Merged PRs (30d)
10

Description

## 🌱 Describe your Feature Request
Adding conversion support from PyTorch not only for LSTM, but also for LSTMCell. At least `unsafe_chunk` op support is required.

## Use cases
LSTMCell is needed for streaming(online) processing on device . For example it will be useful for speech processing or any other time series.

## Describe alternatives you've considered
Conversion PyTorch->ONNX->CoreML is working, but ONNX->CoreML will be deprecated.

## Additional context
Working LSTM conversion
```
import torch
from torch import nn
import coremltools as ct

lstm = nn.LSTM(3, 3)
lstm.eval()
inputs = torch.randn(12, 1, 3)
traced_model = torch.jit.trace(lstm, inputs)

ct.convert(
model=traced_model,
inputs=[
ct.TensorType(name="sequence", shape=(ct.RangeDim(1, 50), 1, 3))
]
)
```

Failed conversion of LSTMCell with error `PyTorch convert function for op 'unsafe_chunk' not implemented.`
```
import torch
from torch import nn
import coremltools as ct

class OneStep(nn.Module):
def __init__(self):
super(OneStep, self).__init__()
self.lstmcell = nn.LSTMCell(input_size=3, hidden_size=3)

def forward(self, sequence, hidden_state, cell_state):
return self.lstmcell(sequence, (hidden_state, cell_state))

lstm_cell = OneStep()
lstm_cell.eval()
dummy_input = [torch.zeros((1, 3)), torch.zeros((1, 3)), torch.zeros((1, 3))]
traced_model = torch.jit.trace(lstm_cell, dummy_input)

ct.convert(
model=traced_model,
inputs=[
ct.TensorType(name="sequence", shape=(1, 3)),
ct.TensorType(name="hidden_state", shape=(1, 3)),
ct.TensorType(name="cell_state", shape=(1, 3)),
]
)
```

Contributor guide

Open the contributing guide

Research direction

Start by examining the PyTorch conversion handling for the unsupported `unsafe_chunk` operation and compare it with the working LSTM conversion shown in the issue. Run the supplied traced LSTMCell example; done means conversion succeeds for LSTMCell inputs and preserves the expected outputs.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning
Issue type
Feature
Difficulty
3/5
Estimated time
1-2 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
50/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.