sagemaker_sklearn_estimator_iris_training_quickstart.py

1import sagemaker
2from sagemaker.sklearn.estimator import SKLearn
3
4# Get the SageMaker session and execution role
5sagemaker_session = sagemaker.Session()
6role = sagemaker.get_execution_role()
7
8# Define the Scikit-learn estimator
9# Note: 'entry_point' should be the path to your training script (e.g., 'train.py')
10sklearn_estimator = SKLearn(
11    entry_point='train.py',
12    role=role,
13    instance_count=1,
14    instance_type='ml.m5.large',
15    framework_version='1.2-1',
16    base_job_name='rf-scikit-iris'
17)
18
19# Start the training job
20# 'train_data' can be an S3 URI or a local path if using Local Mode
21sklearn_estimator.fit({'train': 's3://sagemaker-sample-data/sklearn/iris.csv'})