SageMaker Feature Store

Duration: 5 min

This module delves into Amazon SageMaker Feature Store, a fully managed service that makes it easy to store, manage, and serve machine learning (ML) features at scale. Understanding how to use Feature Store is crucial for building efficient ML pipelines, as it streamlines the process of feature engineering and management, ultimately leading to more robust and scalable ML models.

Creating and Storing Features

Amazon SageMaker Feature Store allows you to create and store features that can be used for training and serving ML models. Features are stored in a centralized repository, making it easy to manage and version them. This ensures that the features used for training are consistent with those used for inference, reducing the risk of model drift.

import boto3
from sagemaker.feature_store.feature_group import FeatureGroup

# Initialize boto3 session
session = boto3.Session()

sagemaker_client = session.client('sagemaker', region_name='us-west-2')

# Create a Feature Group
feature_group_name = 'example-feature-group'
feature_group = FeatureGroup(name=feature_group_name, sagemaker_session=session)

# Define the feature definitions
feature_definitions = [
    {"FeatureName": "user_id", "FeatureType": "String"},
    {"FeatureName": "age", "FeatureType": "Integral"},
    {"FeatureName": "income", "FeatureType": "Fractional"}
]

# Create the Feature Group
feature_group.create(feature_definitions=feature_definitions, record_identifier_name='user_id', event_time_feature_name='event_time')

print(f'Feature Group {feature_group_name} created.')

Try it in Google Colab:

Feature Group example-feature-group created.

Querying Features for Model Training

Once features are stored in the Feature Store, you can query them for use in model training. SageMaker provides an Athena query interface to retrieve features, allowing you to easily integrate feature data into your training pipelines. This ensures that your models are trained on the most up-to-date and consistent feature data.

import pandas as pd
from sagemaker.feature_store.feature_group import FeatureGroup

# Initialize boto3 session
session = boto3.Session()

sagemaker_client = session.client('sagemaker', region_name='us-west-2')

# Load the Feature Group
feature_group_name = 'example-feature-group'
feature_group = FeatureGroup(name=feature_group_name, sagemaker_session=session)

# Query features for model training
query = feature_group.athena_query()
query.run(f"SELECT user_id, age, income FROM {feature_group_name}")
results = query.as_dataframe()

print(results.head())

💡 Tip: Ensure that the IAM role associated with your SageMaker session has the necessary permissions to access the Feature Store and query data from Athena.

Key Concepts

Check Your Understanding