add modules function

Author: CarloLucibello

Manifest.toml

@@ -383,4 +383,4 @@ version = "0.6.3"
 deps = ["MacroTools"]
 git-tree-sha1 = "9e7a1e8ca60b742e508a315c17eef5211e7fbfd7"
 uuid = "700de1a5-db45-46bc-99cf-38207098b444"
-version = "0.2.1"
+version = "0.2.1"

docs/src/utilities.md

@@ -91,6 +91,7 @@ Flux.outputsize
 ## Model Abstraction
 
 ```@docs
+Flux.modules
 Flux.destructure
 Flux.nfan
 ```

src/layers/normalise.jl

@@ -158,7 +158,7 @@ end
 
 function Base.show(io::IO, l::LayerNorm)
   print(io, "LayerNorm($(l.size)")
-  a.λ == identity || print(io, ", $(a.λ)")
+  l.λ == identity || print(io, ", $(l.λ)")
   hasaffine(l) || print(io, ", affine=false")
   print(io, ")")
 end

src/utils.jl

@@ -607,3 +607,50 @@ function throttle(f, timeout; leading=true, trailing=false)
     return result
   end
 end
+
+
+"""
+    modules(m)
+
+Return an iterator over non-leaf objects
+that can be reached from `m` through recursion
+on the children given by [`trainable`](@ref).
+
+It can be used to apply a regularization
+over certain specific modules or subsets of
+the parameters (e.g. the weights but not the biases).
+
+# Examples
+
+```jldoctest
+julia> m1 = Chain(Dense(28^2, 64), BatchNorm(64, relu))
+Chain(Dense(784, 64), BatchNorm(64, relu))
+
+julia> m2 = Chain(model1, Dense(64, 10))
+Chain(Chain(Dense(784, 64), BatchNorm(64, relu)), Dense(64, 10))
+
+julia> Flux.modules(m2)
+5-element Vector{Any}:
+ Chain(Chain(Dense(784, 64), BatchNorm(64, relu)), Dense(64, 10))
+ Chain(Dense(784, 64), BatchNorm(64, relu))
+ Dense(784, 64)
+ BatchNorm(64, relu)
+ Dense(64, 10)
+
+julia> L2(model) = sum(sum(abs2, m.weight) for m in Flux.modules(model) if m isa Dense)
+```
+"""
+modules(m) = [x for x in traverse_trainables(m) if !isleaflike(x)]
+
+function traverse_trainables(x, cache=[])
+  x in cache && return cache
+  push!(cache, x)
+  foreach(y -> traverse_trainables(y, cache), trainable(x))
+  return cache
+end
+
+isleaflike(x) = trainable(x) === ()
+isleaflike(::Tuple{Vararg{<:Number}}) = true
+isleaflike(::Tuple{Vararg{<:AbstractArray}}) = true
+
+@nograd modules

test/utils.jl

@@ -344,3 +344,19 @@ end
   trainmode!(c)
   @test !c[1].testing
 end
+
+@testset "modules" begin
+  m1 = Conv((2,3), 4=>5; pad=6, stride=7)
+  m2 = LayerNorm(8)
+  m3 = m2.diag
+  m4 = SkipConnection(m1, +)
+  m5 =  Chain(m4, m2)
+  modules = Flux.modules(m5)
+  # Depth-first descent
+  @test length(modules) == 5
+  @test modules[1] === m5
+  @test modules[2] === m4
+  @test modules[3] === m1
+  @test modules[4] === m2
+  @test modules[5] === m3
+end