aws / aws/sagemaker-huggingface-inference-toolkit
Sagemaker inference not loading model weight from s3
- 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
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