sagemaker_sklearn_estimator_training_job_quickstart.py

1import sagemaker
2from sagemaker.sklearn.estimator import SKLearn
3
4# Get the execution role and session
5role = sagemaker.get_execution_role()
6session = sagemaker.Session()
7
8# Define the Scikit-learn estimator
9# Note: 'entry_point' is 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    py_version='py3',
17    hyperparameters={'epochs': 10}
18)
19
20# Launch the training job
21# 'inputs' can be an S3 URI or a SageMaker TrainingInput object
22sklearn_estimator.fit({'train': 's3://my-bucket/my-training-data/'})