Supported Machine Learning Libraries

This library currently supports several different machine learning libraries. To save models trained with them, you should use the upload function:

model_store.upload("domain", <kwargs>)
Supported machine learning libraries

Library

Required kwargs

Example code

Annoy

model

Annoy Example

CatBoost

model, pool (for classification)

Catboost Example

FastAI

learner

FastAI Example

Gensim

model

Word2vec Example

Keras

model, optimizer

Keras Example

LightGBM

model

LightGBM Example

Mxnet

model, epoch

Mxnet Example

Onnx

model

Onnx Example

PyTorch

model, optimizer

PyTorch Example

PyTorch Lightning

model, trainer

PyTorch Lightning Example

scikit-learn

model

scikit-learn Example

skorch

model

skorch Example

shap

explainer

shap Example

Tensorflow

model

Tensorflow Example

Transformers

config, model, tokenizer

Transformers Example

XGBoost

model

XGBoost Example

What to do if a library is not supported

If you are using a machine learning library that is not listed above, you can still use model store to upload and version your models. You will not be able to use load() but you will be able to download() them back.

model_path = save_model()

model_store.upload("my-domain", model=model_path)

You can also:

Uploading more than one model file

There are some cases where you may want to upload two models together.

This library supports uploading multiple models, as long as their keyword arguments do not overlap.

For example, you might want to upload a classifier and a shap explainer together:

clf = RandomForestClassifier()
clf.fit(X_train, y_train)

explainer = shap.TreeExplainer(model)

model_store.upload("my-domain", model=model, explainer=explainer)

When you load these models, model store returns a dictionary with both models:

models = modelstore.load(model_domain, model_id)
clf = models["sklearn"]
explainer = models["shap"]

Uploading extra files with the model

There are some cases where you may want to upload more files alongside the model.

This library supports uploading a model with one or more extra files.

For example, you might want to upload a classifier and the predictions it made on the test set.

clf = RandomForestClassifier()
clf.fit(X_train, y_train)

predictions = clf.predict(X_test)
file_path = "predictions.csv"
numpy.savetxt(file_path, predictions, delimiter=",")

modelstore.upload("my-domain", model=model, extras=file_path)

When you load these models, the extra files are not loaded into memory:

clf = modelstore.load(model_domain, model_id)