Bidirectional RNN #708

NeroBlackstone · 2024-06-16T12:17:01Z

issue #687

Please confirm whether the interface meets the requirements. Thank you.

codecov · 2024-06-16T12:27:18Z

Codecov Report

Attention: Patch coverage is 0% with 6 lines in your changes missing coverage. Please review.

Project coverage is 76.81%. Comparing base (1a61165) to head (e7ae725).

Files	Patch %	Lines
src/layers/recurrent.jl	0.00%	6 Missing ⚠️

❗ There is a different number of reports uploaded between BASE (1a61165) and HEAD (e7ae725). Click for more details.

HEAD has 25 uploads less than BASE

Flag BASE (1a61165) HEAD (e7ae725)

38 13

Additional details and impacted files

@@             Coverage Diff             @@
##             main     #708       +/-   ##
===========================================
- Coverage   96.40%   76.81%   -19.60%     
===========================================
  Files          54       54               
  Lines        2726     2730        +4     
===========================================
- Hits         2628     2097      -531     
- Misses         98      633      +535

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

NeroBlackstone · 2024-06-17T16:53:39Z

test/layers/recurrent_tests.jl

+ if mode != "AMDGPU"
+ # gradients test failed after vcat
+ # __f = p -> sum(Base.Fix1(sum, abs2), first(bi_rnn(x, p, st)))
+ # @eval @test_gradients $__f $ps atol=1e-2 rtol=1e-2 gpu_testing=$ongpu


The gradient test is incorrect for the matrix after vcat. I'm not sure if the gradient test code is written incorrectly, or if my Bidirectional implementation is incorrect

Which backend was failing? The basic pipeline is compute with Zygote. Then compute with FiniteDifferences, ReverseDiff, Tracker,... and compare them to Zygote.

NeroBlackstone · 2024-06-17T16:55:30Z

@avik-pal Hi! I have completed the implementation of Bidirectional and written test code, trying to be as equivalent to Keras' API. Please review the changes, thank you!

NeroBlackstone · 2024-06-19T03:03:24Z

😭😭😭Any suggestions? I will fix it today

src/layers/recurrent.jl

avik-pal · 2024-06-19T03:16:13Z

src/layers/recurrent.jl

+
+## Parameters
+
+ - Same as `cell` and `backward_cell`.


It should be NamedTuple with cell and backward_cell

src/layers/recurrent.jl

avik-pal · 2024-06-19T03:18:00Z

src/layers/recurrent.jl

+ backward_cell::Union{AbstractRecurrentCell, Nothing}=nothing,
+ merge_mode::Union{Function, Nothing}=vcat,
+ ordering::AbstractTimeSeriesDataBatchOrdering=BatchLastIndex())
+ if !isnothing(backward_cell) && cell.in_dims != backward_cell.in_dims


check via backward_cell !== nothing, there used to be some performance consideration for not using isnothing (not sure if that is true for current julia versions but better be safe)

src/layers/recurrent.jl

avik-pal · 2024-06-19T03:21:35Z

src/layers/recurrent.jl

+ backward_rnn_layer = isnothing(backward_cell) ? deepcopy(layer) :
+ Recurrence(backward_cell; return_sequence=true, ordering)
+ fuse_op = isnothing(merge_mode) ? nothing : Broadcast.BroadcastFunction(merge_mode)
+ return Parallel(


I was thinking maybe we create a struct

struct BidirectionalRNN <: .... model end (rnn::...)(...) = rnn.model(....)

This is mostly to be consistent with other wrapper layers in Lux which return the same type and not Parallel (as in this case)

avik-pal · 2024-06-19T03:23:55Z

test/layers/recurrent_tests.jl

+ if mode != "AMDGPU"
+ # gradients test failed after vcat
+ # __f = p -> sum(Base.Fix1(sum, abs2), first(bi_rnn(x, p, st)))
+ # @eval @test_gradients $__f $ps atol=1e-2 rtol=1e-2 gpu_testing=$ongpu


Which backend was failing? The basic pipeline is compute with Zygote. Then compute with FiniteDifferences, ReverseDiff, Tracker,... and compare them to Zygote.

NeroBlackstone · 2024-06-20T14:47:13Z

test/layers/recurrent_tests.jl

+ y_, st__ = bi_rnn_no_merge(x, ps, st)
+
+ @jet bi_rnn(x, ps, st)
+ @jet bi_rnn_no_merge(x, ps, st)


@avik-pal Hi, I have updated the implementation, but the @jet test fails and I can't find the reason. Could you please take a look at it when you have time? 😭😭😭

Yeah the tests wont really tell much here. Install JET.jl and run @report_call bi_rnn_no_merge(x, ps, st) target_modules=[Lux, LuxLib, LuxCore] and @report_opt bi_rnn_no_merge(x, ps, st) target_modules=[Lux, LuxLib, LuxCore]. This will give you a list of function calls that cause the test to fail

@avik-pal Hi, I tried @report_call bi_rnn_no_merge(x, ps, st) target_modules=[Lux, LuxLib, LuxCore] but it report @report_call expects only one non-keyword argument, so I remove target_modules=[Lux, LuxLib, LuxCore]. And I replace @jet to @report_call and @report_opt, all test passed and no error throw.

Is this LuxTestUtils bug?,,, Could you please check this branch locally? Thank you very much!!

no don't replace @jet with report_* functions, you need to run those 2 in the REPL / VSCode. That will point to the correct line. @report_call target_modules=[Lux, LuxLib, LuxCore] bi_rnn_no_merge(x, ps, st) https://aviatesk.github.io/JET.jl/dev/tutorial/#Analyse-methods-with-@report_call

NeroBlackstone · 2024-06-21T20:47:34Z

@avik-pal 🥹🥹Hi, could you please help me review the PR and find the reason why @jet test failed? I haven't any idea about that.

Thank you very much ♥️♥️♥️♥️

NeroBlackstone · 2024-06-23T17:18:01Z

src/layers/recurrent.jl

+ model::Parallel
+end
+
+(rnn::BidirectionalRNN)(x, ps, st::NamedTuple) = rnn.model(x, ps, st)


@avik-pal Hi! JET.jl @report_opt bi_rnn(x, ps, st) report:

(::BidirectionalRNN)(x::Array{…}, ps::@NamedTuple{…}, st::@NamedTuple{…}) @ Lux ./Lux.jl/src/layers/recurrent.jl:690 │ runtime dispatch detected: %1::Parallel(x::Array{Float32, 3}, ps::@NamedTuple{layer_1::@NamedTuple{…}, layer_2::@NamedTuple{…}}, st::@NamedTuple{layer_1::@NamedTuple{…}, layer_2::@NamedTuple{…}})::Any

I have no idea why it say runtime dispatch here... Chain has same function signature (c::Chain)(x, ps, st::NamedTuple) = applychain(c.layers, x, ps, st)

It should have more information in the report

🥲 I'm sorry, it has only two lines report, like I post above

NeroBlackstone · 2024-06-25T11:52:19Z

@avik-pal Sorry to bother you. I still don't know how to solve runtime dispatch error for this julia code...

(rnn::BidirectionalRNN)(x, ps, st::NamedTuple) = rnn.model(x, ps, st)

julia> @report_opt bi_rnn(x, ps, st)
═════ 1 possible error found ═════
┌ (::BidirectionalRNN)(x::Array{…}, ps::@NamedTuple{…}, st::@NamedTuple{…}) @ Lux ./Lux.jl/src/layers/recurrent.jl:690
│ runtime dispatch detected: %1::Parallel(x::Array{Float32, 3}, ps::@NamedTuple{layer_1::@NamedTuple{…}, layer_2::@NamedTuple{…}}, st::@NamedTuple{layer_1::@NamedTuple{…}, layer_2::@NamedTuple{…}})::Any
└────────────────────

avik-pal · 2024-06-25T23:43:39Z

I will take a look on the weekend

NeroBlackstone · 2024-06-30T06:57:04Z

I will take a look on the weekend

😥Hi, could you help me review this PR?..

avik-pal · 2024-06-30T19:30:39Z

Fix the gradient tests and it should be fine. They are probably originating from lazy reverse rrules for Zygote not being defined for GPU arrays

NeroBlackstone · 2024-07-01T11:34:47Z

Fix the gradient tests and it should be fine. They are probably originating from lazy reverse rrules for Zygote not being defined for GPU arrays

Thank you for your help!

Some gradient tests still failed at here, I have no idea about how to set lazy reverse rrules... I guess the rest of the implementation may have to be left to you, thanks!

19:22:03 | maxrss 20.0% | mem 67.2% | DONE  (1/1) test item "Bidirectional" 112.9 secs (68.5% compile, 0.1% recompile, 6.2% GC), 188.79 M allocs (13.948 GB)
Test Summary:                          | Pass  Error  Total     Time
ReTestItem Tests                       |  120     12    132  2m08.7s
  Bidirectional                        |   60      6     66  2m03.7s
    cpu                                |   30      3     33  1m08.8s
      cell: RNNCell                    |   10      1     11    52.2s
      cell: LSTMCell                   |   10      1     11     8.0s
      cell: GRUCell                    |   10      1     11     8.6s
    cuda                               |   30      3     33    43.5s
      cell: RNNCell                    |   10      1     11    27.5s
      cell: LSTMCell                   |   10      1     11     7.8s
      cell: GRUCell                    |   10      1     11     8.3s
  Lux                                  |   60      6     66  2m05.6s
    test                               |   60      6     66         
      test/layers                      |   60      6     66         
        test/layers/recurrent_tests.jl |   60      6     66         
          Bidirectional                |   60      6     66  2m03.7s
            cpu                        |   30      3     33  1m08.8s
              cell: RNNCell            |   10      1     11    52.2s
              cell: LSTMCell           |   10      1     11     8.0s
              cell: GRUCell            |   10      1     11     8.6s
            cuda                       |   30      3     33    43.5s
              cell: RNNCell            |   10      1     11    27.5s
              cell: LSTMCell           |   10      1     11     7.8s
              cell: GRUCell            |   10      1     11     8.3s
ERROR: LoadError: Some tests did not pass: 120 passed, 0 failed, 12 errored, 0 broken.
in expression starting at /home/nero/Documents/github/Lux.jl/test/runtests.jl:75

NeroBlackstone · 2024-07-03T04:56:35Z

@avik-pal Hi, I'm sorry to bother you again. I want to continue to push this PR, but it seems that there is a limit to what I can do. Is "lazy reverse rrules" a feature that Zygote is missing? Do I need to open an issue for Zygote?

avik-pal · 2024-07-04T02:37:03Z

Kind of, but not worth opening a Zygote issue for this.

If you look at https://buildkite.com/julialang/lux-dot-jl/builds/2994#01906a93-9e6b-4704-a6a1-d3b8e82bb694/350-1748, it is saying that we doing a broadcast vcat of Vector{<:CuArray} and Iterators.Reverse{<:CuArray}. Zygote doesn't have a rule for that.

The easiest way to resolve this would be to make the Iterators.Reverse into a Vector, since that is effectively materializing a vector of pointers, it is not expensive either.

NeroBlackstone · 2024-07-04T10:03:40Z

make the Iterators.Reverse into a Vector,

@avik-pal Thank you very much for pointing me in the right direction, but I can't really understand "make the Iterators.Reverse into a Vector" here, I think Iterators.Reverse already support vector:

julia> vec = [1,2,3,4,5]
julia> foreach(println, Iterators.reverse(vec))
5
4
3
2
1

Could you give me a few inputs and outputs as examples, and the signature of the function to define? I can implement it if I could

NeroBlackstone changed the title ~~[WIP] Bidirectional RNN~~ Bidirectional RNN Jun 17, 2024

NeroBlackstone commented Jun 17, 2024

View reviewed changes

avik-pal reviewed Jun 19, 2024

View reviewed changes

avik-pal force-pushed the Bidirectional branch from 6f75695 to a73bbf8 Compare June 19, 2024 03:24

avik-pal linked an issue Jun 19, 2024 that may be closed by this pull request

Feature request: Bidirectional for RNN layer. #687

Open

NeroBlackstone commented Jun 20, 2024

View reviewed changes

NeroBlackstone commented Jun 23, 2024

View reviewed changes

avik-pal force-pushed the main branch from dcab297 to ced869f Compare June 28, 2024 17:00

avik-pal force-pushed the Bidirectional branch 5 times, most recently from b97e29a to e7ae725 Compare June 30, 2024 19:15

NeroBlackstone added 6 commits July 3, 2024 22:17

designe Bidirectional interface

dd8b7a5

add applybidirectional() TODO

2d11c59

update

07940cb

update Bidirectional implementation

b8d8392

update Bidirectional implementation

1531020

update

8dea8d3

NeroBlackstone and others added 6 commits July 3, 2024 22:17

update test

1c987b3

add Bidirectional test

52cd434

update implementation

996a56f

remove isnothing

82d4bd3

update birnn docs

541e130

Re-add the tests

39e3703

avik-pal force-pushed the Bidirectional branch from e7ae725 to 39e3703 Compare July 4, 2024 05:17

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Bidirectional RNN #708

Bidirectional RNN #708

NeroBlackstone commented Jun 16, 2024

codecov bot commented Jun 16, 2024 •

edited

Loading

NeroBlackstone Jun 17, 2024

avik-pal Jun 19, 2024

NeroBlackstone commented Jun 17, 2024

NeroBlackstone commented Jun 19, 2024

avik-pal Jun 19, 2024

avik-pal Jun 19, 2024

avik-pal Jun 19, 2024

avik-pal Jun 19, 2024

NeroBlackstone Jun 20, 2024 •

edited

Loading

avik-pal Jun 22, 2024

NeroBlackstone Jun 23, 2024

avik-pal Jun 23, 2024

NeroBlackstone commented Jun 21, 2024

NeroBlackstone Jun 23, 2024

avik-pal Jun 23, 2024

NeroBlackstone Jun 24, 2024

NeroBlackstone commented Jun 25, 2024

avik-pal commented Jun 25, 2024

NeroBlackstone commented Jun 30, 2024

avik-pal commented Jun 30, 2024

NeroBlackstone commented Jul 1, 2024

NeroBlackstone commented Jul 3, 2024

avik-pal commented Jul 4, 2024

NeroBlackstone commented Jul 4, 2024

Bidirectional RNN #708

Are you sure you want to change the base?

Bidirectional RNN #708

Conversation

NeroBlackstone commented Jun 16, 2024

codecov bot commented Jun 16, 2024 • edited Loading

Codecov Report

Choose a reason for hiding this comment

Choose a reason for hiding this comment

NeroBlackstone commented Jun 17, 2024

NeroBlackstone commented Jun 19, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

NeroBlackstone Jun 20, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

NeroBlackstone commented Jun 21, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

NeroBlackstone commented Jun 25, 2024

avik-pal commented Jun 25, 2024

NeroBlackstone commented Jun 30, 2024

avik-pal commented Jun 30, 2024

NeroBlackstone commented Jul 1, 2024

NeroBlackstone commented Jul 3, 2024

avik-pal commented Jul 4, 2024

NeroBlackstone commented Jul 4, 2024

codecov bot commented Jun 16, 2024 •

edited

Loading

NeroBlackstone Jun 20, 2024 •

edited

Loading