From 5391927c58e01694e3eeff6cb8c7e866fb20f6a4 Mon Sep 17 00:00:00 2001
From: Brian Bjarke Jensen <bbj@skov.dk>
Date: Tue, 30 Apr 2024 21:06:53 +0200
Subject: [PATCH] unfinished work

---
 model/src/train.py | 49 +++++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 46 insertions(+), 3 deletions(-)

diff --git a/model/src/train.py b/model/src/train.py
index c95e37f..20c7370 100644
--- a/model/src/train.py
+++ b/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