implemented load_model function

Dockerfile.model

@@ -44,7 +44,7 @@ COPY --from=builder ${VIRTUAL_ENV} ${VIRTUAL_ENV}
 ENV PATH="${VIRTUAL_ENV}/bin:${PATH}"
 
 # create home directory and app user
-RUN mkdir -p /home/app
+RUN mkdir -p $APP_HOME
 
 # add code while changing ownership
 WORKDIR $APP_HOME

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()
+    model: VisualCommunicationModel = load_model(client=minio_client)
-    # 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)
     # show model weights
     summary(model)
     print('loaded model')

model/src/utils/__init__.py

@@ -1 +1 @@
-from .get_model_name import get_model_name
+from .load_model import load_model

model/src/utils/load_model.py

@@ -0,0 +1,40 @@
+"""Definition of load_model function."""
+
+import logging
+from pathlib import Path
+
+import torch
+from minio import Minio
+from models import VisualCommunicationModel
+
+from shared.data_store import get_model
+
+from .get_model_name import get_model_name
+
+DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+
+def load_model(
+    client: Minio,
+) -> VisualCommunicationModel:
+    """Instantiate model with weights loaded from latest model saved in
+    MinIO."""
+    assert isinstance(client, Minio)
+    # 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 = get_model(
+        client=client,
+        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

shared/data_store/get_model.py

@@ -23,7 +23,6 @@ def get_model(
         client=client,
         object_name=object_name,
     )
-    logging.debug('buffer size: %s', buffer.getbuffer().nbytes)
     # convert data to model checkpoint
     buffer.seek(0)
     model_content = torch.load(buffer)