39 lines
1.1 KiB
Python
39 lines
1.1 KiB
Python
"""Definition of load_model function."""
|
|
|
|
import logging
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
|
|
from model.src.models import VisualCommunicationModel
|
|
from shared.datastore import Datastore
|
|
|
|
from .get_model_name import get_model_name
|
|
|
|
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
|
|
|
|
|
def load_model(
|
|
datastore: Datastore,
|
|
) -> VisualCommunicationModel:
|
|
"""Instantiate model with weights loaded from latest model saved in
|
|
MinIO."""
|
|
assert isinstance(datastore, Datastore)
|
|
# instantiate model
|
|
model = VisualCommunicationModel()
|
|
# get model object name
|
|
model_name_path = Path('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 = datastore.get_model(
|
|
object_name=model_object_name,
|
|
)
|
|
model.load_state_dict(model_checkpoint)
|
|
# clean memory
|
|
model_checkpoint.clear()
|
|
# move model to selected device
|
|
model = model.to(DEVICE)
|
|
logging.debug('finished')
|
|
return model
|