Compatibility with ProtoStruct.jl, and LayerFactory ideas for custom layers

[MilesCranmer]

There's this really nice package ProtoStruct.jl that lets you create structs which can be revised. I think this is extremely useful for developing custom models in Flux.jl using Revise.jl, since otherwise I would need to restart every time I want to add a new property in my model.

Essentially the way it works is to transform:

@proto struct MyLayer
    chain1::Chain
    chain2::Chain
end

into (regardless of the properties)

struct MyLayer{NT<:NamedTuple}
    properties::NT
end

and, inside the macro, set up constructors based on your current defined properties.

However, right now it doesn't work with Flux.jl. When I try to get the parameters from a model, I see the error: NamedTuple has no field properties. Here is a MWE:

# Desired API for FluxML
using Flux
using Flux: params, @functor
using ProtoStructs

@proto struct ResidualDense
    w1::Dense
    w2::Dense
    act::Function
end

"""Residual layer."""
function (r::ResidualDense)(x)
    dx = r.w2(r.act(r.w1(x)))
    return r.act(dx + x)
end

@functor ResidualDense

function ResidualDense(in, out; hidden=128, act=relu)
    ResidualDense(Dense(in, hidden), Dense(hidden, out), act)
end

# Chain of linear layers:
mlp = Chain(
    Dense(5 => 128),
    ResidualDense(128, 128),
    ResidualDense(128, 128),
    Dense(128 => 1),
);

p = params(mlp);  # Fails

and here is the error:

ERROR: type NamedTuple has no field properties
Stacktrace:
  [1] getproperty
    @ ./Base.jl:38 [inlined]
  [2] getproperty(o::ResidualDense{NamedTuple{(:w1, :w2, :act), Tuple{Dense, Dense, Function}}}, s::Symbol)
    @ Main ~/.julia/packages/ProtoStructs/4sIVY/src/ProtoStruct.jl:134
  [3] functor(#unused#::Type{ResidualDense{NamedTuple{(:w1, :w2, :act), Tuple{Dense, Dense, Function}}}}, x::ResidualDense{NamedTuple{(:w1, :w2, :act), Tuple{Dense, Dense, Function}}})
    @ Main ~/.julia/packages/Functors/V2McK/src/functor.jl:19
  [4] functor(x::ResidualDense{NamedTuple{(:w1, :w2, :act), Tuple{Dense, Dense, Function}}})
    @ Functors ~/.julia/packages/Functors/V2McK/src/functor.jl:3
  [5] trainable(x::ResidualDense{NamedTuple{(:w1, :w2, :act), Tuple{Dense, Dense, Function}}})
    @ Optimisers ~/.julia/packages/Optimisers/GKFy2/src/interface.jl:153
  [6] params!
    @ ~/.julia/packages/Flux/nJ0IB/src/functor.jl:46 [inlined]
  [7] params!(p::Zygote.Params{Zygote.Buffer{Any, Vector{Any}}}, x::Tuple{Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, ResidualDense{NamedTuple{(:w1, :w2, :act), Tuple{Dense, Dense, Function}}}, ResidualDense{NamedTuple{(:w1, :w2, :act), Tuple{Dense, Dense, Function}}}, ResidualDense{NamedTuple{(:w1, :w2, :act), Tuple{Dense, Dense, Function}}}, ResidualDense{NamedTuple{(:w1, :w2, :act), Tuple{Dense, Dense, Function}}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}, seen::Base.IdSet{Any}) (repeats 3 times)
    @ Flux ~/.julia/packages/Flux/nJ0IB/src/functor.jl:47
  [8] params!
    @ ~/.julia/packages/Flux/nJ0IB/src/functor.jl:40 [inlined]
  [9] params(m::Chain{Tuple{Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, ResidualDense{NamedTuple{(:w1, :w2, :act), Tuple{Dense, Dense, Function}}}, ResidualDense{NamedTuple{(:w1, :w2, :act), Tuple{Dense, Dense, Function}}}, ResidualDense{NamedTuple{(:w1, :w2, :act), Tuple{Dense, Dense, Function}}}, ResidualDense{NamedTuple{(:w1, :w2, :act), Tuple{Dense, Dense, Function}}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}})
    @ Flux ~/.julia/packages/Flux/nJ0IB/src/functor.jl:87
 [10] top-level scope
    @ ~/desired_julia_api/nice_api.jl:37

How hard would it be to make this compatible? I think it would be extremely useful to be able to quickly revise model definitions!

---

(Sorry for the spam today, by the way)

[mcabbott]

This is because:

julia> fieldnames(ResidualDense)
(:properties,)

julia> propertynames(ResidualDense)
(:var, :body)

and Functors isn't careful: it checks fieldnames but calls getproperty:

https://github.com/FluxML/Functors.jl/blob/v0.3.0/src/functor.jl#L11-L16

One fix is:

https://github.com/FluxML/Flux.jl/pull/1932/files#diff-34ea0dd342eeb0864015f5ab93bd9db67f6f46ec46d4bf8dbebe9b8c1a433b04R131-R132

Another fix would be for ProtoStructs to make getproperty work for the extra (internal) name. It could have a weird, possibly generated, name to avoid clashes.

[MilesCranmer]

Nice! Thanks for figuring this out.

Aside: I wonder if something like @proto should be used by default for Flux.jl custom layers? It seems like such a common use-case to modify structs when you are developing neural nets. I suppose it wouldn't hurt performance either.

Here's a terrible idea I had, just to see what if it spurs any useful ideas for others:

@layer struct ResBlock
    w1::Dense
    act::Function
    forward = (self, x, y) -> begin
        self.act(self.w1(x)))
    end
end

where @layer would basically do both what @proto and @functor (as well as @kwdef, for the forward) currently do. The forward property of ResBlock here would let you embed the forward function actually inside the struct declaration. I guess it's kind of against Julia style, but it seems intuitive for quickly building deep neural nets.

So, from this, the macro would generate the following code:

struct ResBlock{NT<:NamedTuple}
    properties::NT
end

function getproperty
    ...
end
# etc., etc.

function (self::ResBlock)(x, y)
    self.forward(self, x, y)
end

[MilesCranmer]

(Maybe I am thinking of trying this out: https://github.com/Suzhou-Tongyuan/ObjectOriented.jl - it's definitely useful for DL-type models)

[darsnack]

I don't know if we want that as default @layer, because eventually you want to remove the @proto specification (this is explicitly stated even in the ProtoStructs.jl readme).

Supporting ProtoStruct layers is good, because sometimes writing your own layer is unavoidable. But Flux has some features that make using different types (classes) for each sub-layer not the norm. For example, Metalhead.jl, the vision model library, can build almost all types of residual networks without defining a ResBlock type. This is because of two things in Flux that most frameworks lack:

• a built-in Parallel layer (and special case SkipConnection) which encompasses most branching network architectures
• the ability to name layers in a Chain or Parallel and access sub-layers via those names (so that you don't just see an ocean of Parallel)

(Btw this is not to shoot anything down; if what we have doesn't really do what you want, then we want to know. Love the enthusiasm!)

[ToucheSir]

To build on Kyle's point: Flux's philosophy towards layer types is very much "come as you are". We'd like to make it easier to define/register existing structs as layers and to do so without depending on Flux itself. So while we should absolutely try to support libraries like ProtoStruct and ObjectOriented.jl if we can, we also want to keep the barrier of entry for working with the module system as low as possible (if you know how to define a callable struct, you can make a layer).

[MilesCranmer]

Thanks for all the comments on my admittedly half-thought-out ideas! Did not know about Parallel - nice!

I don't know if we want that as default @layer, because eventually you want to remove the @proto specification (this is explicitly stated even in the ProtoStructs.jl readme).

Good point, agreed!

I wonder if you would want to mention Revise.jl+ProtoStructs.jl on the custom layer page (or even @reexport some of the key functionality, so you could do Flux.revise() at the top of the REPL to turn it on), since it seems almost required when developing neural nets; otherwise the startup time really hurts when working on complex networks.

For a lot of users, my sense is that Flux.jl may be the very first Julia package they try out.

We also want to keep the barrier of entry for working with the module system as low as possible (if you know how to define a callable struct, you can make a layer).

The compatibility makes a lot of sense. I am just trying to think if there's any way to simplify the current custom layer declaration method for new users. Right now you need to call struct MyStruct, function (m::MyStruct)(...), and @functor, and then also construct the components of the layer separately. This feels unwieldy.

Not only is it four separate calls, it's four different types of calls. I need to create (1) a stuct, then (2) a method, (3) call a macro, then separately (4) construct the layer's components. And I need to do that for every single custom layer. In Haiku and PyTorch, you would only create a (1) class and (2) method – sometimes you can even create a custom layer with just a method. Using four different ideas to make a simple NN layer just seems a bit heavy, and the self methods in Julia (like function (m::MyType)(x)) are admittedly ugly, and seem more appropriate for use by package developers rather than end users. That functionality always seemed more suitable for designing convenience methods in a package, rather than working on user objects.

Even something like

struct MyLayer
    x::Dense
    y::Dense
end

@layermethod (m::MyLayer)(z) -> relu(m.x(z) + m.y(z))

might make for a slightly cleaner convenience API. Though the standard methods would of course also be available.

Thoughts? I am aware I could simply be too comfortable with the Python DL ecosystem, though, and this isn't Julia-esque enough. No worries if that is the case.

I think my dream layer creation method would really be something like

@layerfactory function my_layer(n_in, n_out)
    w1 = Dense(n_in, 128)
    w2 = Dense(128, n_out)
    return (x) -> begin
        y = relu(w1(x)) + x
        w2(y)
    end
end

which would let you construct the layer struct, the layer's components, and the forward pass, all in one go. (But that would obviously be tricky to implement).

Best,
Miles

[mcabbott]

create (1) a stuct, then (2) a method, (3) call a macro, then separately (4) construct the layer's components.

What's nice about 1,2,4 is that there is nothing Flux-specific about them. They are completely ordinary Julia code.

Making a special DSL isn't impossible, but it's one more thing you have to learn, and it will have limitations. This is a little bit like the train! story, where saving a few lines (compared to writing out the ordinary Julia code of the for loop) comes at the cost of weird API to memorise, and limitations so that you later have to know the other way too.

Most things for which it's worth defining a new layer will want something extra. If I'm reading correctly the example here is an easy combination of existing layers:

Chain(SkipConnection(Dense(n_in => 128, relu), +), Dense(128 => nout))

Some of us would like to remove @functor and make it just recurse by default. Then perhaps @layer would be an optional way to add fancy printing (or to customise which bits are trainable) but not needed for the simplest layers.

I agree we should mention Revise.jl & maybe ProtoStructs.jl. On the ecosystem page for sure. Maybe ProtoStructs.jl ought to be on the advanced layer building page too? (Which could use some cleaning up.)

[MilesCranmer]

Thanks, I see your point. I just wonder if there's a way to still use ordinary Julia code but require less manual labor to build custom layers, at the expense of being slightly less generic. Recursing by default would definitely be an improvement!

If I'm reading correctly the example here is an easy combination of existing layers:

This was just a MWE; in practice there would probably not be an existing layer for my applications.

In some ways the layers like Chain/SkipConnection/Dense already forms a DSL. I guess what I am looking for a simpler way to extend this DSL to new layers so I can quickly prototype models for my research. With certain assumptions about what would go into the layer (e.g., similar assumptions to Torch/Flax methods), I think a custom layer factory could be written in a condensed form. Currently the way to make custom layers is 100% generic, and as a result is a bit more code - maybe there could also be a partially-generic method.

[MilesCranmer]

Here's one idea for a layer factory:

struct LayerFactory{F<:Function,NT<:NamedTuple}
    forward::F
    layers::NT
end

LayerFactory(f; layers...) = LayerFactory(f, NamedTuple(layers))

function (f::LayerFactory)(args...)
    return f.forward(f.layers, args...)
end

@functor LayerFactory

This makes it super easy to construct custom layers. Watch this:

my_layer = LayerFactory(; w1=Dense(5, 128), act=relu) do self, x
    self.act(self.w1(x))
end

That's literally all you need! I can construct custom layers in one line, without even changing the underlying structure of Flux.jl. And it works for training and everything.

What do you think about this?

[mcabbott]

This should work fine. What's a nice example of a nontrivial use?

[MilesCranmer]

The use-case I had in mind is graph networks, where you have a set of (nodes, edges, globals) that you need to do scatter operations on - it seems tricky to get that working with a Chain, but it should be doable to set it up with custom layers.

I am really happy about this LayerFactory struct. I'm actually surprised that it just works with this package. You can even change the number of fields in the same runtime, and it still works! Would you be willing to include it in Flux.jl as a simple way to construct custom layers?

e.g., here's another example:

model = LayerFactory(;
    w1=Dense(1, 128), w2=Dense(128, 128), w3=Dense(128, 1), act=relu
) do self, x
    x = self.act(self.w1(x))
    x = self.act(self.w2(x))
    self.w3(x)
end

p = params(model)  # works!

[MilesCranmer]

Here's a simple implementation of a graph network in PyTorch: https://pytorch-geometric.readthedocs.io/en/latest/modules/nn.html#torch_geometric.nn.meta.MetaLayer

[MilesCranmer]

It even works for compositions of LayerFactory!

function MLP(n_in, n_out, nlayers)
    LayerFactory(;
        w1=Dense(n_in, 128), w2=[Dense(128, 128) for i=1:nlayers], w3=Dense(128, n_out), act=relu
    ) do self, x
        embed = self.act(self.w1(x))
        for w in self.w2
            embed = self.act(w(embed))
        end
        self.w3(embed)
    end
end

model = LayerFactory(; mlp1=MLP(1, 128, 2), mlp2=MLP(128, 1, 3)) do self, x
    self.mlp2(self.mlp1(x))
end

[darsnack]

I am not super familiar with GNNs, but you might want to check out GraphNeuralNetworks.jl to see how they handle working with Flux. They do seem to have a custom GNNChain.

Okay, I think I understand what you are saying. If you have a sufficiently complex forward function that involves sub-layers, then writing it from scratch with "base" Julia + Flux is a bunch of text. As would be the case with "base" PyTorch or Jax, but those libraries have utilities built on top like your LayerFactory. So is it right that you are looking for the same in Flux?

While I have no problem with LayerFactory since it is a nice convenience utility, I want to note that if we decide to auto-@functor structs, then LayerFactory comes for free from plain Julia:

mlp(n_in, n_out, nlayers) = let w1 = Dense(n_in, 128), w2 = [Dense(128, 128) for i in 1:nlayers], w3 = Dense(128, n_out)
    return function(x)
        act = relu

        embed = act(w1(x))
        for w in w2
            embed = act(w(embed))
        end
        w3(embed)
    end
end

model = let mlp1 = mlp(1, 128, 2), mlp2 = mlp(128, 1, 3)
    x -> mlp2(mlp1(x))
end

p = params(model) # works too!

---

Below is just for your reference.

Looking at the link you shared, this is what I would write in Flux:

# this is one way avoiding structs completely
EdgeModel(edge_mlp = Chain(...)) = Chain(
    (src, dest, edge_attr, u, batch) -> vcat(src, dest, edge_attr, u[batch]),
    edge_mlp
)

# admittedly, structs seems nice here
Base.@kwdef struct NodeModel{T, S}
    node_mlp_1::T = Chain(...)
    node_mlp_2::S = Chain(...)
end

@functor NodeModel

function (m::NodeModel)((x, edge_index, edge_attr, u, batch))
    row, col = edge_index
    out = vcat(x[row], edge_attr)
    out = m.node_mlp_1(out)
     # not sure what this is doing but we have a NNlib.scatter
    out = scatter_mean(out, col, ...)
    out = vcat(x, out, u[batch])
    return m.node_mlp_2(out)
end

And so on. Modulo the @functor issue, I don't see how defining a class and forward function is shorter than what's above. Seem like just an extra end keyword and a separation between the struct definition and the forward definition. The models I defined above are ready to be passed into a Chain or Parallel (assuming that's what MetaLayer is).

[darsnack]

Or another way of putting it: Chain and friends are a sort-of a DSL for passing arguments between sub-layers. You have a need to define your own argument passing container layer, but you don't want to write all the extra stuff that goes along with a base layer like Conv. The Haiku link you shared shows a mechanism for constructing "layers" that have no type, but they do have fields, a forward pass, and a variable they are bound to. These three things in Julia are exactly what make an anonymous function! The only thing preventing this from just working^(TM) in Flux is that @functor is opt-in.