Skip to content

How to use Neptune in SageMaker training jobs#

To use Neptune with model training jobs in Amazon SageMaker, you need to either

  • A) use a custom Docker container where Neptune is preinstalled, or
  • B) use the jobs with custom scripts and a requirements.txt file that includes neptune

and in both cases you need to run your own code in the training job that uses the Neptune client library.

To demonstrate how to use Neptune in SageMaker jobs, we provide two examples.

Using Neptune in training jobs with custom Docker containers

This tutorial is an adaptation of the official SageMaker tutorial by AWS.

You can run this part of the notebook either locally or from a SageMaker notebook. Additional dependencies include AWS CLI tools and Docker.

Go to example on GitHub 

Using Neptune in training jobs with PyTorch Estimator

Run this example in a SageMaker notebook.

In this guide, we train a model on the MNIST dataset with the SageMaker PyTorch Estimator. We show how to enable Neptune in the workflow and set up Neptune logging in the code.

Go to example on GitHub