destructure returns only trainable params

Author: CarloLucibello

src/functor.jl

@@ -38,6 +38,7 @@ Possible values include:
 """
 trainmode!(m, mode = true) = mode isa Bool ? testmode!(m, !mode) : testmode!(m, mode)
 
+# push!(::Params, x) automatically discards already seen arrays
 params!(p::Params, x::AbstractArray{<:Number}, seen = IdSet()) = push!(p, x)
 
 function params!(p::Params, x, seen = IdSet())

src/utils.jl

@@ -630,9 +630,11 @@ end
 
 function _restructure(m, xs)
   i = 0
-  m̄ = fmap(m) do x
-    x isa AbstractArray || return x
-    x = reshape(xs[i.+(1:length(x))], size(x))
+  filter = (x, c) -> any(y -> c === y, trainable(x))
+  walk = filtered_walk(filter)
+  m̄ = fmap(m; walk) do x
+    x isa AbstractArray{<:Number} || return x
+    x = reshape(xs[i .+ (1:length(x))], size(x))
     i += length(x)
     return x
   end
@@ -673,13 +675,28 @@ modifications to the weight vector (for example, with a hypernetwork).
 """
 function destructure(m)
   xs = Zygote.Buffer([])
-  fmap(m) do x
-    x isa AbstractArray && push!(xs, x)
+  filter = (x, c) -> any(y -> c === y, trainable(x))
+  walk = filtered_walk(filter)
+  fmap(m; walk) do x
+    x isa AbstractArray{<:Number} && push!(xs, x)
     return x
   end
   return vcat(vec.(copy(xs))...), p -> _restructure(m, p)
 end
 
+function filtered_walk(filter)
+  function walk(f, x)
+    children, reconstruct = functor(x)
+    mappedchildren = map(children) do c
+      filter(x, c) ? f(c) : c
+    end
+    reconstruct(mappedchildren)
+  end
+  return walk
+end
+
+@functor Base.RefValue
+
 # Other
 
 """

test/utils.jl

@@ -407,7 +407,24 @@ end
       ∇m = gradient(m -> sum(m(x)), m)[1]
       p, re = destructure(m)
       ∇p = gradient(θ -> sum(re(θ)(x)), p)[1]
-      @test ∇p ≈ destructure(∇m)[1]
+      # @show size(∇p)
+      # @show size(destructure(∇m)[1])
+      # @show norm(∇p - destructure(∇m)[1])
+      @test ∇p ≈ destructure(∇m)[1] atol=1e-4
+    end
+
+    @testset "destructure with buffers" begin
+      p, re = destructure(BatchNorm(10))
+      @test length(p) == 20
+
+      # https://github.com/FluxML/Flux.jl/issues/1727
+      x = rand(Float32, 2, 3)
+      gs, back = Flux.pullback(x, p) do x, p
+          vec(re(p)(x))
+      end
+      @test_nowarn b = back(a)
+      @test b[1] == size(x)
+      @test b[2] == size(p)
     end
   end
 end