Move `dropout` to NNlib #2150

mcabbott · 2023-01-03T22:47:16Z

julia> model = Chain(Dense(10=>100), Dropout(0.5)); data = rand(Float32, 10, 100);

julia> @btime $model($data);
  min 4.536 μs, mean 11.995 μs (4 allocations, 78.22 KiB)
  min 4.692 μs, mean 14.468 μs (4 allocations, 78.22 KiB)

julia> @btime gradient(m -> sum(m($data)), $model);
  min 45.750 μs, mean 77.092 μs (176 allocations, 251.69 KiB)  # before
  min 32.125 μs, mean 55.435 μs (110 allocations, 209.89 KiB)  # after

julia> trainmode!(model);

julia> @btime $model($data);
  min 14.500 μs, mean 25.182 μs (17 allocations, 156.67 KiB)
  min 10.375 μs, mean 24.783 μs (6 allocations, 117.33 KiB)   # after

mcabbott · 2023-01-03T22:55:58Z

src/layers/normalise.jl

+function Dropout(p::Real; dims=:, rng = default_rng_value())
+  0 ≤ p ≤ 1 || throw(ArgumentError("Dropout expexts 0 ≤ p ≤ 1"))
  Dropout(p, dims, nothing, rng)
 end


Since Dropout is mutable, you can conceivably do d = Dropout(0.0); d.p = 0.1. Is this supported?

If not, then the constructor could do this:

p == 0 && return identity p == 1 && return zero

That would mean that some model construction function which happens to have p==0 will do the right thing, without imposing a branch etc.

I don't know how to feel about this. @darsnack might be able to give a more definite answer because we ran into this question (when is it appropriate for layer constructors to return a different type) a couple times during the Metalhead rework.

6c84a6c adds this for Dropout(0), which you cannot mutate, but not Dropout(0.0). Not very sure that's desirable.

We discussed this during the ML community call. The decision is to keep things simple at the level of Flux/NNlib, i.e. we should return Dropout and not optimize it away. These kinds of optimization could be part of a tool that cuts and stitches the model back together with optimized parts (e.g. in Fluxperimental.jl).

mcabbott · 2023-01-04T00:09:40Z

src/deprecations.jl

+  if active
+    NNlib.dropout(rng, x, p; dims)
+  else
+    Base.depwarn("Flux.dropout(...; active=false) is deprecated. Please branch outside the function, or call dropout(x, 0) if you must.", :dropout)


I made this to preserve the active keyword. Not entirely sure whether use of that was supported outside Flux.

The exported function is this one, not the NNlib one, and I think it lacks a docstring at present.

I added it to FluxML/NNlib.jl#452 because backends need it, but to my knowledge none do for dropout so this seems fine.

You think it might be OK to just go without the deprecation?

Not sure. A quick JuliaHub search didn't turn up anything, but I've enabled downstream tests in case we want to check.

Now 0e396a6 removes this entirely. At least to see if anything changes.

src/layers/normalise.jl

mcabbott · 2023-01-07T19:12:10Z

Downstream:

AtomicGraphNets unrelated
Metalhead loadmodel!(dst::Tuple{Chain{ "but the structures do not match.", unclear where the mismatch is & whether it's related. (Can that have a better error message? IIRC it's recursive, and could tell you what inner layer failed, rather than the whole outermost model.)
GeometricFlux fails as its ∇indexed_softmax says "UndefVarError: within_grad not defined" from here. Maybe Add within_gradient NNlib.jl#434 should have deprecated that or maybe it was obviously internal.

ToucheSir · 2023-01-07T19:28:55Z

Metalhead load failure is known and can be ignored. Adding a better error message is tricky because loadmodel! is general-purpose, but that's not my area of expertise.

For GeometricFlux, I missed that the NNlib PR renamed the function. A courtesy deprecation seems fine here.

src/utils.jl

This reverts commit 0e396a6.

codecov-commenter · 2023-01-08T03:02:56Z

Codecov Report

Base: 86.36% // Head: 83.61% // Decreases project coverage by -2.76% ⚠️

Coverage data is based on head (fc9855b) compared to base (7997174).
Patch coverage: 75.00% of modified lines in pull request are covered.

📣 This organization is not using Codecov’s GitHub App Integration. We recommend you install it so Codecov can continue to function properly for your repositories. Learn more

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #2150      +/-   ##
==========================================
- Coverage   86.36%   83.61%   -2.76%     
==========================================
  Files          19       19              
  Lines        1482     1459      -23     
==========================================
- Hits         1280     1220      -60     
- Misses        202      239      +37

Impacted Files	Coverage Δ
src/deprecations.jl	`46.37% <ø> (ø)`
src/layers/normalise.jl	`86.71% <75.00%> (-2.11%)`	⬇️
src/cuda/cudnn.jl	`0.00% <0.00%> (-90.91%)`	⬇️
src/functor.jl	`44.44% <0.00%> (-42.13%)`	⬇️
src/layers/conv.jl	`87.93% <0.00%> (-0.07%)`	⬇️

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

☔ View full report at Codecov.
📢 Do you have feedback about the report comment? Let us know in this issue.

src/layers/normalise.jl

darsnack · 2023-01-27T17:20:25Z

We discussed this on call and resolved the outstanding default RNG issue (i.e. leave unchanged). I commented on other issues with the decision on call. Other than those discussions, this PR appears ready.

mcabbott · 2023-01-27T17:28:40Z

Ok. In its present state this PR doesn't change the RNG handling.

Looking briefly again, I wonder if we should remove the branch if _isactive(a, x) && a.p != 0 here (which is visible to Zygote) in favour of dropout(a.rng, x, a.p * _isactive(a, x); dims=a.dims). FluxML/NNlib.jl#462 made me more worried about branches.

Edit, some quick times:

julia> using Flux

julia> model = Chain(Dense(10=>100), Dropout(0.5)); data = rand(Float32, 10, 100);

julia> @btime $model($data);
  min 6.667 μs, mean 15.595 μs (4 allocations, 78.22 KiB)  # tagged
  min 6.917 μs, mean 18.080 μs (4 allocations, 78.22 KiB)  # PR yesterday
  min 7.014 μs, mean 15.138 μs (4 allocations, 78.22 KiB)  # without branch

julia> @btime gradient(m -> sum(m($data)), $model);
  min 79.792 μs, mean 115.690 μs (176 allocations, 251.69 KiB) # tagged
  min 49.500 μs, mean 97.422 μs (110 allocations, 209.89 KiB)  # PR yesterday
  min 43.625 μs, mean 73.894 μs (70 allocations, 209.02 KiB)   # without branch

julia> trainmode!(model);

julia> @btime $model($data);
  min 26.500 μs, mean 48.083 μs (17 allocations, 156.67 KiB)
  min 15.042 μs, mean 31.965 μs (6 allocations, 117.33 KiB)
  min 14.875 μs, mean 35.244 μs (6 allocations, 117.33 KiB)

Co-authored-by: Carlo Lucibello <carlo.lucibello@gmail.com>

mcabbott · 2023-01-27T17:53:21Z

One more thought. Should we make test mode visible? It's a bit weird that the model has invisible state. Perhaps we should add a keyword to the constructor, and print this:

julia> Flux.trainmode!(m)
Chain(  
  Dense(2 => 3),                        # 9 parameters
  Dropout(0.4, active=true),
)

With more effort, we could make show check _isactive and write say # train mode, 9 parameters but not sure that's better.

darsnack · 2023-01-27T20:27:17Z

Perhaps we should add a keyword to the constructor

I agree, and a keyword seems the best option to me.

ToucheSir · 2023-01-27T21:31:17Z

We could make the active keyword behave like {train,test}mode!, where users can pass :auto or nothing. That would also help with display.

mcabbott · 2023-01-28T00:05:01Z

Yes. Is active the right name? This is only used internally at present. Maybe it should follow {train,test}mode! and be called train=true for the active state?

ToucheSir · 2023-01-28T19:24:29Z

I recall doing a quick survey of this at some point, but not where I posted the results. PyTorch uses training as inherited from torch.nn.Module. TF relies on a training kwarg at call time which the framework automagically sets. Flax has a deterministic field which it expects users to set themselves and provides no automatic functionality for.

mcabbott · 2023-01-29T14:53:48Z

Thanks for digging. Maybe my vote is to leave this for now, and do it later to BatchNorm too.

src/layers/normalise.jl

darsnack · 2023-01-29T16:33:49Z

src/layers/normalise.jl

  p::F
  dims::D
  active::Union{Bool, Nothing}
  rng::R
 end
-Dropout(p, dims, active) = Dropout(p, dims, active, default_rng_value())
+Dropout(p::Real, dims, active) = Dropout(p, dims, active, default_rng_value())


Not sure if this is intentional but the error checking seems to only apply to the keyword based constructor.

I think that's the only "public" one. I have no idea why we have this 3-arg constructor Dropout(p, dims, active), e.g. Functors will use the 4-arg one. Maybe it was in case someone was relying on it from before the rng field was added?

Yeah that's my recollection.

src/layers/normalise.jl

Co-authored-by: Kyle Daruwalla <daruwalla.k.public@icloud.com>

mcabbott · 2023-01-31T16:31:26Z

Maybe this is ready?

mcabbott commented Jan 3, 2023

View reviewed changes

mcabbott commented Jan 4, 2023

View reviewed changes

mcabbott closed this Jan 5, 2023

mcabbott reopened this Jan 5, 2023

mcabbott added 4 commits January 6, 2023 19:38

use NNlib.dropout, deprecate Flux.dropout

e96e4ef

improve Dropout's docstring

fec2d8e

make Dropout(0) === identity, cannot mutate

6c84a6c

NNlibCUDA = 0.2.5

eab0b15

mcabbott force-pushed the dropout branch from 66f9015 to eab0b15 Compare January 7, 2023 00:39

ToucheSir added the run downstream test label Jan 7, 2023

mcabbott commented Jan 7, 2023

View reviewed changes

src/layers/normalise.jl Show resolved Hide resolved

mcabbott mentioned this pull request Jan 7, 2023

overload NNlib._rng_compat_array FluxML/NNlibCUDA.jl#62

Merged

NNlibCUDA = 0.2.6

4ab93b3

mcabbott force-pushed the dropout branch from d1ef6e8 to 9e61486 Compare January 7, 2023 19:36

simplify default_rng etc

0e396a6

mcabbott force-pushed the dropout branch from 9e61486 to 0e396a6 Compare January 7, 2023 19:38

mcabbott commented Jan 8, 2023

View reviewed changes

src/utils.jl Outdated Show resolved Hide resolved

mcabbott added 2 commits January 7, 2023 20:59

Revert "simplify default_rng etc"

f42f475

This reverts commit 0e396a6.

un-revert the removal of the active=true method

d7cc49d

mcabbott force-pushed the dropout branch from 8dcc5ba to d7cc49d Compare January 8, 2023 02:26

CarloLucibello reviewed Jan 8, 2023

View reviewed changes

src/layers/normalise.jl Outdated Show resolved Hide resolved

mcabbott and others added 2 commits January 27, 2023 12:38

avoid a branch

28ac4c4

Update src/layers/normalise.jl

9e99422

Co-authored-by: Carlo Lucibello <carlo.lucibello@gmail.com>

mcabbott force-pushed the dropout branch from 49327b1 to 9e99422 Compare January 27, 2023 17:49

darsnack reviewed Jan 29, 2023

View reviewed changes

mcabbott commented Jan 29, 2023

View reviewed changes

src/layers/normalise.jl Outdated Show resolved Hide resolved

mcabbott commented Jan 29, 2023

View reviewed changes

src/layers/normalise.jl Outdated Show resolved Hide resolved

Apply suggestions from code review

fc9855b

Co-authored-by: Kyle Daruwalla <daruwalla.k.public@icloud.com>

ToucheSir approved these changes Feb 1, 2023

View reviewed changes

mcabbott merged commit e33de0c into FluxML:master Feb 1, 2023

mcabbott deleted the dropout branch February 1, 2023 01:33

mcabbott mentioned this pull request Mar 31, 2023

Print the state of Dropout etc. #2222

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Move `dropout` to NNlib #2150

Move `dropout` to NNlib #2150

mcabbott commented Jan 3, 2023 •

edited

Loading

mcabbott Jan 3, 2023

ToucheSir Jan 3, 2023 •

edited

Loading

mcabbott Jan 7, 2023

darsnack Jan 27, 2023

mcabbott Jan 4, 2023

ToucheSir Jan 4, 2023

mcabbott Jan 7, 2023

ToucheSir Jan 7, 2023

mcabbott Jan 7, 2023

mcabbott commented Jan 7, 2023 •

edited

Loading

ToucheSir commented Jan 7, 2023

codecov-commenter commented Jan 8, 2023 •

edited

Loading

darsnack commented Jan 27, 2023

mcabbott commented Jan 27, 2023 •

edited

Loading

mcabbott commented Jan 27, 2023

darsnack commented Jan 27, 2023

ToucheSir commented Jan 27, 2023

mcabbott commented Jan 28, 2023

ToucheSir commented Jan 28, 2023

mcabbott commented Jan 29, 2023

darsnack Jan 29, 2023

mcabbott Jan 29, 2023

darsnack Jan 29, 2023

mcabbott commented Jan 31, 2023

Move dropout to NNlib #2150

Move dropout to NNlib #2150

Conversation

mcabbott commented Jan 3, 2023 • edited Loading

Choose a reason for hiding this comment

ToucheSir Jan 3, 2023 • 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

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

mcabbott commented Jan 7, 2023 • edited Loading

ToucheSir commented Jan 7, 2023

codecov-commenter commented Jan 8, 2023 • edited Loading

Codecov Report

darsnack commented Jan 27, 2023

mcabbott commented Jan 27, 2023 • edited Loading

mcabbott commented Jan 27, 2023

darsnack commented Jan 27, 2023

ToucheSir commented Jan 27, 2023

mcabbott commented Jan 28, 2023

ToucheSir commented Jan 28, 2023

mcabbott commented Jan 29, 2023

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

mcabbott commented Jan 31, 2023

Move `dropout` to NNlib #2150

Move `dropout` to NNlib #2150

mcabbott commented Jan 3, 2023 •

edited

Loading

ToucheSir Jan 3, 2023 •

edited

Loading

mcabbott commented Jan 7, 2023 •

edited

Loading

codecov-commenter commented Jan 8, 2023 •

edited

Loading

mcabbott commented Jan 27, 2023 •

edited

Loading