Add initial implementation

Author: darsnack

Project.toml

@@ -29,7 +29,7 @@ Adapt = "3.0"
 ArrayInterface = "3.1, 4, 5"
 CUDA = "3"
 ChainRulesCore = "1.12"
-Functors = "0.2.1"
+Functors = "0.2.8"
 MLUtils = "0.2"
 MacroTools = "0.5"
 NNlib = "0.8.2"

src/Flux.jl

@@ -46,6 +46,8 @@ include("layers/normalise.jl")
 include("layers/upsample.jl")
 include("layers/show.jl")
 
+include("loading.jl")
+
 include("outputsize.jl")
 
 include("data/Data.jl")

src/deprecations.jl

@@ -15,6 +15,8 @@ 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...)

src/functor.jl

@@ -85,14 +85,6 @@ function params(m...)
   return ps
 end
 
-function loadparams!(m, xs)
-  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
-
 struct FluxCUDAAdaptor end
 adapt_storage(to::FluxCUDAAdaptor, x) = CUDA.cu(x)
 adapt_storage(to::FluxCUDAAdaptor, x::Zygote.FillArrays.AbstractFill) = CUDA.cu(collect(x))

src/loading.jl

@@ -0,0 +1,46 @@
+_loadleaf(x) = isleaf(x)
+for T in [:Dense, :Diagonal, :Bilinear, :Embedding,
+          :Conv, :ConvTranspose, :DepthwiseConv, :CrossCor]
+  @eval _loadleaf(::$T) = true
+end
+
+loadto!(x, 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
+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
+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
+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)
+    size(p) == size(x) ||
+      error("Expected param size $(size(p)), got $(size(x))")
+    copyto!(p, x)
+  end
+end
+loadmodel!(m, xs::AbstractVector) = loadmodel!(m, params(xs))
+loadmodel!(m, m̄) = fmap(loadto!, m, m̄; exclude = _loadleaf)

test/utils.jl

@@ -373,11 +373,28 @@ end
     weights(m) = mapreduce(l -> [l.weight], vcat, m)
     @testset "Bias type $bt" for bt in (Flux.zeros32, nobias)
       m = dm(bt)
-      loadparams!(m, params(m))
+      loadmodel!(m, params(m))
       testdense(m, bt)
     end
   end
 
+  @testset "loadmodel!(m, m̄)" begin
+    import Flux: loadmodel!
+
+    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))
+
+    loadmodel!(m1, m2)
+    @test m1[1].weight == m2[1].weight
+    @test m1[1].bias == m2[1].bias
+    @test m1[2].σ == relu
+
+    @test_throws ArgumentError loadmodel!(m1, m3)
+    @test_throws DimensionMismatch loadmodel!(m1, m4)
+  end
+
   @testset "destructure" begin
     import Flux: destructure
     @testset "Bias type $bt" for bt in (zeros, nobias)