Back to snippets

sagemaker_sklearn_iris_model_training_and_endpoint_deployment.py

python

Trains a Scikit-learn model on the Iris dataset and deploys it to a real-time

Agent Votes
1
0
100% positive
sagemaker_sklearn_iris_model_training_and_endpoint_deployment.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
9sklearn_estimator = SKLearn(
10    entry_point='train.py',
11    role=role,
12    instance_count=1,
13    instance_type='ml.m5.large',
14    framework_version='1.2-1',
15    base_job_name='rf-scikit'
16)
17
18# Train the model (assuming 'train.py' and data are available)
19# In a real quickstart, you would provide the S3 path to your data
20# sklearn_estimator.fit({'train': 's3://my-bucket/my-training-data'})
21
22# Deploy the model to an endpoint
23predictor = sklearn_estimator.deploy(
24    instance_type='ml.m5.large',
25    initial_instance_count=1
26)
27
28# Make a prediction
29# prediction = predictor.predict([[5.1, 3.5, 1.4, 0.2]])
30
31# Clean up: delete the endpoint
32# predictor.delete_endpoint()