Overload NNlib.within_gradient #136

mcabbott · 2023-01-06T02:11:10Z

Does not solve #133 but the error is different:

julia> using Flux, Zygote, Tracker

julia> bn = BatchNorm(2); x = rand(Float32, 2, 3);

julia> sum(deepcopy(bn)(x))
2.1757803f0

julia> bn.active === nothing
true

julia> bn.active = false
false

julia> _bn = deepcopy(bn); Zygote.withgradient(sum∘_bn, x)[1]
2.1757803f0

julia> _bn = deepcopy(bn); Tracker.withgradient(sum∘_bn, x)[1]
2.1757803f0

julia> bn.active = true
true

julia> _bn = deepcopy(bn); Zygote.withgradient(sum∘_bn, x)[1]
-2.0489097f-7

julia> _bn = deepcopy(bn); Tracker.gradient(sum∘_bn, x)[1]
ERROR: MethodError: no method matching Float32(::Tracker.TrackedReal{Float32})
...
Stacktrace:
  [1] convert(#unused#::Type{Float32}, x::Tracker.TrackedReal{Float32})
    @ Base ./number.jl:7
  [2] setindex!
    @ ./array.jl:971 [inlined]
...
  [9] materialize!
    @ ./broadcast.jl:881 [inlined]
 [10] accum!(x::Matrix{Float32}, Δ::Matrix{Tracker.TrackedReal{Float32}})
    @ Tracker ~/.julia/packages/Tracker/a9oj5/src/back.jl:45
 [11] back(x::Tracker.Tracked{Matrix{Float32}}, Δ::Matrix{Tracker.TrackedReal{Float32}}, once::Bool)
    @ Tracker ~/.julia/packages/Tracker/a9oj5/src/back.jl:48
...
 [22] back(x::Tracker.Tracked{Matrix{Tracker.TrackedReal{Float32}}}, Δ::Matrix{Tracker.TrackedReal{Float32}}, once::Bool)

Does not make the example from FluxML/Flux.jl#2122 work, but perhaps the failure is unrelated:

julia> let
       using Flux
       using Random
       Random.seed!(123)

       model = Chain(
                 Conv((3, 3), 3 => 5, pad=1, bias=false), 
                 BatchNorm(5, relu), 
                 Conv((3, 3), 5 => 3, stride=16),
               )
       image = rand(Float32, 224, 224, 3, 1);
       @show sum(model(image));

       loss(m, x) = sum(m(x))

       opt = Flux.setup(Flux.Adam(0.001f0,  (0.9f0, 0.999f0), 1.1920929f-7), model)

       val, grads = Tracker.withgradient(model) do m
           loss(m, image)
       end

       Flux.update!(opt, model, grads[1])
       @show loss(model, image);
       end;
sum(model(image)) = -0.33076355f0
ERROR: UndefRefError: access to undefined reference
Stacktrace:
  [1] getindex
    @ ./essentials.jl:14 [inlined]
  [2] conv_direct!(y::Array{Tracker.TrackedReal{Float32}, 5}, x::Array{Tracker.TrackedReal{Float32}, 5}, w::Array{Float32, 5}, cdims::DenseConvDims{3, 3, 3, 6, 3}, ::Val{(3, 3, 1)}, ::Val{3}, ::Val{(0, 0, 0, 0, 0, 0)}, ::Val{(1, 1, 1)}, ::Val{(16, 16, 1)}, fk::Val{false}; alpha::Tracker.TrackedReal{Float32}, beta::Bool)
    @ NNlib ~/.julia/packages/NNlib/QJIIj/src/impl/conv_direct.jl:111
  [3] kwcall(::NamedTuple{(:alpha, :beta), Tuple{Tracker.TrackedReal{Float32}, Bool}}, ::typeof(NNlib.conv_direct!), y::Array{Tracker.TrackedReal{Float32}, 5}, x::Array{Tracker.TrackedReal{Float32}, 5}, w::Array{Float32, 5}, cdims::DenseConvDims{3, 3, 3, 6, 3}, ::Val{(3, 3, 1)}, ::Val{3}, ::Val{(0, 0, 0, 0, 0, 0)}, ::Val{(1, 1, 1)}, ::Val{(16, 16, 1)}, fk::Val{false})
    @ NNlib ~/.julia/packages/NNlib/QJIIj/src/impl/conv_direct.jl:59
  [4] conv_direct!(y::Array{Tracker.TrackedReal{Float32}, 5}, x::Array{Tracker.TrackedReal{Float32}, 5}, w::Array{Float32, 5}, cdims::DenseConvDims{3, 3, 3, 6, 3}; alpha::Tracker.TrackedReal{Float32}, beta::Bool)
    @ NNlib ~/.julia/packages/NNlib/QJIIj/src/impl/conv_direct.jl:50
  [5] conv_direct!
    @ ~/.julia/packages/NNlib/QJIIj/src/impl/conv_direct.jl:47 [inlined]
  [6] #conv!#301
    @ ~/.julia/packages/NNlib/QJIIj/src/conv.jl:288 [inlined]
  [7] conv!
    @ ~/.julia/packages/NNlib/QJIIj/src/conv.jl:280 [inlined]
  [8] conv!(y::Array{Tracker.TrackedReal{Float32}, 4}, x::Array{Tracker.TrackedReal{Float32}, 4}, w::Array{Float32, 4}, cdims::DenseConvDims{2, 2, 2, 4, 2}; kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ NNlib ~/.julia/packages/NNlib/QJIIj/src/conv.jl:145
  [9] conv!
    @ ~/.julia/packages/NNlib/QJIIj/src/conv.jl:140 [inlined]
 [10] conv(x::Array{Tracker.TrackedReal{Float32}, 4}, w::Array{Float32, 4}, cdims::DenseConvDims{2, 2, 2, 4, 2}; kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ NNlib ~/.julia/packages/NNlib/QJIIj/src/conv.jl:88
 [11] conv
    @ ~/.julia/packages/NNlib/QJIIj/src/conv.jl:83 [inlined]
 [12] #_forward#636
    @ ~/.julia/packages/Tracker/a9oj5/src/lib/array.jl:520 [inlined]
 [13] _forward
    @ ./none:0 [inlined]
 [14] #track#1
    @ ~/.julia/packages/Tracker/a9oj5/src/Tracker.jl:58 [inlined]
 [15] track
    @ ~/.julia/packages/Tracker/a9oj5/src/Tracker.jl:57 [inlined]
 [16] #conv#633
    @ ~/.julia/packages/Tracker/a9oj5/src/lib/array.jl:516 [inlined]
 [17] conv
    @ ~/.julia/packages/Tracker/a9oj5/src/lib/array.jl:516 [inlined]
 [18] (::Conv{2, 4, typeof(identity), TrackedArray{…,Array{Float32, 4}}, TrackedArray{…,Vector{Float32}}})(x::TrackedArray{…,Array{Tracker.TrackedReal{Float32}, 4}})
    @ Flux ~/.julia/dev/Flux/src/layers/conv.jl:200
 [19] macro expansion
    @ ~/.julia/dev/Flux/src/layers/basic.jl:53 [inlined]
 [20] _applychain(layers::Tuple{Conv{2, 4, typeof(identity), TrackedArray{…,Array{Float32, 4}}, Bool}, BatchNorm{typeof(relu), TrackedArray{…,Vector{Float32}}, Float32, Vector{Float32}}, Conv{2, 4, typeof(identity), TrackedArray{…,Array{Float32, 4}}, TrackedArray{…,Vector{Float32}}}}, x::Array{Float32, 4})
    @ Flux ~/.julia/dev/Flux/src/layers/basic.jl:53
 [21] Chain
    @ ~/.julia/dev/Flux/src/layers/basic.jl:51 [inlined]
 [22] (::var"#loss#31")(m::Chain{Tuple{Conv{2, 4, typeof(identity), TrackedArray{…,Array{Float32, 4}}, Bool}, BatchNorm{typeof(relu), TrackedArray{…,Vector{Float32}}, Float32, Vector{Float32}}, Conv{2, 4, typeof(identity), TrackedArray{…,Array{Float32, 4}}, TrackedArray{…,Vector{Float32}}}}}, x::Array{Float32, 4})
    @ Main ./REPL[28]:14
 [23] (::var"#30#32"{var"#loss#31", Array{Float32, 4}})(m::Chain{Tuple{Conv{2, 4, typeof(identity), TrackedArray{…,Array{Float32, 4}}, Bool}, BatchNorm{typeof(relu), TrackedArray{…,Vector{Float32}}, Float32, Vector{Float32}}, Conv{2, 4, typeof(identity), TrackedArray{…,Array{Float32, 4}}, TrackedArray{…,Vector{Float32}}}}})
    @ Main ./REPL[28]:19
 [24] withgradient(f::Function, xs::Chain{Tuple{Conv{2, 4, typeof(identity), Array{Float32, 4}, Bool}, BatchNorm{typeof(relu), Vector{Float32}, Float32, Vector{Float32}}, Conv{2, 4, typeof(identity), Array{Float32, 4}, Vector{Float32}}}})
    @ Tracker ~/.julia/packages/Tracker/a9oj5/src/back.jl:218
 [25] top-level scope
    @ REPL[28]:18

Edit: fixed, closes #133, closes #137

coveralls · 2023-01-06T02:37:02Z

Pull Request Test Coverage Report for Build 3852665337

13 of 16 (81.25%) changed or added relevant lines in 2 files are covered.
1 unchanged line in 1 file lost coverage.
Overall coverage increased (+0.09%) to 72.439%

Changes Missing Coverage	Covered Lines	Changed/Added Lines	%
src/lib/real.jl	0	1	0.0%
src/lib/array.jl	13	15	86.67%

Files with Coverage Reduction	New Missed Lines	%
src/lib/array.jl	1	67.49%

Totals
Change from base Build 3010278444:	0.09%
Covered Lines:	502
Relevant Lines:	693

💛 - Coveralls

ToucheSir · 2023-01-06T03:25:10Z

For the first error, guessing it needs something like https://github.com/FluxML/Flux.jl/blob/master/src/layers/normalise.jl#L278-L280.

For the second, this is FluxML/NNlib.jl#405. Although we probably ought to make the NNlib side more robust, the conv layer and conv functions shouldn't be seeing Array{<:TrackedReal} in the first place!

mcabbott · 2023-01-06T03:34:49Z

For the first, I think that after f098dd9 the error is not from within _track_stats!.

For the second, indeed the error is exactly that, thanks. But why an array of tracked numbers is made I don't know. A smaller example is this:

julia> Flux.trainmode!(BatchNorm(3))(param(rand(2,3,1)))
Tracked 2×3×1 Array{Tracker.TrackedReal{Float64}, 3}:
[:, :, 1] =
 -0.994696  -0.999959   0.912711
  0.994696   0.999959  -0.912711

julia> Flux.testmode!(BatchNorm(3))(param(rand(2,3,1)))
Tracked 2×3×1 Array{Float64, 3}:
[:, :, 1] =
 0.664405  0.684316  0.321995
 0.834291  0.463685  0.370812

... which is #137.

ToucheSir

Only remaining test failures on nightly are the known kron ones.

coveralls · 2024-11-05T13:41:15Z

Pull Request Test Coverage Report for Build 3852093460

Warning: This coverage report may be inaccurate.

This pull request's base commit is no longer the HEAD commit of its target branch. This means it includes changes from outside the original pull request, including, potentially, unrelated coverage changes.

For more information on this, see Tracking coverage changes with pull request builds.
To avoid this issue with future PRs, see these Recommended CI Configurations.
For a quick fix, rebase this PR at GitHub. Your next report should be accurate.

Details

12 of 15 (80.0%) changed or added relevant lines in 2 files are covered.
38 unchanged lines in 1 file lost coverage.
Overall coverage decreased (-0.03%) to 72.328%

Changes Missing Coverage	Covered Lines	Changed/Added Lines	%
src/lib/real.jl	0	1	0.0%
src/lib/array.jl	12	14	85.71%

Files with Coverage Reduction	New Missed Lines	%
src/lib/array.jl	38	67.14%

Totals
Change from base Build 3010278444:	-0.03%
Covered Lines:	494
Relevant Lines:	683

💛 - Coveralls

coveralls · 2024-11-05T13:41:15Z

Pull Request Test Coverage Report for Build 3852630410

Warning: This coverage report may be inaccurate.

This pull request's base commit is no longer the HEAD commit of its target branch. This means it includes changes from outside the original pull request, including, potentially, unrelated coverage changes.

For more information on this, see Tracking coverage changes with pull request builds.
To avoid this issue with future PRs, see these Recommended CI Configurations.
For a quick fix, rebase this PR at GitHub. Your next report should be accurate.

Details

12 of 14 (85.71%) changed or added relevant lines in 2 files are covered.
30 unchanged lines in 1 file lost coverage.
Overall coverage increased (+0.2%) to 72.518%

Changes Missing Coverage	Covered Lines	Changed/Added Lines	%
src/lib/array.jl	12	13	92.31%
src/lib/real.jl	0	1	0.0%

Files with Coverage Reduction	New Missed Lines	%
src/lib/array.jl	30	67.67%

Totals
Change from base Build 3010278444:	0.2%
Covered Lines:	504
Relevant Lines:	695

💛 - Coveralls

mcabbott added 3 commits January 5, 2023 21:09

use NNlib.within_gradient

0cbd576

and one more thing to help Flux out

f098dd9

forgot to save

988136d

fix FluxML#137

53b500f

mcabbott force-pushed the within_grad branch from 06f712d to 53b500f Compare January 6, 2023 04:36

ToucheSir approved these changes Jan 6, 2023

View reviewed changes

mcabbott merged commit 00393e2 into FluxML:master Jan 6, 2023

mcabbott deleted the within_grad branch January 6, 2023 05:40

mcabbott mentioned this pull request Jan 6, 2023

Use NNlib's within_gradient FluxML/Flux.jl#2152

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Overload NNlib.within_gradient #136

Overload NNlib.within_gradient #136

mcabbott commented Jan 6, 2023 •

edited

Loading

coveralls commented Jan 6, 2023 •

edited

Loading

ToucheSir commented Jan 6, 2023

mcabbott commented Jan 6, 2023 •

edited

Loading

ToucheSir left a comment

coveralls commented Nov 5, 2024

coveralls commented Nov 5, 2024 •

edited

Loading

Overload NNlib.within_gradient #136

Overload NNlib.within_gradient #136

Conversation

mcabbott commented Jan 6, 2023 • edited Loading

coveralls commented Jan 6, 2023 • edited Loading

Pull Request Test Coverage Report for Build 3852665337

💛 - Coveralls

ToucheSir commented Jan 6, 2023

mcabbott commented Jan 6, 2023 • edited Loading

ToucheSir left a comment

Choose a reason for hiding this comment

coveralls commented Nov 5, 2024

Pull Request Test Coverage Report for Build 3852093460

Warning: This coverage report may be inaccurate.

Details

💛 - Coveralls

coveralls commented Nov 5, 2024 • edited Loading

Pull Request Test Coverage Report for Build 3852630410

Warning: This coverage report may be inaccurate.

Details

💛 - Coveralls

mcabbott commented Jan 6, 2023 •

edited

Loading

coveralls commented Jan 6, 2023 •

edited

Loading

mcabbott commented Jan 6, 2023 •

edited

Loading

coveralls commented Nov 5, 2024 •

edited

Loading