SageMaker sklearn-container example for training

SageMaker sklearn-container example

# S3 prefix
prefix = 'sagemaker/scikit-iris'

import sagemaker
from sagemaker import get_execution_role
sagemaker_session = sagemaker.Session()

# Get a SageMaker-compatible role used by this Notebook Instance.
role = get_execution_role()

import numpy as np
import os
from sklearn import datasets

# Load Iris dataset, then join labels and features
iris = datasets.load_iris()
joined_iris = np.insert(iris.data, 0, iris.target, axis=1)

# Create directory and write csv
os.makedirs('./data', exist_ok=True)
np.savetxt('./data/iris.csv', joined_iris, delimiter=',', fmt='%1.1f, %1.3f, %1.3f, %1.3f, %1.3f')

WORK_DIRECTORY = 'data'
train_input = sagemaker_session.upload_data(WORK_DIRECTORY, key_prefix="{}/{}".format(prefix, WORK_DIRECTORY))

account = sagemaker_session.boto_session.client('sts').get_caller_identity()['Account']
region = sagemaker_session.boto_session.region_name
image_full = '{}.dkr.ecr.{}.amazonaws.com/sklearn-container:latest'.format(account, region)

clf = sagemaker.estimator.Estimator(image_full, role, 1, 'ml.c4.2xlarge', 
                                    output_path="s3://{}/output".format(sagemaker_session.default_bucket()),
                                    sagemaker_session=sagemaker_session)
 
# training with the gradient boosting classifier model
clf.fit(train_input)