How to manually register a sci-kit model with TRAINS python auto-magical experiment manager?

178 views Asked by At

I'm working mostly with scikit-learn, as far as I understand, the TRAINS auto-magic doesn't catch scikit-learn model store/load automatically.

How do I manually register the model after I have 'pickled' it.

For Example:

import pickle
with open("model.pkl", "wb") as file:  
    pickle.dump(my_model, file)
1

There are 1 answers

2
Martin.B On

Assuming you are referring to TRAINS experiment manager: https://github.com/allegroai/trains (which I'm one of the maintainers)

from trains import Task, OutputModel
OutputModel(Task.current_task()).update_weights(weights_filename="model.pkl")

Or, if you have information you want to store together with the pickled model file, you can do:

from trains import Task, OutputModel
model_parameters = {'threshold': 0.123}
OutputModel(Task.current_task(), config_dict=model_parameters).update_weights(weights_filename="model.pkl")

Now, you should see in the UI an output model registered with the experiment. The model contains a link to the pickel file, together with the configuration dictionary.