sagemaker_sklearn_model_training_quickstart.py

1import sagemaker
2from sagemaker.sklearn.estimator import SKLearn
3
4# Initialize a SageMaker Session
5sagemaker_session = sagemaker.Session()
6
7# Get the execution role for the notebook instance
8role = sagemaker.get_execution_role()
9
10# Define the location of the training data
11train_input = sagemaker_session.upload_data(
12    path='data', 
13    key_prefix='input/sklearn-iris'
14)
15
16# Configure the Scikit-learn estimator
17sklearn_estimator = SKLearn(
18    entry_point='train.py',
19    role=role,
20    instance_count=1,
21    instance_type='ml.m5.large',
22    framework_version='1.2-1',
23    py_version='py3',
24    hyperparameters={'epochs': 20}
25)
26
27# Launch the training job
28sklearn_estimator.fit({'train': train_input})