FluxML / FluxML/Functors.jl

Taggable functors

Open
#41 4 comments 1 reaction 0 assignees View on GitHub
Dominant language
Julia
Stars
123
Forks
18
Avg merge
47m
Merged PRs (30d)
1

Description

In Flux, we have `trainable` to designate a subset of leaves as nodes to walk when updating parameters for training. In FluxPrune.jl, I defined `pruneable` to designate a subset of leaves for pruning (note that these cannot be the same as the `trainable` nodes).

Right now this creates an unfortunate circumstance as discussed in https://github.com/FluxML/Flux.jl/issues/1946. Users need to `@functor` their types, remember to define `trainable` if necessary. Potentially, to use FluxPrune.jl, they might want to remember to define `pruneable`. On the developer side of things, we can use the `walk` keyword of `fmap` to walk the differently labeled leaf nodes. But this usually requires defining a separate walk function based on the subset that you are hoping to target.

An alternative would be to build this information directly into what `@functor` defines. Right now, each child of a functor has a name and a value. I propose adding "tags" which would be a tuple of symbols. Then we could do something like
```julia
@functor Conv trainable=(weight, bias) pruneable=(weight,)
```
Ideally, this mechanism should be dynamic, meaning that if Flux.jl already defines the trainable leaves of a type, then another package like FluxPrune.jl should be able to add a pruneable tag on top of that.

My hope is that we make it easier on users by only having one line for making your type Flux-compatible. And we make it easier on developers by making it easy to filter nodes when walking by tag. I haven't spent a lot of time on the implementation aspect, but I just wanted to float the notion of tags first and get some feedback.

Contributor guide

No contributing guide indexed for this repository

Research direction

Start by reading the referenced Flux.jl issue and the existing @functor and fmap walk behavior described here. Compare how trainable and pruneable are currently defined in Flux and FluxPrune.jl. Done would require an agreed design for tagged children and dynamic tag extension, followed by an implementation plan and tests; this issue does not yet identify files or tests.

Written by the indexing model from the issue text.

Assessment

Tech stack
julia
Domain
machine-learning
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.