Support LSTMCell layer from PyTorch
- 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
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