Implement APIs of freeze parameters and freeze layers

docs/src/models/advanced.md

@@ -62,12 +62,32 @@ During training, the gradients will only be computed for (and applied to) the la
 Flux.params(m[1], m[3:end])
 ```
 
-Sometimes, a more fine-tuned control is needed. 
-We can freeze a specific parameter of a specific layer which already entered a `Params` object `ps`, 
+Sometimes, a more fine-tuned control is needed.
+We can freeze a specific parameter of a specific layer which already entered a `Params` object `ps`,
 by simply deleting it from `ps`:
 
 ```julia
 ps = params(m)
-delete!(ps, m[2].b) 
+delete!(ps, m[2].b)
 ```
 
+The `freezelayers!` function prevents parameters of multiple layers from being updated. The following example stops the parameters of the first and third layer from being updated.
+
+```julia
+m = Chain(
+    Dense(4, 2),
+    Dense(2, 3),
+    Dense(3, 4),
+    Dense(4, 3),
+    softmax
+)
+
+freezed_layer_indexes = [1, 3]
+
+train!(
+    (x, y) -> crossentropy(m(x), y),
+    freezelayers!(params(m), m, freezed_layer_indexes),
+    [(rand(4, 10), rand(1, 10))],
+    ADAM(0.005)
+)
+```

docs/src/training/training.md

@@ -75,7 +75,7 @@ Training data can be conveniently  partitioned for mini-batch training using the
 ```julia
 X = rand(28, 28, 60000)
 Y = rand(0:9, 60000)
-data = DataLoader(X, Y, batchsize=128) 
+data = DataLoader(X, Y, batchsize=128)
 ```
 
 Note that, by default, `train!` only loops over the data once (a single "epoch").

src/optimise/Optimise.jl

@@ -1,8 +1,8 @@
 module Optimise
 
-export train!, update!,
+export train!, update!, freezelayers!,
 	SGD, Descent, ADAM, Momentum, Nesterov, RMSProp,
-	ADAGrad, AdaMax, ADADelta, AMSGrad, NADAM, ADAMW,RADAM, 
+	ADAGrad, AdaMax, ADADelta, AMSGrad, NADAM, ADAMW,RADAM,
 	InvDecay, ExpDecay, WeightDecay, stop, Optimiser
 
 include("optimisers.jl")

src/optimise/train.jl

@@ -6,13 +6,13 @@ import Zygote: Params, gradient
     update!(opt, p, g)
     update!(opt, ps::Params, gs)
 
-Perform an update step of the parameters `ps` (or the single parameter `p`) 
+Perform an update step of the parameters `ps` (or the single parameter `p`)
 according to optimizer `opt`  and the gradients `gs` (the gradient `g`).
 
-As a result, the parameters are mutated and the optimizer's internal state may change. 
+As a result, the parameters are mutated and the optimizer's internal state may change.
 
   update!(x, x̄)
-  
+
 Update the array `x` according to `x .-= x̄`.
 """
 function update!(x::AbstractArray, x̄)
@@ -30,6 +30,19 @@ function update!(opt, xs::Params, gs)
   end
 end
 
+# Fix layers
+function freezelayers!(ps::Params, m, layer_indexes::Vector{Int64})
+    for layer_index in layer_indexes
+        layer = m.layers[layer_index]
+        param_names = fieldnames(typeof(layer))
+        for param_name in param_names
+            delete!(ps, getfield(layer, param_name))
+        end
+    end
+
+    return ps
+end
+
 # Callback niceties
 call(f, xs...) = f(xs...)
 runall(f) = f
@@ -61,7 +74,7 @@ end
 For each datapoint `d` in `data` computes the gradient of `loss(d...)` through
 backpropagation and calls the optimizer `opt`.
 
-In case datapoints `d` are of numeric array type, assumes no splatting is needed 
+In case datapoints `d` are of numeric array type, assumes no splatting is needed
 and computes the gradient of `loss(d)`.
 
 Takes a callback as keyword argument `cb`. For example, this will print "training"

test/freezelayers.jl

@@ -0,0 +1,51 @@
+using Flux: Chain, Dense, softmax, params, train!
+using Flux: crossentropy, ADAM, freezelayers!
+using Test
+
+function collect_params(m, layer_indexes::Vector{Int64})
+    ps = []
+    for layer_index in layer_indexes
+        layer = m.layers[layer_index]
+        param_names = fieldnames(typeof(layer))
+        for param_name in param_names
+            p = deepcopy(getfield(layer, param_name))
+            push!(ps, p)
+        end
+    end
+
+    return ps
+end
+
+@testset "FreezeLayers" begin
+    # Fixed layers
+    freezed_layer_indexes = [1, 3]
+    mutable_layer_indexes = [2, 4]
+
+    # Model
+    m = Chain(
+        Dense(4, 2),
+        Dense(2, 3),
+        Dense(3, 4),
+        Dense(4, 3),
+        softmax
+    )
+
+    # Original params
+    ps_freezed = collect_params(m, freezed_layer_indexes)
+    ps_mutable = collect_params(m, mutable_layer_indexes)
+
+    # Update params
+    train!(
+        (x, y) -> crossentropy(m(x), y),
+        freezelayers!(params(m), m, freezed_layer_indexes),
+        [(rand(4, 10), rand(1, 10))],
+        ADAM(0.005)
+    )
+
+    # Params after update
+    new_ps_freezed = collect_params(m, freezed_layer_indexes)
+    new_ps_mutable = collect_params(m, mutable_layer_indexes)
+
+    @test new_ps_freezed == ps_freezed
+    @test new_ps_mutable != ps_mutable
+end

test/runtests.jl

@@ -1,6 +1,6 @@
-using Flux 
+using Flux
 using Flux.Data
-using Test 
+using Test
 using Random, Statistics, LinearAlgebra
 using Documenter
 using IterTools: ncycle
@@ -21,6 +21,10 @@ Random.seed!(0)
     include("optimise.jl")
   end
 
+  @testset "FreezeLayers" begin
+    include("freezelayers.jl")
+  end
+
   @testset "Data" begin
     include("data.jl")
   end