CODEWITHSUNDEEP
amazon-sagemaker
DanielR
DanielR

Reputation: 596

Create a custom AWS Sagemaker Estimator with a configurable entrypoint

I'm writing a custom Estimator in AWS Sagemaker , for a framework that isn't supported out-of-the-box. I have my own docker image for training, with the training code bundled into the image, which forces me to rebuild the image every time the code changes.

What I would like to do, is create an Estimator that uses this image, and accepts a file as an entry-point, like the built-in framework estimators do (e.g Tensorflow).

From reading the source code for the Sagemaker python SDK, I found the sagemaker.estimator.Framework class, which accepts the entry_point argument, and from which the built-in frameworks estimators inherit. However, the documentation doesn't really show how to inherit from that class in my own code.

Is it possible to write a custom Estimator class that inherits from Framework, or is there another way to create a custom estimator that receives an entry-point argument?

Upvotes: 5

Views: 3502

Answers (2)

Harish Rao
Harish Rao

Reputation: 11

Instead of using an Estimator, use a TensorFlow estimator and then pass your custom image as an input to image_uri. This way you have the best of the both worlds. And if you have a requirements.txt file in your source_dir, that gets installed before training.

Here is a custom image docker file

FROM python:3.9-buster

RUN apt-get -y update && apt-get install -y --no-install-recommends \
     wget \
     python3-pip \
     python3-setuptools \
     nginx \
     ca-certificates \
     vim \
&& rm -rf /var/lib/apt/lists/*

#RUN ln -s /usr/bin/python3 /usr/bin/python
RUN ln -s /usr/bin/pip3 /usr/bin/pip

ENV PYTHONUNBUFFERED=TRUE
ENV PYTHONDONTWRITEBYTECODE=TRUE
ENV PATH="/opt/program:${PATH}"

RUN pip install sagemaker-training

and here is how you call your custom image

tf_estimator = TensorFlow(entry_point='mnist_keras_tf.py',
                        source_dir='.',
                        role=role,
                        instance_count=1, 
                        instance_type='ml.p3.2xlarge',
                        image_uri='<arn of custom image>',
                        script_mode=True,
                        hyperparameters={
                            'epochs': 2,
                            'batch-size': 256,
                            'learning-rate': 0.01}

Upvotes: 1

lauren
lauren

Reputation: 513

The existing framework estimators in the SageMaker Python SDK can be a good place to start, e.g. https://github.com/aws/sagemaker-python-sdk/blob/master/src/sagemaker/pytorch/estimator.py

However, it would probably be easier, since you've already bundled your training code into the image, to set an entrypoint in your Dockerfile that runs the training code. For more about how Amazon SageMaker runs your training image, see https://docs.aws.amazon.com/sagemaker/latest/dg/your-algorithms-training-algo-dockerfile.html

Upvotes: 0

Related Questions