add freeze/thaw

Author: mcabbott

docs/src/api.md

@@ -35,6 +35,8 @@ Optimisers.setup
 Optimisers.update
 Optimisers.update!
 Optimisers.adjust(::Any, ::Real)
+Optimisers.freeze!
+Optimisers.thaw!
 ```
 
 Calling `Functors.@functor` on your model's layer types by default causes

src/adjust.jl

@@ -1,3 +1,56 @@
+###
+### freeze!
+###
+
+"""
+    Optimisers.freeze!(tree)
+
+Temporarily alters the state `tree = setup(rule, model)` so that parameters will not be updated.
+Can be applied to the state corresponding to only part of a model, for instance `model.layers[1]`.
+Un-done by [`thaw!`](@ref Optimisers.thaw).
+
+# Example
+```jldoctest
+julia> m = (x = ([1.0], 2.0), y = [3.0]);
+
+julia> s = Optimisers.setup(Momentum(), m);
+
+julia> Optimisers.freeze!(s.x)
+
+julia> Optimisers.update!(s, m, (x = ([pi], 10pi), y = [100pi]));  # with fake gradient
+
+julia> m
+(x = ([1.0], 2.0), y = [-0.14159258336972558])
+
+julia> s  # Leaf(..., true) means frozen
+(x = (Leaf(Momentum{Float32}(0.01, 0.9), [0.0], true), ()), y = Leaf(Momentum{Float32}(0.01, 0.9), [3.14159]))
+
+julia> Optimisers.thaw!(s)
+
+julia> s.x[1]
+Leaf(Momentum{Float32}(0.01, 0.9), [0.0])
+```
+"""
+freeze!(tree) = (fmapstructure(freeze!, tree; exclude = x -> x isa Leaf); nothing)
+freeze!(ℓ::Leaf) = (ℓ.frozen = true; nothing)
+
+"""
+    Optimisers.thaw!(tree)
+
+Un-does [`freeze!`](@ref Optimisers.freeze!) for all parameters,
+mutating every `Leaf(rule, state, true)` to `Leaf(rule, state, false)`.
+"""
+thaw!(tree) = (fmapstructure(thaw!, tree; exclude = x -> x isa Leaf); nothing)
+thaw!(ℓ::Leaf) = (ℓ.frozen = false; nothing)
+
+freeze!(::Union{Number, AbstractArray{<:Number}}) = throw(ArgumentError(
+  "`freeze!` must not be applied to a model, only to the state tree from `setup`"))
+thaw!(::Union{Number, AbstractArray{<:Number}}) = throw(ArgumentError(
+  "`thaw!` must not be applied to a model, only to the state tree from `setup`"))
+
+###
+### adjust
+###
 
 """
     Optimisers.adjust(tree, η) -> tree

src/interface.jl

@@ -10,10 +10,12 @@ abstract type AbstractRule end
 ### setup
 ###
 
-mutable struct Leaf{R,S}  # mutable so that its identity encodes parameter sharing
+mutable struct Leaf{R,S}  # mutable so that its identity encodes parameter sharing...
   rule::R
   state::S
+  frozen::Bool  # ... and to allow freeze! to act on this.
 end
+Leaf(rule, state) = Leaf(rule, state, false)
 
 @functor Leaf
 
@@ -46,6 +48,7 @@ function Base.show(io::IO, ℓ::Leaf)  # show method is mostly to hide its long
   ioc = IOContext(io, :compact => true)
   print(ioc, "Leaf(", ℓ.rule, ", ")
   show(ioc, ℓ.state)
+  ℓ.frozen && print(ioc, ", true")
   print(ioc, ")")
 end
 
@@ -83,6 +86,7 @@ function _update!(tree, x; grads, params)
 end
 function _update!(ℓ::Leaf, x; grads, params)
   haskey(params, (ℓ,x)) && return params[(ℓ,x)]
+  ℓ.frozen && return x
   params[(ℓ,x)] = if haskey(grads, ℓ)
     ℓ.state, x̄′ = apply!(ℓ.rule, ℓ.state, x, grads[ℓ]...)
     subtract!(x, x̄′)