ConvRNN Cell and ConvLSTM Cell about flux.jl

[darsnack]

This is not true for the current LSTM implementation but could be done for Recur. Is this something people use or can we close as a non-priority?

[MikeInnes]

Our LSTM implementation might be general enough that you can just throw a convolution into it. That would be a nice bit of reuse.

[tbenst]

Definitely something that people use. Have a need for a quick implementation of this for a image denoising application in a Julia codebase, and would love to see in Flux :). seems like PyCall + pytorch is my best bet for now (e.g. https://towardsdatascience.com/video-prediction-using-convlstm-with-pytorch-lightning-27b195fd21a2)

[ToucheSir]

At the very least, I think it would be a good contribution to the model zoo. Whether ConvLSTM should be a core layer would require more discussion (do you do ConvRNN? ConvGRU? Is it possible to generalize the cells so they works with conv and matmul, etc.)

[ToucheSir]

Following up on this, Flax does an interesting thing by using Dense layers instead of explicitly writing out the matmuls: https://github.com/google/flax/blob/main/flax/linen/recurrent.py#L113-L118. My understanding is that the ConvLSTM paper simply subs out those matmuls for convolutions, so it may be possible to parameterize cells in terms of sublayer(s) that compute results from W + x or h.