Better docs

Author: darsnack

docs/src/saving.md

@@ -85,10 +85,23 @@ This ensures that the model loaded from `"mymodel.bson"` matches the structure o
 
 ```@docs
 Flux.loadmodel!
+Flux.loadto!
 Flux.isloadleaf
 Flux.loadleaf!
 ```
 
+### Customizing `loadmodel!` for a custom layer
+
+By default, [`loadmodel!`](@ref) will recursively walk a nested model (like a `Chain`) using [`Functors.fmap`](@ref) until it encounters a loading *leaf node*. A leaf node is defined as any node for which [`Flux.isloadleaf`](@ref) returns `true`. For example, consider the model
+
+```julia
+model = Chain(Dense(10 => 5), Parallel(+, Dense(5 => 2), Dense(5 => 2)))
+```
+
+Here, the `Chain` and `Parallel` layers are not leaf nodes, but all the `Dense` layers are leaf nodes. This makes sense, because `Dense` layers are the ones with parameters that we need to copy. The default behavior for [`Flux.isloadleaf`](@ref) should work for most custom layers, but you can override this function for your type.
+
+Once a pair of leaf nodes is encountered, `loadmodel!` will call [`Flux.loadto!](@ref) on them. By default, this just copies the parameters from one leaf node to the other, but you can customize the behavior by overriding `loadto!` for your pair of types.
+
 ## Checkpointing
 
 In longer training runs it's a good idea to periodically save your model, so that you can resume if training is interrupted (for example, if there's a power cut). You can do this by saving the model in the [callback provided to `train!`](training/training.md).

src/loading.jl

@@ -27,7 +27,17 @@ function loadleaf!(x::AbstractArray, x̄::AbstractArray, err)
   copyto!(x, x̄)
 end
 
-function _loadto!(m, m̄)
+"""
+    loadto!(m, m̄)
+
+Load a leaf node `m̄` into `m`.
+
+By default, call [`Flux.loadleaf!`](@ref) on each pair of children
+in `zip(Functors.children(m), Functors.children(m̄))`.
+"""
+function loadto!(m::T, m̄::S) where {T, S}
+  (nameof(T) == nameof(S)) || throw(ArgumentError("Tried to load $m̄ into $m."))
+  
   ls, _ = functor(m)
   l̄s, _ = functor(m̄)
   (keys(ls) == keys(l̄s)) ||
@@ -38,10 +48,6 @@ function _loadto!(m, m̄)
 
   return 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
 
 """
     loadmodel!(m, m̄)
@@ -56,8 +62,38 @@ throwing an error whenever:
 - `x` and `x̄` do not share the same fields
 - the parameter sizes are mismatched between `x` and `x̄`
 
-See [`loadleaf!`](@ref) for more details on the copy behavior.
-See [`isloadleaf`](@ref) for more details on which layers are considered leaves.
+```julia
+julia> using Flux: loadmodel!
+
+julia> m = Chain(Dense(Flux.ones32(2, 5)), Dense(2 => 1))
+Chain(
+  Dense(5 => 2),                        # 12 parameters
+  Dense(2 => 1),                        # 3 parameters
+)                   # Total: 4 arrays, 15 parameters, 316 bytes.
+
+julia> m̄ = Chain(Dense(5 => 2), Dense(2 => 1));
+
+julia> all(isone, m[1].weight)
+true
+
+julia> m = loadmodel!(m, m̄)
+Chain(
+  Dense(5 => 2),                        # 12 parameters
+  Dense(2 => 1),                        # 3 parameters
+)                   # Total: 4 arrays, 15 parameters, 316 bytes.
+
+julia> all(isone, m[1].weight)
+false
+
+julia> m[1].weight == m̄[1].weight
+true
+
+julia> m[2].bias == m̄[2].bias
+true
+```
+
+See [`Flux.loadleaf!`](@ref) for more details on the copy behavior.
+See [`Flux.isloadleaf`](@ref) for more details on which layers are considered leaves.
 
 !!! warning
     This function allows `m̄` to be a vector or `Params` for backwards-compatibility.