Supported Machine Learning Libraries

This library currently supports:

The common pattern, across all supported libraries, is to:

# Create an instance of the model store
from modelstore import ModelStore

ms = ModelStore.from_gcloud(
   project_name="my-project",
   bucket_name="my-bucket",
)

# Export your model by calling `create_archive()`
archive = ms.<library-name>.create_archive(**kwargs)

# Upload your model by calling `upload()`
meta_data = ms.upload("model-domain", archive)

CatBoost

To export a CatBoost model, use:

# Train your model
model = ctb.CatBoostClassifier(loss_function="MultiClass")
model.fit(x, y)

# Create an archive
archive = ms.catboost.create_archive(model=clf, pool=df)

This will add a multiple formats of your model to the archive:

  • CatBoost binary format

  • JSON

  • Onnx

The pool argument is required if you are training a multi class model.

The archive will also contain a model_attributes.json file with all of the attributes of the model.

Keras

To export a Keras model, use:

# Train your model
model = keras.Model(inputs, outputs)
model.compile(optimizer="adam", loss="mean_squared_error")
model.fit(X_train, y_train, epochs=10)
# ...

# Create and upload an archive
archive = ms.keras.create_archive(model=net, optimizer=optim)
rsp = ms.upload("model-domain", archive)

This will add two dumps of the model into the archive; based on calling model.to_json() and model.save().

PyTorch

To export a PyTorch model, use:

# Train your model
net = ExampleNet()
optim = ExampleOptim()
# ...

# Create and upload an archive
archive = ms.pytorch.create_archive(model=net, optimizer=optim)
rsp = ms.upload("model-domain", archive)

This will add two dumps of the model into the archive; a checkpoint.pt that contains the net and optimizer’s state (e.g., to continue training at a later date), and a model.pt that is the result of torch.save with the model only (e.g., for inference).

Scikit-Learn

To export a scikit-learn model, use:

# Train your model
clf = RandomForestClassifier(n_estimators=10)
clf = clf.fit(X, Y)

# Create an archive
archive = ms.sklearn.create_archive(model=clf)

This will add a joblib dump of the model into the archive.

Tensorflow

To export a tensorflow model, use:

# Train your model
model = tf.keras.models.Sequential(
    [
        tf.keras.layers.Dense(5, activation="relu", input_shape=(10,)),
        tf.keras.layers.Dropout(0.2),
        tf.keras.layers.Dense(1),
    ]
)
model.compile(optimizer="adam", loss="mean_squared_error")
model.fit(X_train, y_train, epochs=10)

# Create an archive
archive = model_store.tensorflow.create_archive(model=model)

This will both save the weights (as a checkpoint file) and export/save the entire model.

Transformers

To export a transformers model, use:

# Get a pre-trained model and fine tune it
model_name = "distilbert-base-cased"
config = AutoConfig.from_pretrained(
    model_name, num_labels=2, finetuning_task="mnli",
)
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(
    model_name, config=config,
)

# Create an archive
archive = model_store.transformers.create_archive(
    config=config, model=model, tokenizer=tokenizer,
)

The config and tokenizer parameters are optional. This will use the save_pretrained() function to save your model.

XGBoost

To export an XGBoost model, use:

# Train your model
bst = xgb.train(param, dtrain, num_round)

# Create and upload an archive
archive = ms.xgboost.create_archive(model=bst)
rsp = ms.upload("model-domain", archive)

This will add two dumps of the model into the archive; a model dump (in an interchangeable format, for loading again later), and a model save (in JSON format, which, to date, is experimental).