CODEWITHSUNDEEP
amazon-sagemaker
Luis Quiñones
Luis Quiñones

Reputation: 25

How can I load a trained model from the Model Registry into a SageMaker pipeline to make inference?

I am following a SageMaker Pipelines notebook from the aws-samples repository. At the bottom of this notebook, there is Batch Transform section which does a task that I need to do, which is run inference on a dataset from SageMaker pipelines. My problem is that I am trying to replace the training step with loading a model that is in the Model Registry, trained within a pipeline.

This is the cell that I want to replace to load the model instead of taking the model from the train step:

# First, create a SageMaker model for the transform
model = Model(
    image_uri=xgboost_image_uri,        
    model_data=step_train.properties.ModelArtifacts.S3ModelArtifacts,
    name=f"from-idea-to-prod-xgboost-model",
    sagemaker_session=session,
    role=sm_role,
)

'''
Insert load model from Model Registry
'''

# Define create model step
step_create_model = ModelStep(
    name=f"{pipeline_name}-model",
    step_args=model.create(instance_type="ml.m5.large", accelerator_type="ml.eia1.medium"),
)

Please let me know whether this is possible or I am taking the wrong approach. The requirement here is to be able to trigger this inference pipeline via a lambda function.

I have tried to do the ModelStep or even CreateModelStep from the model package arn without any luck.

Upvotes: 1

Views: 1245

Answers (1)

Ram Vegiraju
Ram Vegiraju

Reputation: 377

Is there a training step in your Pipeline, are you then registering this trained model and then deploying for Batch Inference? You can first do a CreateModel step with the trained model artifacts and then take this step and it's metadata and register it

model = Model(
    image_uri=image_uri,
    model_data=training_step.properties.ModelArtifacts.S3ModelArtifacts,
    role=role,
    sagemaker_session=pipeline_session
)

create_model_step = ModelStep(
    name="CreateXGBoostModel",
    step_args=model.create(),
)

# registering model object
from sagemaker.workflow.step_collections import RegisterModel
register_step = RegisterModel(
 name="AbaloneRegisterModel",
 model=model,
 content_types=["text/csv"],
 response_types=["text/csv"],
 inference_instances=["ml.t2.medium", "ml.m5.xlarge"],
 transform_instances=["ml.m5.xlarge"],
 model_package_group_name='batchgroup',
)

You can then take the create model metadata and run Batch Inference:

from sagemaker.workflow.pipeline_context import PipelineSession

from sagemaker.transformer import Transformer
from sagemaker.inputs import TransformInput
from sagemaker.workflow.steps import TransformStep

transformer = Transformer(model_name=create_model_step.properties.ModelName,
                          instance_count=1, instance_type=batch_transform_param, 
                          assemble_with="Line", accept="text/csv", 
                          sagemaker_session=PipelineSession())

transform_step = TransformStep(
    name="AbaloneTransform",
    step_args=transformer.transform(data=test_data_param, 
                                    content_type = "text/csv"),
)

For the entire end to end example please reference the following Github: https://github.com/RamVegiraju/sagemaker-pipelines-examples/blob/master/train-batch-transform/pipelines-train-transform.ipynb

Upvotes: 2

Related Questions