Skip to content

Save a SageMaker model to the model registry#

This guide uses code snippets from the official Amazon SageMaker Examples repository.1

Let's say you have trained the knn model in SageMaker as below:

# set up the estimator
knn = sagemaker.estimator.Estimator(
    get_image_uri(boto3.Session().region_name, "knn"),
    get_execution_role(),
    instance_count=1,
    instance_type="ml.m5.large",
    output_path=output_path,
    sagemaker_session=sagemaker.Session(),
)
knn.set_hyperparameters(**hyperparams)

# train a model. fit_input contains the locations of the train and test data
fit_input = {"train": s3_train_data}
if s3_test_data is not None:
    fit_input["test"] = s3_test_data
knn.fit(fit_input)

To store the model in the Neptune model registry, you first need to create a new model.

import neptune

model = neptune.init_model(key="AWS")

Next, create a version of the model.

model_version = neptune.init_model_version(model="???-AWS") # (1)!
  1. ??? stands for the project key, which together with the model key forms the model ID. You need the full ID when you create new versions of the model.

Then, you can save all the relevant model artifacts to the model registry. For example:

model_version["model"].track_files(knn.model_data)  # tar.gz with the trained model
model_version["hyperparameters"] = knn.hyperparameters()
model_version["training_image_uri"] = knn.training_image_uri()

Related

  1. In particular, the notebook showing how to train k-nearest neighbors algorithm in SageMaker . ↩