Add Support for torch.nn.LSTM and torch.nn.LSTMCell Modules

[coreylammie]

Support should be added for torch.nn.LSTM and torch.nn.LSTMCell modules.

[frmar440]

Hi! Has the implementation of LSTM modules been started?

[coreylammie]

Hi @frmar440,

I have implemented the base code/logic for a RNN module here: https://github.com/coreylammie/MemTorch/blob/master/memtorch/mn/RNN.py. Although currently, only legacy forward passes are currently enabled, it should be fairly straight forward to extend support for LSTM modules.

The required forward pass for both RNN and LSTM modules can be implemented similarly to as done here for linear layers. The code/logic above (i.e., https://github.com/coreylammie/MemTorch/blob/master/memtorch/mn/RNN.py) deconstructs the inference routine of a RNN to a set of Vector Matrix Multiplications (VMMs), which can be performed in-memory, and other supporting operations/logic.

Please let me know if you have any further questions! If someone is able to work on adding this support, I'm more than happy to help out where possible.

Kind Regards,

Corey.