From 8b1893b30a283ae4c5f93a1a1305a4d3c2b0ae3c Mon Sep 17 00:00:00 2001 From: Brian Bjarke Jensen Date: Wed, 24 Jul 2024 09:13:28 +0200 Subject: [PATCH] implemented load_model function --- model/src/main.py | 23 +++-------------------- 1 file changed, 3 insertions(+), 20 deletions(-) diff --git a/model/src/main.py b/model/src/main.py index 888131f..e776beb 100644 --- a/model/src/main.py +++ b/model/src/main.py @@ -1,42 +1,25 @@ """Main script to be run by service.""" -import logging from pathlib import Path -import torch from dotenv import load_dotenv from models import VisualCommunicationModel from torchinfo import summary -from utils import get_model_name +from utils import load_model -from shared.data_store import connect, get_model +from shared.data_store import connect from shared.utils import setup_logging -DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu') - if __name__ == '__main__': # load in env file env_path = Path(__file__).parent.parent.parent / 'server.env' - assert env_path.exists() load_dotenv(env_path) # setup logging setup_logging() # connect to minio minio_client = connect() # instantiate model - model = VisualCommunicationModel() - # get model object name - model_name_path = Path(__file__).parent / 'model_name.txt' - model_object_name = get_model_name(path=model_name_path) - logging.info('using model: %s', model_object_name) - # load model from minio - model_checkpoint = get_model( - client=minio_client, - object_name=model_object_name, - ) - model.load_state_dict(model_checkpoint) - # move model to selected device - model = model.to(DEVICE) + model: VisualCommunicationModel = load_model(client=minio_client) # show model weights summary(model) print('loaded model')