Custom Inference.py script for Vision Transformer

I am trying to deploy a Vision Transformer huggingface model to a sagemaker endpoint. I am trying to make an inference.py script to go along with the model artificats, so that I can have my own custom output from predict_fn.

However, the endpoint keeps returning this error when passing the image through the feature_extractor:

ModelError: An error occurred (ModelError) when calling the InvokeEndpoint operation: Received client error (400) from model with message "{
“code”: 400,
“type”: “InternalServerException”,
“message”: “an integer is required (got type str)”

I’ve checked and confirmed that the image is indeed a PIL image (which works when I pass it through the endpoint regularly, without the inference script, just without returning the desired output format). I have even tried to convert the image to a numpy array, but I still get the same error.

This is the script:
def model_fn(model_dir):
logging.info(‘within model_fn’)
model = ViTForImageClassification.from_pretrained(model_dir)
feature_extractor = ViTFeatureExtractor.from_pretrained(model_dir)
logging.info(‘Extracted Model and Extractor’)
return model, feature_extractor

def predict_fn(data, model_extractor):
# destruct model and tokenizer
logging.info(‘Inside Prediction_fn’)
model, feature_extractor = model_extractor

inputs = feature_extractor(images=data['inputs'], return_tensors="pt")
outputs = model(**inputs)
logits = outputs.logits

# return {'outputs': outputs, 'logits': logits}
return {'logits': logits}

Thanks in advance!

So an answer, but comes with a follow up question: The error only occurs when I load the feature extractor from the trained artifacts I have saved. When I input the huggingface hub name of the original feature extractor I used though (google ViT), the error goes away.

This only seems to be an error when I downgrade to transformers version 4.17, required for sagemaker deployments.

Any solutions on how to load it from the artifacts, in version 4.17?

If your issue is related to the fact that you have to use version 4.17 for Sagemaker deployment - there is a nice little trick to use higher versions in Sagemaker. You can upgrade the container image to install the latest version of transformers and use that for your deployment.

Details in this blog post: Unlock the Latest Transformer Models with Amazon SageMaker | by Heiko Hotz | Dec, 2022 | Towards Data Science