Refactor to allow better support for loading errors with custom models

Author: darsnack

src/loading.jl

@@ -1,40 +1,27 @@
-_loadleaf(x) = isleaf(x)
-for T in [:Dense, :Diagonal, :Bilinear, :Embedding,
-          :Conv, :ConvTranspose, :DepthwiseConv, :CrossCor]
-  @eval _loadleaf(::$T) = true
-end
+isloadleaf(x) = all(Functors.isleaf, Functors.children(x))
 
-loadto!(x, x̄) = x
-loadto!(x::Zeros, x̄) = x
-loadto!(x::AbstractArray, x̄::AbstractArray) = copyto!(x, x̄)
-for T in [:Dense, :Bilinear, :Conv, :ConvTranspose, :DepthwiseConv, :CrossCor]
-  @eval begin
-    function loadto!(m::$T, m̄::$T)
-      if (size(m.weight) != size(m̄.weight)) || (size(m.bias) != size(m̄.bias))
-        throw(DimensionMismatch("Tried to load $m̄ into $m but the parameter sizes do not match."))
-      else
-        return fmap(loadto!, m, m̄)
-      end
-    end
-    loadto!(m::$T, m̄) = throw(ArgumentError("Tried to load $m̄ into $m."))
-  end
+loadnumeric!(x, x̄, err) = x
+loadnumeric!(x::Zeros, x̄, err) = x
+function loadnumeric!(x::AbstractArray, x̄::AbstractArray, err)
+    (size(x) == size(x̄)) || throw(err)
+    copyto!(x, x̄)
 end
-function loadto!(m::Diagonal, m̄::Diagonal)
-  if (size(m.α) != size(m̄.α)) || (size(m.β) != size(m̄.β))
-    throw(DimensionMismatch("Tried to load $m̄ into $m but the parameter sizes do not match."))
-  else
-    return fmap(loadto!, m, m̄)
-  end
+
+function _loadto!(m, m̄)
+    ls, _ = functor(m)
+    l̄s, _ = functor(m̄)
+    (keys(ls) == keys(l̄s)) ||
+        throw(ArgumentError("Tried to load $m̄ into $m but the structures do not match."))
+
+    err = DimensionMismatch("Tried to load $m̄ into $m but the parameter sizes do not match.")
+    foreach((l, l̄) -> loadnumeric!(l, l̄, err), ls, l̄s)
+
+    return m
 end
-loadto!(m::Diagonal, m̄) = throw(ArgumentError("Tried to load $m̄ into $m."))
-function loadto!(m::Embedding, m̄::Embedding)
-  if size(m.weight) != size(m̄.weight)
-    throw(DimensionMismatch("Tried to load $m̄ into $m but the parameter sizes do not match."))
-  else
-    return fmap(loadto!, m, m̄)
-  end
+function loadto!(m::T, m̄::S) where {T, S}
+    (nameof(T) == nameof(S)) || throw(ArgumentError("Tried to load $m̄ into $m."))
+    _loadto!(m, m̄)
 end
-loadto!(m::Embedding, m̄) = throw(ArgumentError("Tried to load $m̄ into $m."))
 
 function loadmodel!(m, xs::Params)
   for (p, x) in zip(params(m), xs)
@@ -44,4 +31,4 @@ function loadmodel!(m, xs::Params)
   end
 end
 loadmodel!(m, xs::AbstractVector) = loadmodel!(m, params(xs))
-loadmodel!(m, m̄) = fmap(loadto!, m, m̄; exclude = _loadleaf)
+loadmodel!(m, m̄) = fmap(loadto!, m, m̄; exclude = isloadleaf)