Save a SageMaker model to Neptune

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

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 Neptune, create a run:

from neptune.integrations.aws import init_run 

run = neptune.init_run()

Then, save all the relevant model artifacts. For example:

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

Related

• How to use Neptune in SageMaker training jobs – you can save the outputs as described on this page.
• Models in Neptune

---

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