facebookresearch / facebookresearch/fvcore
flop count analysis of LSTM layers
- Dominant language
- Python
- Stars
- 2.3k
- Forks
- 236
- PR merge metrics
- No merged PRs in 30d
Description
Hello, it seems that LSTM layers are not yet supported for the `fvcore.nn.FlopCountAnalysis` method:
```python
import torch
from fvcore.nn import FlopCountAnalysis
from torch import nn
class ToyLSTMModel(nn.Module):
def __init__(self):
super().__init__()
self.rnn = nn.LSTM(10, 20, 1)
def forward(self, x):
h0 = torch.randn(1, 3, 20)
c0 = torch.randn(1, 3, 20)
output, _ = self.rnn(x, (h0, c0))
return output
model = ToyLSTMModel()
example_input = torch.randn(5, 3, 10)
print(FlopCountAnalysis(model, example_input).by_module())
```
gives:
```
Unsupported operator aten::randn encountered 2 time(s)
Unsupported operator aten::lstm encountered 1 time(s)
Counter({'': 0, 'rnn': 0})
```
While the same works an LSTM cell:
```python
import torch
from fvcore.nn import FlopCountAnalysis
from torch import nn
class ToyLSTMModel(nn.Module):
def __init__(self):
super().__init__()
self.rnn = nn.LSTMCell(10, 20)
def forward(self, x):
hx = torch.randn(3, 20)
cx = torch.randn(3, 20)
output = []
for i in range(x.size()[0]):
hx, cx = self.rnn(x[i], (hx, cx))
output.append(hx)
output = torch.stack(output, dim=0)
return output
model = ToyLSTMModel()
example_input = torch.randn(5, 3, 10)
print(FlopCountAnalysis(model, example_input).by_module())
```
output:
```
Unsupported operator aten::randn encountered 2 time(s)
Unsupported operator aten::add_ encountered 10 time(s)
Unsupported operator aten::unsafe_chunk encountered 5 time(s)
Unsupported operator aten::sigmoid_ encountered 15 time(s)
Unsupported operator aten::tanh_ encountered 5 time(s)
Unsupported operator aten::mul encountered 15 time(s)
Unsupported operator aten::tanh encountered 5 time(s)
Counter({'': 36000, 'rnn': 36000})
```
Is there any particular reason for that? The number of FLOPS of the LSTM layer should be the same than from one LSTM cell times the number of time steps.
Contributor guide
Assessment
This issue has not been assessed yet.