Introduction to JumpStart - Image Embedding
This notebook’s CI test result for us-west-2 is as follows. CI test results in other regions can be found at the end of the notebook.
Note: This notebook was tested on ml.t3.medium instance in Amazon SageMaker Studio with Python 3 (Data Science) kernel and in Amazon SageMaker Notebook instance with conda_python3 kernel.
1. Set Up
[ ]:
!pip install sagemaker ipywidgets --upgrade --quiet
Permissions and environment variables
[ ]:
import sagemaker, boto3, json
from sagemaker import get_execution_role
aws_role = get_execution_role()
aws_region = boto3.Session().region_name
sess = sagemaker.Session()
2. Select a model
Here, we download jumpstart model_manifest file from the jumpstart s3 bucket, filter-out all the Image Embedding models and select a model for inference. ***
[ ]:
from ipywidgets import Dropdown
# download JumpStart model_manifest file.
boto3.client("s3").download_file(
f"jumpstart-cache-prod-{aws_region}", "models_manifest.json", "models_manifest.json"
)
with open("models_manifest.json", "rb") as json_file:
model_list = json.load(json_file)
# filter-out all the Image Embedding models from the manifest list.
image_embedding_models = []
for model in model_list:
model_id = model["model_id"]
if "-icembedding-" in model_id and model_id not in image_embedding_models:
image_embedding_models.append(model_id)
# display the model-ids in a dropdown to select a model for inference.
model_dropdown = Dropdown(
options=image_embedding_models,
value="tensorflow-icembedding-bit-m-r101x1-ilsvrc2012-featurevector-1",
description="Select a model",
style={"description_width": "initial"},
layout={"width": "max-content"},
)
Chose a model for Inference
[ ]:
display(model_dropdown)
[ ]:
# model_version="*" fetches the latest version of the model
model_id, model_version = model_dropdown.value, "*"
3. Retrieve JumpStart Artifacts & Deploy an Endpoint
We start by retrieving the deploy_image_uri
, deploy_source_uri
, and model_uri
for the pre-trained model. To host the pre-trained model, we create an instance of `sagemaker.model.Model
<https://sagemaker.readthedocs.io/en/stable/api/inference/model.html>`__ and deploy it. This may take a few minutes. ***
[ ]:
from sagemaker import image_uris, model_uris, script_uris, hyperparameters
from sagemaker.model import Model
from sagemaker.predictor import Predictor
from sagemaker.utils import name_from_base
endpoint_name = name_from_base(f"jumpstart-example-infer-{model_id}")
inference_instance_type = "ml.p2.xlarge"
# Retrieve the inference docker container uri. This is the base Tensorflow container image for the default model above.
deploy_image_uri = image_uris.retrieve(
region=None,
framework=None, # automatically inferred from model_id
image_scope="inference",
model_id=model_id,
model_version=model_version,
instance_type=inference_instance_type,
)
# Retrieve the inference script uri. This includes all dependencies and scripts for model loading, inference handling etc.
deploy_source_uri = script_uris.retrieve(
model_id=model_id, model_version=model_version, script_scope="inference"
)
# Retrieve the model uri. This includes the model and model parameters.
model_uri = model_uris.retrieve(
model_id=model_id, model_version=model_version, model_scope="inference"
)
# Create the SageMaker model instance
model = Model(
image_uri=deploy_image_uri,
source_dir=deploy_source_uri,
model_data=model_uri,
entry_point="inference.py", # entry point file in source_dir and present in deploy_source_uri
role=aws_role,
predictor_cls=Predictor,
name=endpoint_name,
)
# deploy the Model. Note that we need to pass Predictor class when we deploy model through Model class,
# for being able to run inference through the sagemaker API.
model_predictor = model.deploy(
initial_instance_count=1,
instance_type=inference_instance_type,
predictor_cls=Predictor,
endpoint_name=endpoint_name,
)
We download example images from the JumpStart assets S3 bucket.
[ ]:
region = boto3.Session().region_name
s3_bucket = f"jumpstart-cache-prod-{region}"
key_prefix = "inference-notebook-assets"
s3 = boto3.client("s3")
def download_from_s3(key_filenames):
for key_filename in key_filenames:
s3.download_file(s3_bucket, f"{key_prefix}/{key_filename}", key_filename)
cat_jpg, dog_jpg = "cat.jpg", "dog.jpg"
download_from_s3(key_filenames=[cat_jpg, dog_jpg])
5. Query endpoint and parse response
[ ]:
def query(model_predictor, image_file_name):
"""Query the model predictor."""
with open(image_file_name, "rb") as file:
input_img_rb = file.read()
query_response = model_predictor.predict(
input_img_rb,
{
"ContentType": "application/x-image",
"Accept": "application/json",
},
)
return query_response
def parse_response(query_response):
"""Parse response and return the embedding."""
model_predictions = json.loads(query_response)
translation_text = model_predictions["embedding"]
return translation_text
[ ]:
from IPython.core.display import HTML
for img_name in [cat_jpg, dog_jpg]:
query_response = query(model_predictor, img_name)
embedding = parse_response(query_response)
first_5element_embeddings = "{" + ", ".join([str(id) for id in embedding[:5]]) + "}"
display(
HTML(
f'<img src={img_name} alt={img_name} align="left" style="width: 250px;"/>'
f"<figcaption>First-5 elements of the feature vector (embedding) are: {first_5element_embeddings}</figcaption>"
f"<figcaption>Total length of the feature vector (embedding) is: {len(embedding)}</figcaption>"
)
)
[ ]:
# Delete the SageMaker endpoint
model_predictor.delete_model()
model_predictor.delete_endpoint()
This notebook was tested in multiple regions. The test results are as follows, except for us-west-2 which is shown at the top of the notebook.