added code

code/new_realtime/models/utils.py

@@ -0,0 +1,137 @@
+import os
+from typing import Callable
+import math
+
+import numpy as np
+from bson import ObjectId
+import torch
+from torch import nn
+
+from utils.data_utils import get_collated_batch_for_key
+
+
+def simple_get_y(collated_batch):
+    """
+    Get y from collated batch
+    Args:
+        collated_batch: collated batch to get y from
+
+    Returns:
+        np.ndarray: y
+    """
+    return collated_batch[1]
+
+
+def simple_model_save(model: nn.Module,
+                      training_configuration: dict) -> None:
+    """
+    Save model to disk
+    Args:
+        model: model to save
+        training_configuration: training configuration
+    Returns:
+        None
+    """
+    model_path = os.path.join(training_configuration["training_dir"], "model.pt")
+    torch.save(model.state_dict(), model_path)
+
+
+def simple_model_load(model_configuration: dict,
+                      training_configuration: dict,
+                      sample_key: str | ObjectId,
+                      device: str,
+                      model_creation_fn: Callable,
+                      *args, **kwargs) -> nn.Module:
+    """
+    Load model from disk
+    Args:
+        model_configuration: model configuration
+        training_configuration: training configuration
+        sample_key: sample key to get sample data batch with
+        device: device to load model on
+        model_creation_fn: function to create model
+    Returns:
+        model: loaded model
+    """
+    training_dir = training_configuration["training_dir"]
+    model_state_path = os.path.join(training_dir, "model.pt")
+    model = model_creation_fn(model_configuration,
+                              sample_key, *args, **kwargs)
+    model.to(device)
+    model.load_state_dict(torch.load(model_state_path))
+    return model
+
+
+def simple_model_creation(model_configuration: dict,
+                          sample_key: str | ObjectId,
+                          lmdb_env=None) -> nn.Module:
+    """
+    Create model from configuration
+    Args:
+        model_configuration: model configuration
+        sample_key: sample key to get sample data batch with
+        lmdb_env: LMDB environment to use for getting sample data batch
+
+    Returns:
+        model: created model
+    """
+    sample_item = get_collated_batch_for_key(sample_key, model_configuration, lmdb_env=lmdb_env)
+    input_size = sample_item[0].shape[2]
+    output_size = sample_item[1].shape[2]
+    model_class = model_configuration["model_class"]
+    cnn_model = model_class(input_dim=input_size,
+                            output_dim=output_size,
+                            **model_configuration["model_parameters"], )
+    return cnn_model
+
+
+def simple_x_y_predict(model: nn.Module,
+                       collated_batch: tuple[torch.Tensor, torch.Tensor],
+                       device: str,
+                       batch_size: int,
+                       *args, **kwargs) -> torch.Tensor:
+    """
+    Predict y from x
+    Args:
+        model: model to use for prediction, must have 'predict' method
+        collated_batch: collated data batch to predict with,
+            must be a tuple of (x, y) where x is the input data and y is the target data
+        device: device to use for prediction
+        *args: additional arguments to pass to the model's predict method
+        **kwargs: additional keyword arguments to pass to the model's predict method
+
+    Returns:
+        torch.Tensor: predicted y as numpy array on CPU
+    """
+
+    model.eval()
+    x = collated_batch[0].to(device).float()
+    batches = list()
+    num_batches = math.ceil(len(x) / batch_size)
+    for i in range(num_batches):
+        start = i * batch_size
+        end = (i + 1) * batch_size
+        if end > len(x):
+            end = len(x)
+        batch_slice = x[start:end]
+        batches.append(batch_slice)
+
+    preds = list()
+    with torch.no_grad():
+        for batch in batches:
+            batch = batch.to(device).float()
+            pred = model(batch, *args, **kwargs)
+            preds.append(pred.cpu().numpy())
+    return np.concatenate(preds)
+
+
+def simple_get_y(collated_batch: tuple[torch.Tensor, torch.Tensor], ) -> torch.Tensor:
+    """
+    Get y from collated batch
+    Args:
+        collated_batch: collated batch to get y from
+
+    Returns:
+        torch.Tensor: y
+    """
+    return collated_batch[1]