Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Adding support for folding RNNs over 3d arrays #1686

Merged
merged 15 commits into from
Sep 14, 2021
Merged

Adding support for folding RNNs over 3d arrays #1686

Changes from 8 commits
Commits
Show all changes
15 commits
Select commit Hold shift + click to select a range
f0589d5
Adding FoldedRNN
mkschleg Jul 31, 2021
68219e7
Merge branch 'master' of github.com:FluxML/Flux.jl into FoldedRNN
mkschleg Aug 2, 2021
e84d043
Changed some of how folded works.
mkschleg Aug 2, 2021
40dbdd3
Changed which hs are concatenated.
mkschleg Aug 2, 2021
e273fc4
Minor changes.
mkschleg Aug 2, 2021
f42f174
Adding support for tuple states.
mkschleg Aug 2, 2021
76ebec7
Tested, changed `1:size` to use `axes`.
mkschleg Aug 2, 2021
bf07ee1
Removed Folded, implemented with mapslices for now.
mkschleg Aug 15, 2021
25a3a3c
Apply suggestions from code review
mkschleg Aug 17, 2021
f4076b9
Changed to stack implementation, removed folded apis for now.
mkschleg Aug 17, 2021
5734435
Changing to reduce implementation per the comparison.
mkschleg Aug 23, 2021
2ce5618
Added docs and tests. GPU test needs to be vetted.
mkschleg Aug 24, 2021
808517e
Apply suggestions from code review
mkschleg Aug 25, 2021
226524c
Cleaned up BPTT-3D test.
mkschleg Aug 25, 2021
d1b1daf
use views
DhairyaLGandhi Aug 30, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 12 additions & 0 deletions src/layers/recurrent.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,7 @@ rnn.state = hidden(rnn.cell)
reset!(m::Recur) = (m.state = m.cell.state0)
reset!(m) = foreach(reset!, functor(m)[1])


# TODO remove in v0.13
function Base.getproperty(m::Recur, sym::Symbol)
if sym === :init
Expand All @@ -67,6 +68,13 @@ end

flip(f, xs) = reverse(f.(reverse(xs)))

function (m::Recur)(x::AbstractArray{<:Number, 3})
mkschleg marked this conversation as resolved.
Show resolved Hide resolved
h_ret = mapslices(m, x, dims=[1, 3])
mkschleg marked this conversation as resolved.
Show resolved Hide resolved
sz = size(x)
mkschleg marked this conversation as resolved.
Show resolved Hide resolved
return h_ret
end


# Vanilla RNN

struct RNNCell{F,A,V,S}
Expand Down Expand Up @@ -102,6 +110,7 @@ The most basic recurrent layer; essentially acts as a `Dense` layer, but with th
output fed back into the input each time step.
"""
RNN(a...; ka...) = Recur(RNNCell(a...; ka...))
FoldedRNN(args...; kwargs...) = FoldedRecur(Flux.RNNCell(args...; kwargs...))
Recur(m::RNNCell) = Recur(m, m.state0)

# TODO remove in v0.13
Expand Down Expand Up @@ -162,6 +171,7 @@ See [this article](https://colah.github.io/posts/2015-08-Understanding-LSTMs/)
for a good overview of the internals.
"""
LSTM(a...; ka...) = Recur(LSTMCell(a...; ka...))
FoldedLSTM(args...; kwargs...) = FoldedRecur(Flux.LSTMCell(args...; kwargs...))
mkschleg marked this conversation as resolved.
Show resolved Hide resolved
Recur(m::LSTMCell) = Recur(m, m.state0)

# TODO remove in v0.13
Expand Down Expand Up @@ -227,6 +237,7 @@ See [this article](https://colah.github.io/posts/2015-08-Understanding-LSTMs/)
for a good overview of the internals.
"""
GRU(a...; ka...) = Recur(GRUCell(a...; ka...))
FoldedGRU(args...; kwargs...) = FoldedRecur(Flux.GRUCell(args...; kwargs...))
Recur(m::GRUCell) = Recur(m, m.state0)

# TODO remove in v0.13
Expand Down Expand Up @@ -281,6 +292,7 @@ See [this article](https://colah.github.io/posts/2015-08-Understanding-LSTMs/)
for a good overview of the internals.
"""
GRUv3(a...; ka...) = Recur(GRUv3Cell(a...; ka...))
FoldedGRUv3(args...; kwargs...) = FoldedRecur(Flux.GRUv3Cell(args...; kwargs...))
Recur(m::GRUv3Cell) = Recur(m, m.state0)


Expand Down