aws / aws/sagemaker-huggingface-inference-toolkit

Sagemaker inference not loading model weight from s3

Open
#116 2 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
270
Forks
60
PR merge metrics
No merged PRs in 30d

Description

I have the following model_fn where I'm trying to load 2 pth files using torch sagemaker is able to load one .pth file but when coming to the second file it fails I've also tried deploying only the second .pth file but it fails.But the model seems to load when its downloaded from source but when i download it from source and make it fetch from s3 it fails again.Both the files are in model_dir
Here is my model_fn
def model_fn(model_dir):
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)
print("Model Loading...")
det_model = db_resnet50(pretrained=False, pretrained_backbone=False)
det_model_path = os.path.join(model_dir, 'db_resnet50-ac60cadc.pt')
det_params = torch.load(det_model_path, map_location=device)
det_model.load_state_dict(det_params)
print("loading second file")
model = crnn_vgg16_bn(pretrained=True, pretrained_backbone=False)
#failing on this step
rec_model_path = os.path.join(model_dir, 'crnn_vgg16_bn-9762b0b0.pt')
reco_params = torch.load(rec_model_path, map_location=device)
reco_model.load_state_dict(reco_params)
model = ocr_predictor(det_arch=det_model, reco_arch=reco_model, pretrained=False)
model.to(device=device)
print("model_loaded")
return model

Contributor guide

Open the contributing guide

Research direction

Start with the shown model_fn and reproduce the SageMaker deployment while inspecting the runtime logs around the second torch.load call. Verify the model_dir contents and loading behavior for db_resnet50-ac60cadc.pt and crnn_vgg16_bn-9762b0b0.pt; done means both weights load successfully in the deployed inference container.

Written by the indexing model from the issue text.

Assessment

Tech stack
aws, python, pytorch
Domain
backend, machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.