Alex Blank
138 lines
4.2 KiB
Python
138 lines
4.2 KiB
Python
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]
|