Merge pull request 'add_model_resnet18' (#34) from add_model_resnet18 into main

Reviewed-on: http://192.168.1.2:3000/brian/visual_critical_discourse_analysis/pulls/34

core/database/__init__.py

@@ -1,10 +1,10 @@
 """Database module content."""
 from __future__ import annotations
 
-from .classes import Dataset
-from .classes import NoDocumentFoundException
-from .classes import VisualCommunication
-from .utils import connect
+from .classes.dataset import Dataset
+from .classes.exceptions import NoDocumentFoundException
+from .classes.visual_communication import VisualCommunication
+from .utils.connect import connect
 from .utils.count_documents import count_documents
 from .utils.get_visual_communication import get_visual_communication
 from .utils.list_names import list_names

core/database/utils/get_dataset.py

@@ -0,0 +1,13 @@
+from __future__ import annotations
+
+from typing import Literal
+
+from pymongo.collection import Collection
+
+
+def get_dataset(
+    collection: Collection,
+    type: Literal['train', 'test', 'validation'],
+) -> list[str]:
+    """Get list of data names for the corresponding type."""
+    return []

core/database/utils/save_dataset.py

@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+import logging
+
+from pymongo.collection import Collection
+
+from core.database import Dataset
+
+
+def save_dataset(
+    collection: Collection,
+    dataset: Dataset,
+) -> None:
+    """Save dataset to database."""
+    res = collection.insert_one(
+        document=dataset.model_dump(),
+    )
+    logging.debug('inserted document: %s', res)
+
+
+if __name__ == '__main__':
+    from dotenv import load_dotenv
+    load_dotenv('local.env')
+    from core.database import list_names, connect
+    # connect to database
+    collection, db, client = connect()
+    print(client.server_info())
+
+    name_list = list_names(collection=collection, only_with_annotation=True)
+    ds = Dataset.new_from_name_list(name_list=name_list)
+
+    print(ds)
+    # save_dataset(
+    #     collection=collection,
+    #     dataset=ds
+    # )

model/src/train.py

@@ -1,6 +1,49 @@
 from __future__ import annotations
 
-from dataloader import VCDADataset  # noqa: F401
-from loss_fn import setup_criterion  # noqa: F401
+from datetime import datetime
 
-from model import ResNet18Head  # noqa: F401
+import torch
+import torch.nn as nn
+from dataloader import VCDADataset  # noqa: F401
+
+from .models import VisualCommunicationModel
+
+PRE_WARMUP_LR = 1e-10
+POST_WARMUP_LR = 1e-5
+BATCH_SIZE = 32
+MODEL_NAME = 'visual_communication_model_v1'
+MAX_EPOCHS = 200
+RUN_NAME = datetime.now().strftime('%Y-%m-%d-%H%M') + '_' + MODEL_NAME
+DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+
+# create loss function
+criterion = nn.MSELoss()
+
+
+def loss_fn(pred, target):
+    """Create loss function."""
+    return criterion(pred, target.unsqueeze(-1))
+
+
+# create model
+model = VisualCommunicationModel().to(DEVICE)
+
+# create optimizer
+optim = torch.optim.Adam(model.parameters(), lr=POST_WARMUP_LR)
+
+# create datasets
+# train_dataset
+# validation_dataset
+
+# create dataloaders
+
+# create trainer and evaluator
+
+# setup progressbar
+
+# setup lr scheduler
+
+# setup checkpoint saving
+
+# load in checkpoint if exists

pyproject.toml

@@ -5,7 +5,7 @@ description = ""
 authors = ["Brian Bjarke Jensen <[email protected]>"]
 readme = "README.md"
 packages = [
-    { include = "src" },
+    { include = "core" },
 ]
 
 [tool.poetry.dependencies]