File size: 862 Bytes
64a41ee 1c9d3b0 64a41ee 1c9d3b0 64a41ee 1c9d3b0 64a41ee 1c9d3b0 b0dca77 1c9d3b0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 |
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 |