Advanced Model Building and Customisation
Here we will try and describe usage of some more advanced features that Flux provides to give more control over model building.
Customising Parameter Collection for a Model
Taking reference from our example Affine
layer from the basics.
By default all the fields in the Affine
type are collected as its parameters, however, in some cases it may be desired to hold other metadata in our "layers" that may not be needed for training, and are hence supposed to be ignored while the parameters are collected. With Flux, it is possible to mark the fields of our layers that are trainable in two ways.
The first way of achieving this is through overloading the trainable
function.
julia> @functor Affine
julia> a = Affine(rand(3,3), rand(3))
Affine{Array{Float64,2},Array{Float64,1}}([0.66722 0.774872 0.249809; 0.843321 0.403843 0.429232; 0.683525 0.662455 0.065297], [0.42394, 0.0170927, 0.544955])
julia> Flux.params(a) # default behavior
Params([[0.66722 0.774872 0.249809; 0.843321 0.403843 0.429232; 0.683525 0.662455 0.065297], [0.42394, 0.0170927, 0.544955]])
julia> Flux.trainable(a::Affine) = (a.W,)
julia> Flux.params(a)
Params([[0.66722 0.774872 0.249809; 0.843321 0.403843 0.429232; 0.683525 0.662455 0.065297]])
Only the fields returned by trainable
will be collected as trainable parameters of the layer when calling Flux.params
.
Another way of achieving this is through the @functor
macro directly. Here, we can mark the fields we are interested in by grouping them in the second argument:
Flux.@functor Affine (W,)
However, doing this requires the struct
to have a corresponding constructor that accepts those parameters.
Freezing Layer Parameters
When it is desired to not include all the model parameters (for e.g. transfer learning), we can simply not pass in those layers into our call to params
.
Consider a simple multi-layer perceptron model where we want to avoid optimising the first two Dense
layers. We can obtain this using the slicing features Chain
provides:
m = Chain(
Dense(784, 64, relu),
Dense(64, 64, relu),
Dense(32, 10)
)
ps = Flux.params(m[3:end])
The Zygote.Params
object ps
now holds a reference to only the parameters of the layers passed to it.
During training, the gradients will only be computed for (and applied to) the last Dense
layer, therefore only that would have its parameters changed.
Flux.params
also takes multiple inputs to make it easy to collect parameters from heterogenous models with a single call. A simple demonstration would be if we wanted to omit optimising the second Dense
layer in the previous example. It would look something like this:
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
, by simply deleting it from ps
:
ps = params(m)
delete!(ps, m[2].b)
Custom multiple input or output layer
Sometimes a model needs to receive several separate inputs at once or produce several separate outputs at once. In other words, there multiple paths within this high-level layer, each processing a different input or producing a different output. A simple example of this in machine learning literature is the inception module.
Naively, we could have a struct that stores the weights of along each path and implement the joining/splitting in the forward pass function. But that would mean a new struct any time the operations along each path changes. Instead, this guide will show you how to construct a high-level layer (like Chain
) that is made of multiple sub-layers for each path.
Multiple inputs: a custom Join
layer
Our custom Join
layer will accept multiple inputs at once, pass each input through a separate path, then combine the results together. Note that this layer can already be constructed using Parallel
, but we will first walk through how do this manually.
We start by defining a new struct, Join
, that stores the different paths and a combine operation as its fields.
using Flux
using CUDA
# custom join layer
struct Join{T, F}
combine::F
paths::T
end
# allow Join(op, m1, m2, ...) as a constructor
Join(combine, paths...) = Join(combine, paths)
Notice that we parameterized the type of the paths
field. This is necessary for fast Julia code; in general, T
might be a Tuple
or Vector
, but we don't need to pay attention to what it specifically is. The same goes for the combine
field.
The next step is to use Flux.@functor
to make our struct behave like a Flux layer. This is important so that calling params
on a Join
returns the underlying weight arrays on each path.
Flux.@functor Join
Finally, we define the forward pass. For Join
, this means applying each path
in paths
to each input array, then using combine
to merge the results.
(m::Join)(xs::Tuple) = m.combine(map((f, x) -> f(x), m.paths, xs)...)
(m::Join)(xs...) = m(xs)
Lastly, we can test our new layer. Thanks to the proper abstractions in Julia, our layer works on GPU arrays out of the box!
model = Chain(
Join(vcat,
Chain(Dense(1, 5),Dense(5, 1)), # branch 1
Dense(1, 2), # branch 2
Dense(1, 1)), # branch 3
Dense(4, 1)
) |> gpu
xs = map(gpu, (rand(1), rand(1), rand(1)))
model(xs)
# returns a single float vector with one value
Using Parallel
Flux already provides Parallel
that can offer the same functionality. In this case, Join
is going to just be syntactic sugar for Parallel
.
Join(combine, paths) = Parallel(combine, paths)
Join(combine, paths...) = Join(combine, paths)
# use vararg/tuple version of Parallel forward pass
model = Chain(
Join(vcat,
Chain(Dense(1, 5),Dense(5, 1)),
Dense(1, 2),
Dense(1, 1)
),
Dense(4, 1)
) |> gpu
xs = map(gpu, (rand(1), rand(1), rand(1)))
model(xs)
# returns a single float vector with one value
Multiple outputs: a custom Split
layer
Our custom Split
layer will accept a single input, then pass the input through a separate path to produce multiple outputs.
We start by following the same steps as the Join
layer: define a struct, use Flux.@functor
, and define the forward pass.
using Flux
using CUDA
# custom split layer
struct Split{T}
paths::T
end
Split(paths...) = Split(paths)
Flux.@functor Split
(m::Split)(x::AbstractArray) = tuple(map(f -> f(x), m.paths))
Now we can test to see that our Split
does indeed produce multiple outputs.
model = Chain(
Dense(10, 5),
Split(Dense(5, 1),Dense(5, 3),Dense(5, 2))
) |> gpu
model(gpu(rand(10)))
# returns a tuple with three float vectors
A custom loss function for the multiple outputs may look like this:
using Statistics
# assuming model returns the output of a Split
# x is a single input
# ys is a tuple of outputs
function loss(x, ys, model)
# rms over all the mse
ŷs = model(x)
return sqrt(mean(Flux.mse(y, ŷ) for (y, ŷ) in zip(ys, ŷs)))
end