Mlj.jl: Saving and loading models

[banshee-44]

I've been playing with MLJ the last couple of days and I really like it! It's solving every issue I ever had with sklearn and having to write endless boilerplate code in Python! :)

Only issue I've encountered is how to save and reload models: using BSON.jl I can successfully save and load a machine, but this obviously takes the data it was trained with along with it.

I can manually set the args field in the machine struct and the fitresult substruct to (nothing,), save, load, predict and get the correct results, but this feels like a bit of a hack.

Is there anything planned for model saving and loading?
In the meantime, is there a better way you can suggest I go about saving and loading a model?

[ablaom]

Thanks for the feedback, which is much appreciated.

Funny you should mention model saving/loading. This was the discussion in our last meeting. Yes, it's on the list.

I suppose there are two attributes to a machine worth loading/saving: the model (container for hyperpamaters) and the fitresult. We could have a universal way of serializing a model (eg, BSON.jl) but allow algorithms with custom fitresult formats to overload some fallback methods save and load we define in MLJBase (dispatched on models, as there are no machines in MLJBase).

How should this look at the user-interface? I don't see a need for the user to save just a model. They could do something like save(mach, fitted=true) to save a machine with or without the fitresult.

@banshee-44 What reasons do you choose BSON.jl for serialisation? Can you say if it compares favorably with other options, eg, JLD.jl?

[banshee-44]

Ha, talk about good timing then!

Oh I misunderstood the names, I guess I am looking to save/load the fitresult in this case then.

I chose BSON over JLD for a number of reasons:

• It's pretty much dependency free, notably it doesn't have the dependency on HDFS. This is super useful because it dramatically simplifies container/web server dependencies, setup complexity and setup time.
• it's super straightforward to use. (The @save and @load macros are super neat)
• I'd previously used it when using Flux
• it also doesn't have the security issue that Python Pickle, JLD, etc formats all have, in which loading the file can cause loading of arbitrary code.

My reason for saving just the model (I guess I should say just the fitresult?) was so that I could do the training on my local machine/server instance, then save the best performing one to cloud storage (like AWS S3) so it could be downloaded by containers/servers and used to serve predictions in a web service. As such, I don't need the hyper-parameters anymore, and I don't want to ship the training data (in the :args fields) to a web server if I've trained on several GB or more data or if the data was sensitive.

[darenasc]

Nice @banshee-44, I'd guess you are deploying a trained model from MLJ to a different (production maybe) environment? That's a neat use case!

I'd guess both types of saving would be needed, saving models for future tuning or re training for example, and deploying trained fitresult for inference or scoring of new data.

[banshee-44]

I am indeed! :)

I'd guess both types of saving would be needed, saving models for future tuning or re training for example

Definitely, Flux has a really nice approach here, where the training function let's you provide a callback that can execute every n seconds so you can do things like log metrics to console (evaluate! does that for you already) or save models and training state to disk so you can resume later! https://fluxml.ai/Flux.jl/stable/saving/

[ablaom]

@banshee-44 Thanks for your critique of BSON. That is very helpful.

I wonder if you be interested in implementing save/load for MLJ machines? I could provide guidance on the existing API?

[ablaom]

@banshee-44 Quick question: Is it safe in BSON to save and load structs that have undefined fields? In particular, is undefined(SomeStruct, :some_field) == false preserved on reloading?

[banshee-44]

@ablaom I'd love to give it a go!

As for the BSON question, I'll test and get back to you!

[banshee-44]

Just tested the BSON undefined loading. It is indeed safe, provided the Struct you're loading exists in the environment. Here's a minimal example demonstrating everything:

using BSON
struct TestStruct
    a::String
    b
end

s1 = TestStruct("some string", nothing)
isdefined(s1, :c) # false
BSON.@save "test_struct.bson" s1

Now, from within the same environment:

mystruct = BSON.load("test_struct.bson")
isdefined(mystruct[:s1], :c) # false

everything works as expected :tada:

Same approach from a new environment:

using BSON
mystruct = BSON.load("test_struct.bson")
\# Julia throws UndefVarError: TestStruct not defined

However, if we have our struct defined in a module like so:

module MyModule
struct TestStruct
    a::String
    b
end
end

and then we instantiate and save the struct as you would normally:

using BSON,
using MyModule
s1 = TestStruct("hello world", 1)
isdefined(s1, :c) # false
BSON.@save "test_struct.bson" s1

Now, in a new environment:

using BSON
using MyModule

mystruct = BSON.load("test_struct.bson")
# works as expected
typeof(mystruct)
# Dict{Symbol,Any} with 1 entry:
# :s1 => TestStruct("hello world", 1)
isdefined(mystruct[:s1], :c) # false, as expected

BSON also has a neat @load macro, where it inserts things directly into the environment, so we could also do this and get the same result:

using BSON
using MyModule

BSON.@load "test_struct.bson" s1
# s1 is inserted into the environment
typeof(s1)
# MyModule.TestStruct("hello world", 1)
isdefined(s1, :c) # false, as expected

Edit: fixing formatting

[ablaom]

Thanks for that. These observations are helpful but may or may not have a bearing on the implementation of machine saving/loading we go for.

I guess we do something like this: We havesave("mach.bson", mach::Machine) (or a macro call) BSON-save mach.model and mach.fitresult to file (the first represents hyperparameters, the second the fitted parameters). We could also save a string representation of the model type (I think str = string(MLJBase.coretype(typeof(model))) will do the trick) so that we can pre-load necessary code when restoring the saved machine.

To restore a machine from file we arrange for load("mach.bson") to: (i) read from file the string representation str of the model type; (ii) use MLJ's existing @load str to load the necessary code (if the code already exists, this is a no-op); (iii) read model and fitresult from the file and instantiate a machine with those fields only defined.

We will need the fit! method for machines to throw an error if the machine args are not defined (which was previously not a possibility) so users who try to fit! a restored machine know why that isn't going to work and suggest a remedy (perhaps calling machine(filename, args...) loads a machine from file and binds new data to the machine).

One annoyance is that @load is already used to load model code. I'd hate to make a breaking change (eg, @loadcode) but you could twist my arm.

@banshee-44 Hows does this sound?

[ablaom]

@banshee-44 Some corrections and clarifications:

• we should save the model name and package name for code reloading. Rather than using my suggestion above, just use info(model)[:name] and info(model)[:package_name] where model = mach.model and mach::Machine is the machine you are saving. We need both because different packages are allowed to have models of the same name.

• for reloading code the syntax is then @load PCA pkg=MultivariateStats . However, @load currently loads the code into the module that makes the call. So your @loadmachine should instead, use _load method (see MLJ/src/loading.jl) where you can specify the module (specifically, the __module__ variable local to your @loadmachine macro). If all this unclear, leave out the code-loading bit and @giordano or I can add later, no sweat.

Thanks @giordano for the _load function.

[banshee-44]

Ok, I'll give it a go, earliest I can probably try is probably this weekend.

What's your preferred method of contributing: should I make a fork or a branch and make pull requests off that?

[ablaom]

Awesome. PR off a fork is good.

[oxinabox]

it also doesn't have the security issue that Python Pickle, JLD, etc formats all have, in which loading the file can cause loading of arbitrary code.

This is not true.
The pure BSON format itself does not cause the loading of arbitrary code.
BSON.jl includes extensions that allow it to serialize arbitrary julia objects, including methods.
I am almost certain it is possible to craft a BSON.jl file that can execute arbitrary code.
It is definitely possible to make one that calls arbitrary constructors (this is the same as JLD.jl)
Security and safety was not a design concern for BSON.jl
Don't load BSON files from untrusted sources.

There was a discussion in #internal the other day about this.
And a proposal to split BSON into two modes: the safe pure BSON stuff that can only deal with JSON-able types + Arrays + Datetimes, and the unsafe extensions that allow you to work with arbitrary julia types.

[ablaom]

I'm revisiting this issue after it has sat dormant for a while. After some reflection, and to make this a more attractive option for a volunteer contribution, I propose we keep this simple, and not worry about auto-loading code or anything. We use BSON.jl to get this behaviour:

• save(file, machine::Machine) saves a machine's model, fitresult and report fields to a BSON file with name file

• machine(file, args...) instantiate a Machine using the BSON file with name file and binds it to optionally supplied data args.... Currently the constructor machine(model, args..) complains if length(args) == 0, in the argument checking logic, so this should be relaxed. There is a constructor without the checks, called MLJBase.Machine(model, args...) but this should not be used in case data is actually supplied. However, for a proof of concept, I use it below:

julia> X, y = @load_boston;
julia> mach = fit!(machine(model, X, y))
julia> mach2 = MLJBase.Machine(model)
julia> mach2.fitresult = mach.fitresult
julia> predict(mach2, selectrows(X, 1:3))
3-element Array{Distributions.Normal{Float64},1}:
 Distributions.Normal{Float64}(μ=22.532806324110698, σ=9.188011545278206)
 Distributions.Normal{Float64}(μ=22.532806324110698, σ=9.188011545278206)
 Distributions.Normal{Float64}(μ=22.532806324110698, σ=9.188011545278206)

[ablaom]

Note that the relevant code should now go in MLJBase/src/machines.jl

[oxinabox]

Update re:

it also doesn't have the security issue that Python Pickle, JLD, etc formats all have, in which loading the file can cause loading of arbitrary code.

This is not true.
The pure BSON format itself does not cause the loading of arbitrary code.
BSON.jl includes extensions that allow it to serialize arbitrary julia objects, including methods.
I am almost certain it is possible to craft a BSON.jl file that can execute arbitrary code.

Examples of abritary code execution have been created for BSON, JLD2 and the Julia Serializer (which means also that they definately exist for JLSO since it is a container that holds BSON / the Julia Serializer)

[ablaom]

@oxinabox I guess the long term plan would be to add a method for models to implement that provides a possibility for saving/loading using a model-specific format, and we just use BSON as the fallback.

Thoughts?

[oxinabox]

Seems legit.
You may want to give JLSO a serious look. (Says one of the authors of JLSO.)
JLSO is not a serialisation format, it is a serialized data container.
It adds metadata (Project/Manifest of package versions etc) + compression, in a wrapper around serialized data (BSON.jl or the builtin serializer)