Custom models

When wrapping a layer with Model, the layer struct's fieldnames are stored inside Model and are automatically classified into one of 3 categories:

model.paramattrs
Fields with trainable parameters. By default all fields of type AbstractArray{T,<:AbstractFloat} are treated as trainable parameters
model.childrenattrs Fields that contain children layers. By default fields that are a Flux layer are treated as children.
model.settingattrs Fields with hyperparameters that define the layer structure, e.g. padding on a convolution. Default field category.

While this should work well for even many custom layers, in some cases you have to define what category a field is yourself.

The classification happens with dispatch on the function attrtype.

ModelUtils.AttrType — Type

AttrType

Enum for layer field types.

ParamAttr
Field with trainable parameters, e.g. Dense.W
ChildAttr
Field that contains children layers, e.g. Chain.layers
SettingAttr
Fields with hyperparameters that define the layer structure, e.g. Conv.pad

source

ModelUtils.attrtype — Function

attrtype(layer::L, name::Val)::AttrType

Dispatch on a layer, and a field name (wrapped in a Val) to determine what AttrType the field is.

If there is no custom method defined, for a layer type and field name, attrtype will additionally dispatch on the field's value to give a reasonable default.

By default the layers will be classified as following:

ParamAttr
Fields of type AbstractArray{T,<:AbstractFloat} are treated as trainable parameters
ChildAttr
Fields that are a Flux layer are treated as children.
SettingAttr
Any other

To customize for your type, define a method like

attrtype(::MyLayer, ::Val{:myfieldname}) = ChildAttr

source

Defining a custom `attrtype` method

Let's say we create a custom layer that's like Chain:

using Flux, ModelUtils # hide
struct MyChain
    layers
    MyChain(layers...) = new(layers)
end
(m::MyChain)(x) = foldl((x, layer) -> layer(x), m.layers, init = x)

Now we want MyChain.layers to be a ChildAttr but we can check that it is not classified as such:

mychain = MyChain(Dense(10, 10), Dense(10, 10), softmax)
ModelUtils.attrtype(mychain, Val(:layers))

So we can define a custom attrtype method:

import ModelUtils: attrtype
ModelUtils.attrtype(::MyChain, ::Val{:layers}) = ModelUtils.ChildAttr

And now we get the result we wanted:

ModelUtils.attrtype(mychain, Val(:layers))

Custom models

Defining a custom attrtype method

Defining a custom `attrtype` method