| import torch | |
| import torch.nn as nn | |
| from torchvision import models | |
| import huggingface_hub | |
| from huggingface_hub import hf_hub_download | |
| import torchvision | |
| model_repo_id = "CSSE416-final-project/faceRecogModel" | |
| weight_file_id = "modelWeights100.bin" | |
| def load_model(repo_id): | |
| # Download the model weights from the repo | |
| weights_path = hf_hub_download(repo_id=model_repo_id, filename=weight_file_id) | |
| # Initialize the ResNet-18 architecture | |
| model = torchvision.models.resnet18(pretrained=False) | |
| num_ftrs = model.fc.in_features | |
| model.fc = nn.Linear(num_ftrs, 100) # Adjust for your task (e.g., 128 classes) | |
| # Load the model weights | |
| state_dict = torch.load(weights_path, map_location=torch.device("cpu")) | |
| model.load_state_dict(state_dict) | |
| model.eval() # Set the model to evaluation mode | |
| return model |