sagemaker_sklearn_estimator_iris_training_quickstart.py

1import sagemaker
2from sagemaker.sklearn.estimator import SKLearn
3
4# Initialize a SageMaker Session
5sagemaker_session = sagemaker.Session()
6
7# Get a SageMaker-compatible role used for training and deploying models
8role = sagemaker.get_execution_role()
9
10# Define the Scikit-learn Estimator
11sklearn_estimator = SKLearn(
12    entry_point='train.py',
13    role=role,
14    instance_count=1,
15    instance_type='ml.m5.large',
16    framework_version='0.23-1',
17    py_version='py3',
18    hyperparameters={'epochs': 20}
19)
20
21# Start the training job
22# Note: 'train.py' should be in the same directory and 's3_data_path' should point to your dataset in S3
23sklearn_estimator.fit({'train': 's3://my-bucket/my-training-data'})