Add better support for loadmodel! w/ tied parameters and address some other review comments

Author: darsnack

src/deprecations.jl

@@ -15,8 +15,6 @@ zeros(T::Type, dims...) = Base.zeros(T, dims...)
 ones32(::Type, dims...) = throw(ArgumentError("Flux.ones32 is always Float32, use Base.ones to specify the element type"))
 zeros32(::Type, dims...) = throw(ArgumentError("Flux.zeros32 is always Float32, use Base.zeros to specify the element type"))
 
-@deprecate loadparams!(m, xs) loadmodel!(m, xs)
-
 # v0.13 deprecations
 
 function Broadcast.broadcasted(f::Recur, args...)
@@ -50,6 +48,16 @@ function Diagonal(size::Tuple; kw...)
   Scale(size...; kw...)
 end
 
+# Deprecate this eventually once saving models w/o structure is no more
+function loadparams!(m, xs)
+  Base.depwarn("loadparams! will be deprecated eventually. Use loadmodel! instead.", :loadparams!)
+  for (p, x) in zip(params(m), xs)
+    size(p) == size(x) ||
+      error("Expected param size $(size(p)), got $(size(x))")
+    copyto!(p, x)
+  end
+end
+
 # Channel notation: Changed to match Conv, but very softly deprecated!
 # Perhaps change to @deprecate for v0.14, but there is no plan to remove these.
 Dense(in::Integer, out::Integer, σ = identity; kw...) =

src/loading.jl

@@ -1,31 +1,42 @@
-"""
-    loadleaf!(dst, src, err)
-
-Copy `src` to `dst` or throw `err` when their sizes are mismatched.
-By default, use `copyto!` when `dst` and `src` are arrays.
-When only `dst` is an array, set every element to `src`.
-Otherwise, just return `dst`.
-"""
 loadleaf!(dst, src, err) = dst
-function loadleaf!(dst::AbstractArray, src, err)
-  dst .= src
+function loadleaf!(dst::AbstractArray, src::Bool, err)
+  if iszero(src)
+    dst .= src
+  else
+    error("Cannot copy boolean parameter == true to non-zero parameter.")
+  end
   return dst
 end
+loadleaf!(dst::Bool, src::AbstractArray, err) = iszero(dst) ? dst :
+  error("Cannot copy non-zero parameter to boolean parameter == true.")
 function loadleaf!(dst::AbstractArray, src::AbstractArray, err)
   (size(dst) == size(src)) || throw(err)
   copyto!(dst, src)
 end
 
+_parent(x) = x
+_parent(x::AbstractArray) = parent(x)
+
+_tie_check(dst::AbstractArray, src::AbstractArray) = dst == src
+_tie_check(dst, src) = true
+
+_bool_tie_check(dst::Bool, src::AbstractArray) = iszero(dst)
+_bool_tie_check(dst::AbstractArray, src::Bool) = iszero(dst) && iszero(src)
+_bool_tie_check(dst, src) = true
+
 """
     loadmodel!(dst, src)
 
 Copy all the parameters (trainable and non-trainable) from `src` to `dst`.
 
 `loadmodel!` recursively walks the [`Functors.children`](@ref) of `dst` and `src`
-calling `loadleaf!` on any pair of children where [`Functors.isleaf`](@ref) is true.
+calling `copyto!` on any pair of children where [`Functors.isleaf`](@ref) is true.
+It also handles "absent" parameters such as `bias == false`.
 It throws an error whenever:
 - `dst` and `src` do not share the same fields (at any level)
 - the sizes of leaf nodes are mismatched between `dst` and `src`
+- `dst` is a "tied" parameter (e.g. `transpose` of another parameter) and
+  loaded into multiple times with mismatched source values
 
 ```julia
 julia> using Flux: loadmodel!
@@ -56,30 +67,29 @@ true
 julia> dst[2].bias == src[2].bias
 true
 ```
-
-See [`Flux.loadleaf!`](@ref) for more details on the copy behavior.
-
-!!! warning
-    This function allows `src` to be a `Params` for backwards-compatibility.
-    You should avoid using `loadmodel!` this way, because it skips most of the structural
-    checking used when `src` is also a nested structure. Silent errors may occur.
 """
-function loadmodel!(m, xs::Params)
-  for (p, x) in zip(params(m), xs)
-    size(p) == size(x) ||
-      error("Expected param size $(size(p)), got $(size(x))")
-    copyto!(p, x)
-  end
-end
-function loadmodel!(dst, src)
+function loadmodel!(dst, src; cache = Base.IdSet())
   ldsts, _ = functor(dst)
   lsrcs, _ = functor(src)
   (keys(ldsts) == keys(lsrcs)) ||
     throw(ArgumentError("Tried to load $src into $dst but the structures do not match."))
 
   err = DimensionMismatch("Tried to load $src into $dst but the parameter sizes do not match.")
   foreach(ldsts, lsrcs) do ldst, lsrc
-    Functors.isleaf(ldst) ? loadleaf!(ldst, lsrc, err) : loadmodel!(ldst, lsrc)
+    if _parent(ldst) in cache # we already loaded this parameter before
+      if !_bool_tie_check(ldst, lsrc) # special case to handle tied + absent parameters
+        error("Encountered tied parameter with boolean source at some nodes and non-boolean sources at others.")
+      elseif _tie_check(ldst, lsrc) # the arrays match and we already loaded (or these are not arrays)
+        return ldst
+      else # tied dst but mismatched src case
+        error("Encountered tied destination parameters with untied and mismatched sources.")
+      end
+    elseif Functors.isleaf(ldst) # our first time loading this leaf
+      push!(cache, ldst)
+      loadleaf!(ldst, lsrc, err)
+    else # this isn't a leaf
+      loadmodel!(ldst, lsrc; cache = cache)
+    end
   end
 
   return dst

test/utils.jl

@@ -2,7 +2,7 @@ using Flux
 using Flux: throttle, nfan, glorot_uniform, glorot_normal,
              kaiming_normal, kaiming_uniform, orthogonal, truncated_normal,
              sparse_init, identity_init, stack, unstack, batch, unbatch,
-             unsqueeze, params, loadmodel!
+             unsqueeze, params, loadparams!, loadmodel!
 using StatsBase: var, std
 using Statistics, LinearAlgebra
 using Random
@@ -366,26 +366,24 @@ end
     @test_skip typeof(l1.bias) === typeof(l2.bias)
   end
 
-  @testset "loadmodel!" begin
+  @testset "loadparams!" begin
     pars(w, b) = [w, b]
     pars(l) = pars(l.weight, l.bias)
     pararray(m) = mapreduce(pars, vcat, m)
     weights(m) = mapreduce(l -> [l.weight], vcat, m)
     @testset "Bias type $bt" for bt in (Flux.zeros32, nobias)
       m = dm(bt)
-      Flux.loadmodel!(m, params(m))
+      Flux.loadparams!(m, params(m))
       testdense(m, bt)
     end
   end
 
   @testset "loadmodel!(dst, src)" begin
-    import Flux: loadmodel!, Zeros
-
     m1 = Chain(Dense(10, 5), Dense(5, 2, relu))
     m2 = Chain(Dense(10, 5), Dense(5, 2))
     m3 = Chain(Conv((3, 3), 3 => 16), Dense(5, 2))
     m4 = Chain(Dense(10, 6), Dense(6, 2))
-    m5 = Chain(Dense(10, 5), Parallel(+, Dense(Flux.ones32(2, 5), Zeros()), Dense(5, 2)))
+    m5 = Chain(Dense(10, 5), Parallel(+, Dense(Flux.ones32(2, 5), false), Dense(5, 2)))
     m6 = Chain(Dense(10, 5), Parallel(+, Dense(5, 2), Dense(5, 2)))
 
     loadmodel!(m1, m2)
@@ -408,6 +406,7 @@ end
     # size mismatches throw an error
     @test_throws DimensionMismatch loadmodel!(m1, m4)
 
+    # tests for BatchNorm and Dropout
     m1 = Chain(Conv((3, 3), 3 => 16), BatchNorm(16), Flux.flatten, Dropout(0.2))
     m2 = Chain(Conv((3, 3), 3 => 16), BatchNorm(16), x -> reshape(x, :, size(x)[end]), Dropout(0.1))
     m2[2].μ .= rand(Float32, size(m2[2].μ)...)
@@ -484,16 +483,32 @@ end
     @test chain2[2].bias != chain1[2].bias
 
     # test shared weights
-    m1 = Chain(Dense(10 => 5), Dense(5 => 2))
-    m2 = Chain(Dense(transpose(m1[2].weight)), Dense(permutedims(m1[1].weight)))
-    m3 = Chain(Dense(m1[1].weight), Dense(m1[2].weight))
-    m2[2].weight .= 1f0
-    loadmodel!(m1, m3)
-    @test m1[2].weight === parent(m2[1].weight)
-    @test m1[2].weight == transpose(m2[1].weight)
-    @test m1[1].weight === m3[1].weight
-    @test m2[2].weight != transpose(m1[1].weight)
-    @test m3[2].weight == transpose(m2[1].weight)
+    encoder_dst = Chain(Dense(10 => 5), Dense(5 => 2))
+    decoder_dst = Chain(Dense(transpose(encoder_dst[2].weight)),
+                        Dense(permutedims(encoder_dst[1].weight)))
+    encoder_src = Chain(Dense(10 => 5), Dense(5 => 2))
+    decoder_src = Chain(Dense(transpose(encoder_src[2].weight)),
+                        Dense(5 => 10))
+    # matched weights are okay
+    m1 = Chain(encoder_dst, decoder_dst)
+    m2 = Chain(encoder_src, decoder_src)
+    loadmodel!(m1, m2)
+    @test m1[1][2].weight === parent(m1[2][1].weight)
+    @test m1[1][1].weight == m2[1][1].weight
+    @test m1[1][1].weight != permutedims(m1[2][2].weight)
+    # mismatched weights are an error
+    m2 = Chain(Chain(Dense(10 => 5), Dense(5 => 2)), Chain(Dense(2 => 5), Dense(5 => 10)))
+    @test_throws ErrorException loadmodel!(m1, m2)
+    # loading into tied weights with absent parameter is okay when the dst == zero
+    b = Flux.zeros32(5)
+    m1 = Chain(Dense(10 => 5; bias = b), Dense(5 => 5; bias = b))
+    m2 = Chain(Dense(10 => 5; bias = Flux.zeros32(5)), Dense(5 => 5; bias = false))
+    loadmodel!(m1, m2)
+    @test m1[1].bias === m1[2].bias
+    @test iszero(m1[1].bias)
+    # loading into tied weights with absent parameter is bad when the dst != zero
+    m2[1].bias .= 1
+    @test_throws ErrorException loadmodel!(m1, m2)
   end
 
   @testset "destructure" begin