From 8ca6ce0c0b8c43099dff771d94d013d493d19d32 Mon Sep 17 00:00:00 2001 From: Carlo Lucibello Date: Tue, 9 Apr 2024 04:11:32 +0200 Subject: [PATCH] add `path` option to `trainables` (#174) * add path=true * fix * fix * fix docs * fix docs * update doc workflow --- .github/workflows/ci.yml | 24 ---- .github/workflows/dependabot.yml | 7 + .github/workflows/documentation.yml | 28 ++++ Project.toml | 2 +- docs/.DS_Store | Bin 6148 -> 6148 bytes docs/Project.toml | 1 + docs/make.jl | 6 +- docs/src/api.md | 14 ++ src/Optimisers.jl | 7 +- src/destructure.jl | 2 +- src/interface.jl | 16 +-- src/trainables.jl | 81 +++++++++-- src/utils.jl | 15 +++ test/trainables.jl | 202 ++++++++++++++++------------ 14 files changed, 268 insertions(+), 137 deletions(-) create mode 100644 .github/workflows/dependabot.yml create mode 100644 .github/workflows/documentation.yml create mode 100644 src/utils.jl diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 5c5a952a..7c3b18cc 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -53,27 +53,3 @@ jobs: file: lcov.info continue-on-error: ${{ matrix.julia-version == 'nightly' }} - docs: - name: Documentation - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v3 - - uses: julia-actions/setup-julia@v1 - with: - version: '1.6' - - run: | - julia --project=docs -e ' - using Pkg - Pkg.develop(PackageSpec(path=pwd())) - Pkg.instantiate()' - - run: | - julia --color=yes --project=docs/ -e ' - using Optimisers - using Documenter - using Documenter: doctest - DocMeta.setdocmeta!(Optimisers, :DocTestSetup, :(using Optimisers); recursive = true) - doctest(Optimisers)' - - run: julia --project=docs docs/make.jl - env: - GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} - DOCUMENTER_KEY: ${{ secrets.DOCUMENTER_KEY }} diff --git a/.github/workflows/dependabot.yml b/.github/workflows/dependabot.yml new file mode 100644 index 00000000..700707ce --- /dev/null +++ b/.github/workflows/dependabot.yml @@ -0,0 +1,7 @@ +# https://docs.github.com/github/administering-a-repository/configuration-options-for-dependency-updates +version: 2 +updates: + - package-ecosystem: "github-actions" + directory: "/" # Location of package manifests + schedule: + interval: "weekly" diff --git a/.github/workflows/documentation.yml b/.github/workflows/documentation.yml new file mode 100644 index 00000000..efb91902 --- /dev/null +++ b/.github/workflows/documentation.yml @@ -0,0 +1,28 @@ +name: Documentation + +on: + push: + branches: + - master # update to match your development branch (master, main, dev, trunk, ...) + tags: '*' + pull_request: + +jobs: + build: + permissions: + contents: write + statuses: write + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: julia-actions/setup-julia@v1 + with: + version: '1.10' + - uses: julia-actions/cache@v1 + - name: Install dependencies + run: julia --project=docs/ -e 'using Pkg; Pkg.develop(PackageSpec(path=pwd())); Pkg.instantiate()' + - name: Build and deploy + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} # If authenticating with GitHub Actions token + DOCUMENTER_KEY: ${{ secrets.DOCUMENTER_KEY }} # If authenticating with SSH deploy key + run: julia --project=docs/ docs/make.jl diff --git a/Project.toml b/Project.toml index 7422f28c..4920eaa9 100644 --- a/Project.toml +++ b/Project.toml @@ -12,7 +12,7 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" [compat] ChainRulesCore = "1" -Functors = "0.4" +Functors = "0.4.9" Statistics = "1" Zygote = "0.6.40" julia = "1.6" diff --git a/docs/.DS_Store b/docs/.DS_Store index e14553107e0e9f4dfeb62498c2852eba47071535..3836d25ddc99a8f8e54c93754f83b9777eda452b 100644 GIT binary patch delta 92 zcmZoMXfc@J&&a(oU^gQp_hcTX`>!Cw%}gXpWX+P rrI}@5Oj~9-#)io;%*IUJk&`Dfi?QikHubHzzj-aQGs|Xnj=%f>@RT8G delta 194 zcmZoMXfc@J&&ahgU^gQp*JK{1`}H~LhQZ1CxdjX$z`O6P{fdoY~##z Wi&(|)Z9c;6%rdb-d^0=8Uw#1h-Z6s! diff --git a/docs/Project.toml b/docs/Project.toml index 47d2b7c1..31e718a6 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -1,5 +1,6 @@ [deps] Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196" +Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" diff --git a/docs/make.jl b/docs/make.jl index 47d64137..55594e8c 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -1,8 +1,9 @@ using Documenter, Optimisers, Zygote, StaticArrays, Functors -DocMeta.setdocmeta!(Optimisers, :DocTestSetup, :(using Optimisers); recursive = true) +DocMeta.setdocmeta!(Optimisers, :DocTestSetup, :(using Optimisers, Functors); recursive = true) +DocMeta.setdocmeta!(Functors, :DocTestSetup, :(using Functors); recursive = true) -makedocs(modules = [Optimisers], +makedocs(modules = [Optimisers, Functors], doctest = false, sitename = "Optimisers.jl", pages = ["Home" => "index.md", @@ -13,6 +14,7 @@ makedocs(modules = [Optimisers], assets = ["assets/flux.css"], prettyurls = get(ENV, "CI", nothing) == "true" ), + checkdocs = :none, # don't check that Functors' docstrings are all reported here ) deploydocs( diff --git a/docs/src/api.md b/docs/src/api.md index 5c203492..378bf72a 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -1,3 +1,6 @@ +```@meta +CollapsedDocStrings = true +``` ## Optimisation Rules @@ -72,3 +75,14 @@ Optimisers.@lazy Optimisers.adjust(::AbstractRule, ::Real) Optimisers.@def ``` + +## KeyPath + +A `KeyPath` is a sequence of keys that can be used to access a value within a nested structure. +It is defined in Functors.jl and re-exported by Optimisers.jl here for convenience. + +```@docs +Functors.KeyPath +Functors.haskeypath +Functors.getkeypath +``` diff --git a/src/Optimisers.jl b/src/Optimisers.jl index 3cc98808..2e115c40 100644 --- a/src/Optimisers.jl +++ b/src/Optimisers.jl @@ -1,11 +1,15 @@ module Optimisers -using Functors: functor, fmap, isleaf, @functor, fmapstructure, children, AbstractWalk +using Functors: functor, fmap, fmap_with_path, + KeyPath, haskeypath, getkeypath, + isleaf, @functor, fmapstructure, children, AbstractWalk using LinearAlgebra include("interface.jl") export AbstractRule +include("utils.jl") + include("adjust.jl") include("destructure.jl") @@ -13,6 +17,7 @@ export destructure include("trainables.jl") export trainables +export KeyPath, haskeypath, getkeypath # from Functors.jl include("rules.jl") export Descent, Adam, Momentum, Nesterov, Rprop, RMSProp, diff --git a/src/destructure.jl b/src/destructure.jl index f9950a92..a6284522 100644 --- a/src/destructure.jl +++ b/src/destructure.jl @@ -78,7 +78,7 @@ end struct TrainableStructWalk <: AbstractWalk end -(::TrainableStructWalk)(recurse, x) = map(recurse, _trainable(x)) +(::TrainableStructWalk)(recurse, x) = mapvalue(recurse, _trainable(x)) _vec(x::Number) = LinRange(x,x,1) _vec(x::AbstractArray) = vec(x) diff --git a/src/interface.jl b/src/interface.jl index aa5447c0..ac9b90bc 100644 --- a/src/interface.jl +++ b/src/interface.jl @@ -45,7 +45,7 @@ function _setup(rule, x; cache) cache[x] = ℓ end else - valuemap(xᵢ -> _setup(rule, xᵢ; cache), _trainable(x)) + mapvalue(xᵢ -> _setup(rule, xᵢ; cache), _trainable(x)) end end @@ -82,7 +82,7 @@ function _update!(tree, x; grads, params) haskey(params, (tree,x)) && return params[(tree,x)] isbits(tree) && return x # means () is not cached, and also (((),),) x′, re = functor(x) - x′′ = re(valuemap((tᵢ, xᵢ) -> _update!(tᵢ, xᵢ; grads, params), tree, x′)) + x′′ = re(mapvalue((tᵢ, xᵢ) -> _update!(tᵢ, xᵢ; grads, params), tree, x′)) if ismutable(x′′) params[(tree,x)] = x′′ else # no ties to preserve between immutable structs, right? @@ -115,7 +115,7 @@ function _grads!(dict::IdDict, tree, x, x̄s...) # functor(typeof(tree), base(x̄)), for things like Transpose x̄s′ = map(x̄ -> functor(typeof(x), base(x̄))[1], x̄s) x′, _ = functor(typeof(x), x) - valueforeach((tᵢ, xᵢ, x̄sᵢ...) -> _grads!(dict, tᵢ, xᵢ, x̄sᵢ...), tree, x′, x̄s′...) + foreachvalue((tᵢ, xᵢ, x̄sᵢ...) -> _grads!(dict, tᵢ, xᵢ, x̄sᵢ...), tree, x′, x̄s′...) end # default all rules to first order calls @@ -172,7 +172,7 @@ _trainable(x) = _trainable(functor(x)[1], trainable(x)) _trainable(ch::NamedTuple, tr::NamedTuple) = merge(map(_ -> nothing, ch), tr) _trainable(ch::Tuple{Vararg{Any,N}}, tr::Tuple{Vararg{Any,N}}) where N = tr _trainable(ch::AbstractArray, tr::AbstractArray) = tr -_trainable(ch::Dict, tr::Dict) = merge(valuemap(_ -> nothing, ch), tr) +_trainable(ch::Dict, tr::Dict) = merge(mapvalue(_ -> nothing, ch), tr) function _trainable(ch::NamedTuple, tr::Tuple) # for old Flux-style no-names tuple @warn "trainable(x) should now return a NamedTuple with the field names, not a Tuple" maxlog=3 @@ -180,14 +180,6 @@ function _trainable(ch::NamedTuple, tr::Tuple) # for old Flux-style no-names tu end -valuemap(f, x...) = map(f, x...) -valuemap(f, x::Dict, ys...) = Dict(k => f(v, (get(y, k, nothing) for y in ys)...) for (k,v) in x) -valueforeach(f, x...) = foreach(f, x...) -valueforeach(f, x::Dict, ys...) = foreach(pairs(x)) do (k, v) - f(v, (get(y, k, nothing) for y in ys)...) -end - - ### ### rule definition helpers ### diff --git a/src/trainables.jl b/src/trainables.jl index 625c5659..370de0f6 100644 --- a/src/trainables.jl +++ b/src/trainables.jl @@ -1,12 +1,17 @@ """ - trainables(x) + trainables(x, path = false) -Return a list over all the trainable parameters in `x`, that is all the numerical +Return an iterable over all the trainable parameters in `x`, that is all the numerical arrays (see [`isnumeric`](@ref Optimisers.isnumeric)) which are reachable through [`trainable`](@ref Optimisers.trainable). Parameters appearing multiple times in the model (tied weights) will be present only once in the output. +If `path = false`, the output is a list of numerical arrays. + +If `path = true`, the output is a list of `(KeyPath, AbstractArray)` pairs, where [`KeyPath`](@ref) is a type +representing the path to the array in the original structure. + See also [`destructure`](@ref) for a similar operation that returns a single flat vector instead. # Examples @@ -33,11 +38,36 @@ julia> trainables(x) 2-element Vector{AbstractArray}: [1.0, 2.0] [3.0] +``` + +```jldoctest +julia> x = (a = [1.0,2.0], b = (Dict("c" => [3.0, 4.0], "d" => 5.0), [6.0,7.0])); + +julia> for (kp, y) in trainables(x, path = true) + println(kp, " => ", y) + end +KeyPath(:a,) => [1.0, 2.0] +KeyPath(:b, 1, "c") => [3.0, 4.0] +KeyPath(:b, 2) => [6.0, 7.0] + +julia> getkeypath(x, KeyPath(:b, 1, "c")) +2-element Vector{Float64}: + 3.0 + 4.0 +``` """ -function trainables(x) +function trainables(x; path = false) + if path + return _trainables_with_path(x) + else + return _trainables(x) + end +end + + +function _trainables(x) arrays = AbstractArray[] - exclude(x) = Optimisers.isnumeric(x) - fmap(x; exclude, walk = Optimisers.TrainableStructWalk()) do y + fmap(x; exclude = isnumeric, walk = TrainableStructWalk()) do y push!(arrays, y) return y end @@ -45,15 +75,50 @@ function trainables(x) end function ∇trainables(x, Δ) - exclude(x) = Optimisers.isnumeric(x) i = 0 - return fmapstructure(x; exclude, walk = TrainableStructWalk()) do _ + return fmapstructure(x; exclude = isnumeric, walk = TrainableStructWalk()) do _ return Δ[i+=1] end end -function ChainRulesCore.rrule(::typeof(trainables), x) +function ChainRulesCore.rrule(::typeof(_trainables), x) y = trainables(x) trainables_back(Δ) = (NoTangent(), ∇trainables(x, unthunk(Δ))) return y, trainables_back end + +function _trainables_with_path(x) + named_params = [] + exclude(kp, x) = isnumeric(x) + fmap_with_path(x; exclude, walk = TrainableStructWalkWithPath()) do kp, y + push!(named_params, (kp, y)) + return y + end + return named_params +end + +struct TrainableStructWalkWithPath <: AbstractWalk end + +function (::TrainableStructWalkWithPath)(recurse, kp::KeyPath, x) + x_children = trainable(x) + kps = mapkey(c -> KeyPath(kp, c), x_children) + return mapvalue(recurse, kps, x_children) +end + +function ChainRulesCore.rrule(::typeof(_trainables_with_path), x) + y = _trainables_with_path(x) + trainables_with_path_back(Δ) = (NoTangent(), ∇trainables_with_path(x, unthunk(Δ))) + return y, trainables_with_path_back +end + +function ∇trainables_with_path(x, Δ) + i = 0 + return fmapstructure(x; exclude = isnumeric, walk = TrainableStructWalk()) do _ + Δi = Δ[i+=1] + if isnothing(Δi) + return nothing + else + return Δi[2] + end + end +end diff --git a/src/utils.jl b/src/utils.jl new file mode 100644 index 00000000..7c6c95be --- /dev/null +++ b/src/utils.jl @@ -0,0 +1,15 @@ + +mapvalue(f, x...) = map(f, x...) +mapvalue(f, x::Dict, ys...) = Dict(k => f(v, (get(y, k, nothing) for y in ys)...) for (k,v) in x) + +mapkey(f, x::NamedTuple{Ks}) where Ks = NamedTuple{Ks}(map(f, Ks)) +mapkey(f, x::Dict) = Dict(k => f(k) for k in keys(x)) +mapkey(f, x::Tuple) = ntuple(i -> f(i), length(x)) +mapkey(f, x::AbstractArray) = [f(i) for i=1:length(x)] + +foreachvalue(f, x...) = foreach(f, x...) + +foreachvalue(f, x::Dict, ys...) = foreach(pairs(x)) do (k, v) + f(v, (get(y, k, nothing) for y in ys)...) +end + diff --git a/test/trainables.jl b/test/trainables.jl index d4b93ce8..e1aa0115 100644 --- a/test/trainables.jl +++ b/test/trainables.jl @@ -14,102 +14,128 @@ mat = Float32[4 6; 5 7] m9 = (a = m1, b = mat, c = [mat, m1]) @testset "trainables" begin - ps = trainables(m1) - @test ps isa Vector - @test length(ps) == 1 - @test ps[1] == m1 - - ps = trainables(m2) - @test ps isa Vector - @test length(ps) == 2 - @test ps[1] == m2[1] - @test ps[2] == m2[2] - - ps = trainables(m3) - @test length(ps) == 2 - @test ps[1] == 1:3 - @test ps[2] == 4:6 - - ps = trainables(m4) - @test length(ps) == 2 - @test ps[1] == 1:3 - @test ps[2] == 4:6 - - ps = trainables(m5) - @test length(ps) == 3 - @test ps[1] == 1:3 - @test ps[2] == 4:6 - @test ps[3] == 4:6 - - ps = trainables(m6) - @test length(ps) == 2 - @test ps[1] == 1:3 - @test ps[2] == ComplexF64[4.0 + 1.0im] - - ps = trainables(m7) - @test length(ps) == 1 - @test ps[1] == [1.0, 2.0, 3.0] - - ps = trainables(m8) - @test length(ps) == 3 - @test ps[1] == 1:3 - @test ps[2] == [4.0] - @test ps[3] == [5.0] - - ps = trainables(m9) - @test length(ps) == 2 - @test ps[1] == 1:3 - @test ps[2] == mat + ps = trainables(m1) + @test ps isa Vector + @test length(ps) == 1 + @test ps[1] == m1 + + ps = trainables(m2) + @test ps isa Vector + @test length(ps) == 2 + @test ps[1] == m2[1] + @test ps[2] == m2[2] + + ps = trainables(m3) + @test length(ps) == 2 + @test ps[1] == 1:3 + @test ps[2] == 4:6 + + ps = trainables(m4) + @test length(ps) == 2 + @test ps[1] == 1:3 + @test ps[2] == 4:6 + + ps = trainables(m5) + @test length(ps) == 3 + @test ps[1] == 1:3 + @test ps[2] == 4:6 + @test ps[3] == 4:6 + + ps = trainables(m6) + @test length(ps) == 2 + @test ps[1] == 1:3 + @test ps[2] == ComplexF64[4.0 + 1.0im] + + ps = trainables(m7) + @test length(ps) == 1 + @test ps[1] == [1.0, 2.0, 3.0] + + ps = trainables(m8) + @test length(ps) == 3 + @test ps[1] == 1:3 + @test ps[2] == [4.0] + @test ps[3] == [5.0] + + ps = trainables(m9) + @test length(ps) == 2 + @test ps[1] == 1:3 + @test ps[2] == mat end @testset "gradient" begin - loss(m) = sum([sum(abs2, p) for p in trainables(m)]) - g = gradient(loss, m1)[1] - @test g == [2.0, 4.0, 6.0] - - g = gradient(loss, m2)[1] - @test g == ([2.0, 4.0, 6.0], [8.0, 10.0, 12.0]) - - g = gradient(loss, m3)[1] - @test g.x == [2.0, 4.0, 6.0] - @test g.y === nothing - @test g.z == [8.0, 10.0, 12.0] - - g = gradient(loss, m4)[1] - @test g == (x = [2.0, 4.0, 6.0], y = [2.0, 4.0, 6.0], z = [8.0, 10.0, 12.0]) - g.x === g.y # shared gradient for shared weights - - g = gradient(loss, m5)[1] - @test g == (a = ((x = [2.0, 4.0, 6.0], y = nothing, z = [8.0, 10.0, 12.0]), nothing), b = ([2.0, 4.0, 6.0], nothing), c = ((x = [2.0, 4.0, 6.0], y = [2.0, 4.0, 6.0], z = [8.0, 10.0, 12.0]), nothing)) - - g = gradient(loss, m6)[1] - @test g == (a = [2.0, 4.0, 6.0], b = ComplexF64[8.0 + 2.0im], c = [2.0, 4.0, 6.0]) - - g = gradient(loss, m7)[1] - @test g == (a = (nothing, [2.0, 4.0, 6.0]), b = nothing, c = nothing) - - g = gradient(loss, m8)[1] - @test g[1] == (x = [2.0, 4.0, 6.0], y = [2.0, 4.0, 6.0]) - @test g[2] == (a = nothing, b = (x = [8.0], y = nothing), c = nothing) - @test g[3] == [[10.0]] - - g = gradient(loss, m9)[1] - @test g == (a = [2.0, 4.0, 6.0], b = Float32[8.0 12.0; 10.0 14.0], c = Array[Float32[8.0 12.0; 10.0 14.0], [2.0, 4.0, 6.0]]) + loss(m) = sum([sum(abs2, p) for p in trainables(m)]) + g = gradient(loss, m1)[1] + @test g == [2.0, 4.0, 6.0] + + g = gradient(loss, m2)[1] + @test g == ([2.0, 4.0, 6.0], [8.0, 10.0, 12.0]) + + g = gradient(loss, m3)[1] + @test g.x == [2.0, 4.0, 6.0] + @test g.y === nothing + @test g.z == [8.0, 10.0, 12.0] + + g = gradient(loss, m4)[1] + @test g == (x = [2.0, 4.0, 6.0], y = [2.0, 4.0, 6.0], z = [8.0, 10.0, 12.0]) + g.x === g.y # shared gradient for shared weights + + g = gradient(loss, m5)[1] + @test g == (a = ((x = [2.0, 4.0, 6.0], y = nothing, z = [8.0, 10.0, 12.0]), nothing), b = ([2.0, 4.0, 6.0], nothing), c = ((x = [2.0, 4.0, 6.0], y = [2.0, 4.0, 6.0], z = [8.0, 10.0, 12.0]), nothing)) + + g = gradient(loss, m6)[1] + @test g == (a = [2.0, 4.0, 6.0], b = ComplexF64[8.0 + 2.0im], c = [2.0, 4.0, 6.0]) + + g = gradient(loss, m7)[1] + @test g == (a = (nothing, [2.0, 4.0, 6.0]), b = nothing, c = nothing) + + g = gradient(loss, m8)[1] + @test g[1] == (x = [2.0, 4.0, 6.0], y = [2.0, 4.0, 6.0]) + @test g[2] == (a = nothing, b = (x = [8.0], y = nothing), c = nothing) + @test g[3] == [[10.0]] + + g = gradient(loss, m9)[1] + @test g == (a = [2.0, 4.0, 6.0], b = Float32[8.0 12.0; 10.0 14.0], c = Array[Float32[8.0 12.0; 10.0 14.0], [2.0, 4.0, 6.0]]) +end + +@testset "dict" begin + d = Dict(:a => rand(2), :b => ones(2)) + ps = trainables(d) + @test length(ps) == 2 + @test ps[1] == d[:a] + @test ps[2] == d[:b] + + g = gradient(d -> sum(trainables(d)[1].^2) /2 + sum(trainables(d)[2]), d)[1] + @test g[:a] == d[:a] + @test_broken g[:b] == [1.0, 1.0] end @testset "second order derivatives" begin - struct DenseLayer - w - b - end + struct DenseLayer + w + b + end + + Functors.@functor DenseLayer - Functors.@functor DenseLayer + loss(m) = sum([sum(abs2, p) for p in trainables(m)]) + + model = DenseLayer([1. 2.; 3. 4.], [0., 0.]) + + g = gradient(m -> loss(gradient(loss, m)), model)[1] + @test g.w == [8.0 16.0; 24.0 32.0] + @test g.b == [0.0, 0.0] +end - loss(m) = sum([sum(abs2, p) for p in trainables(m)]) +@testset "trainables(x, path=true)" begin + loss(m) = sum(abs2, trainables(m, path=true)[1][2]) - model = DenseLayer([1. 2.; 3. 4.], [0., 0.]) + ps = trainables(m4, path=true) + @test length(ps) == 2 + @test ps[1] == (KeyPath(:x,), [1.0, 2.0, 3.0]) + @test ps[2] == (KeyPath(:z,), [4.0, 5.0, 6.0]) - g = gradient(m -> loss(gradient(loss, m)), model)[1] - @test g.w == [8.0 16.0; 24.0 32.0] - @test g.b == [0.0, 0.0] + g = gradient(loss, m4)[1] + @test g.x == [2.0, 4.0, 6.0] + @test g.y == [2.0, 4.0, 6.0] + @test g.z === nothing end