Merge pull request #2041 from darsnack/load-filter

Author: darsnack

src/loading.jl

@@ -28,6 +28,10 @@ _tie_check(dst, src) = true
 
 _bool_tie_check(dst, src) = true
 
+_filter_children(f, children::NamedTuple) =
+  NamedTuple(filter(kv -> f(kv[2]), pairs(children)))
+_filter_children(f, children) = filter(f, children)
+
 """
     loadmodel!(dst, src)
 
@@ -41,7 +45,7 @@ Zero bias vectors and `bias=false` are considered equivalent
 
 # Examples
 ```julia
-julia> dst = Chain(Dense(Flux.ones32(2, 5, tanh)), Dense(2 => 1; bias = [1f0]))
+julia> dst = Chain(Dense(Flux.ones32(2, 5), Flux.ones32(2), tanh), Dense(2 => 1; bias = [1f0]))
 Chain(
   Dense(5 => 2, tanh),                  # 12 parameters
   Dense(2 => 1),                        # 3 parameters
@@ -77,9 +81,9 @@ however, attempting to copy a non-zero array to an inactive parameter will throw
 Likewise, copying a `src` value of `false` to any `dst` array is valid,
 but copying a `src` value of `true` will error.
 """
-function loadmodel!(dst, src; cache = Base.IdSet())
-  ldsts, _ = functor(dst)
-  lsrcs, _ = functor(src)
+function loadmodel!(dst, src; filter = _ -> true, cache = Base.IdSet())
+  ldsts = _filter_children(filter, functor(dst)[1])
+  lsrcs = _filter_children(filter, functor(src)[1])
   (keys(ldsts) == keys(lsrcs)) ||
     throw(ArgumentError("Tried to load $src into $dst but the structures do not match."))
 
@@ -91,7 +95,7 @@ function loadmodel!(dst, src; cache = Base.IdSet())
       push!(cache, ldst)
       loadleaf!(ldst, lsrc, err)
     else # this isn't a leaf
-      loadmodel!(ldst, lsrc; cache = cache)
+      loadmodel!(ldst, lsrc; filter = filter, cache = cache)
     end
   end
 

test/utils.jl

@@ -508,6 +508,42 @@ end
     # loading into tied weights with absent parameter is bad when the dst != zero
     m2[1].bias .= 1
     @test_throws ErrorException loadmodel!(m1, m2)
+
+    @testset "loadmodel! & filter" begin
+      m1 = Chain(Dense(10, 5), Dense(5, 2, relu))
+      m2 = Chain(Dense(10, 5), Dropout(0.2), Dense(5, 2))
+      m3 = Chain(Dense(10, 5), Dense(5, 2, relu))
+
+      # this will not error cause Dropout is skipped
+      loadmodel!(m1, m2; filter = x -> !(x isa Dropout))
+      @test m1[1].weight == m2[1].weight
+      @test m1[2].weight == m2[3].weight
+
+      # this will not error cause Dropout is skipped
+      loadmodel!(m2, m3; filter = x -> !(x isa Dropout))
+      @test m3[1].weight == m2[1].weight
+      @test m3[2].weight == m2[3].weight
+    end
+
+    @testset "loadmodel! & absent bias" begin
+      m0 = Chain(Dense(2 => 3; bias=false, init = Flux.ones32), Dense(3 => 1))
+      m1 = Chain(Dense(2 => 3; bias = Flux.randn32(3)), Dense(3 => 1))
+      m2 = Chain(Dense(Float32[1 2; 3 4; 5 6], Float32[7, 8, 9]), Dense(3 => 1))
+    
+      Flux.loadmodel!(m1, m2)
+      @test m1[1].bias == 7:9
+      @test sum(m1[1].weight) == 21
+    
+      # load from a model without bias -- should ideally recognise the `false` but `Params` doesn't store it
+      m1 = Flux.loadmodel!(m1, m0)
+      @test iszero(m1[1].bias)
+      @test sum(m1[1].weight) == 6  # written before error
+    
+      # load into a model without bias -- should it ignore the parameter which has no home, or error?
+      m0 = Flux.loadmodel!(m0, m2)
+      @test iszero(m0[1].bias)  # obviously unchanged
+      @test sum(m0[1].weight) == 21
+    end
   end
 
   @testset "destructure" begin
@@ -529,26 +565,6 @@ end
   end
 end
 
-@testset "loadmodel! & absent bias" begin
-  m0 = Chain(Dense(2 => 3; bias=false, init = Flux.ones32), Dense(3 => 1))
-  m1 = Chain(Dense(2 => 3; bias = Flux.randn32(3)), Dense(3 => 1))
-  m2 = Chain(Dense(Float32[1 2; 3 4; 5 6], Float32[7, 8, 9]), Dense(3 => 1))
-
-  Flux.loadmodel!(m1, m2)
-  @test m1[1].bias == 7:9
-  @test sum(m1[1].weight) == 21
-
-  # load from a model without bias -- should ideally recognise the `false` but `Params` doesn't store it
-  m1 = Flux.loadmodel!(m1, m0)
-  @test iszero(m1[1].bias)
-  @test sum(m1[1].weight) == 6  # written before error
-
-  # load into a model without bias -- should it ignore the parameter which has no home, or error?
-  m0 = Flux.loadmodel!(m0, m2)
-  @test iszero(m0[1].bias)  # obviously unchanged
-  @test sum(m0[1].weight) == 21
-end
-
 @testset "Train and test mode" begin
   mutable struct DummyLayer
     testing::Bool