LegolasFlux
Note: Upgrading from LegolasFlux v0.1.x to v0.2?
The only change is an update to Legolas v0.5. Be sure to check out the guidance for updating Legolas to v0.5 along with the rest of Legolas's documentation and tour.
LegolasFlux provides some simple functionality to use Legolas.jl's
extensible Arrow schemas as means to serialize Flux models similarly to using Flux's params
and loadparams!
(instead, we export similar functions fetch_weights
and load_weights!
which handle layers like BatchNorm
correctly for this purpose).
The aim is to serialize only the numeric weights, not the code defining the model. This is a very different approach
from e.g. BSON.jl, and hopefully much more robust. Note that in this package, we use weights
to refer to the numeric arrays that are modified over the course of training a model; that includes biases as well as means and variances in e.g. BatchNorms (but not e.g. configuration settings).
With this approach, however, if you change the code such that the weights are no longer valid (e.g. add a layer), you will not be able to load back the same model.
Usage
using Flux
function make_my_model()
return Chain(Dense(1,10), Dense(10, 10), Dense(10, 1))
end
my_model = make_my_model()
# train it? that part is optional ;)
# Now, let's save it!
using LegolasFlux
model_row = LegolasFlux.ModelV1(; weights = fetch_weights(cpu(my_model)),
architecture_version=1)
write_model_row("my_model.model.arrow", model_row)
# Great! Later on, we want to re-load our model weights.
fresh_model = make_my_model()
model_row = read_model_row("my_model.model.arrow")
load_weights!(fresh_model, model_row.weights)
# Now our weights have been loaded back into `fresh_model`.
We can make use of the architecture_version
column to specify a version number for the architectures, in order
to keep track of for which architectures the weights are valid for.
See examples/digits.jl for a larger example, which also saves out extra metadata with the model, by using a Legolas schema extension.
LegolasFlux.ModelV1
A LegolasFlux.ModelV1
is a central object of LegolasFlux. It acts as a Tables.jl-compatible row that can store the weights
of a Flux model in the weights
column, optionally an architecture_version
(defaults to missing
).
ModelV1
is not exported because downstream models likely want to define their own rows which extend the schema provided by LegolasFlux
that might end up being called something similar. See the next section for more on extensibility.
Extensibility
As a Legolas.jl schema, it is meant to be extended. For example, let's say I had an MNIST classification model
that I call Digits
. I am very committed to reproducibility, so I store the commit_sha
of my model's repo
with every training run, and I also wish to save the accuracy and epoch. I might create a DigitsRow
which is
a schema extension of the legolas-flux.model
schema:
using Legolas, LegolasFlux
using Legolas: @schema, @version
@schema "digits-model" DigitsRow
@version DigitsRowV1 > ModelV1 begin
# re-declare this ModelV1 field as parametric for this schema as well
weights::(<:Union{Missing,Weights})
epoch::Union{Missing, Int}
accuracy::Union{Missing, Float32}
commit_sha::Union{Missing, String}
end
Now I can use a DigitsRowV1
much like LegolasFlux's ModelV1
. It has the same required weights
column and optional architecture_version
column, as well as the additional epoch
, accuracy
, and commit_sha
columns. As a naming convention,
one might name files produced by this row as e.g. training_run.digits.model.arrow
.
When writing out a DigitsRowV1
, I'll pass the schema version like so
write_model_row(path, my_digits_row, DigitsRowV1SchemaVersion())
so that later, when I call read_model_row
on this path, I'll get back a DigitsRowV1
instance.
Note in this example the schema is called digits.model
instead of just say digits
, since the package Digits might want to
create other Legolas schemas as well at some point.
Check out the Legolas.jl repo to see more about how its extensible schema system works, and the example at examples/digits.jl.