214 lines
9.2 KiB
Python
214 lines
9.2 KiB
Python
import os
|
|
|
|
import lmdb
|
|
import numpy as np
|
|
import torch
|
|
from torch import nn
|
|
|
|
from utils.dataset_creation import get_scalers_for_model, inverse_scale_feature
|
|
from utils.data_utils import get_collated_batch_for_key
|
|
from utils.dataset_utils import get_dataset_path
|
|
from utils.model_utils import get_model_config_from_file
|
|
from utils.training_utils import get_training_config_from_file, get_data_ids, find_max_batch_size
|
|
|
|
|
|
def load_model(model_dir: str,
|
|
results_base_dir: str,
|
|
lmdb_base_dir: str,
|
|
device: str = "cuda") -> tuple:
|
|
"""
|
|
Load the model from the model directory.
|
|
:param model_dir: model directory
|
|
:param results_base_dir: results base directory
|
|
:param lmdb_base_dir: LMDB base directory
|
|
:param device: device to use
|
|
:return: model configuration, training configuration, dataset directory
|
|
"""
|
|
|
|
training_id = None
|
|
model_configuration = get_model_config_from_file(model_dir, results_base_dir, lmdb_base_dir)
|
|
feature_config = model_configuration["feature_config"]
|
|
dataset_dir = get_dataset_path(lmdb_base_dir, feature_config["feature_set_name"])
|
|
|
|
# get list of trainings
|
|
training_base_dir = os.path.join(model_dir, "trainings")
|
|
# if no training id is given, use the latest training
|
|
if training_id is None:
|
|
training_dirs = os.listdir(training_base_dir)
|
|
training_dirs = [os.path.join(training_base_dir, dir) for dir in training_dirs if
|
|
os.path.isdir(os.path.join(training_base_dir, dir))]
|
|
training_dirs = sorted(training_dirs, key=os.path.getmtime, reverse=True)
|
|
training_dir = training_dirs[0]
|
|
else:
|
|
training_dir = os.path.join(training_base_dir, training_id)
|
|
|
|
training_configuration = get_training_config_from_file(training_dir, results_base_dir, model_configuration)
|
|
|
|
limit = None
|
|
train_ids, val_ids, test_ids = get_data_ids(model_configuration, training_configuration, dataset_dir, limit)
|
|
|
|
dataset_dir = f"{lmdb_base_dir}/{feature_config['feature_set_name']}"
|
|
env = lmdb.open(dataset_dir, readonly=True)
|
|
load_fn = model_configuration["model_load_fn"]
|
|
model = load_fn(model_configuration, training_configuration, test_ids[0], device, lmdb_env=env)
|
|
|
|
# estimate batch size
|
|
sample_batch = get_collated_batch_for_key(train_ids[0], model_configuration, lmdb_env=env)
|
|
sample_input = sample_batch[0][0] # get the first item in the batch
|
|
input_shape = sample_input.shape
|
|
# send model to proper device
|
|
model.to(device)
|
|
batch_size = find_max_batch_size(model,
|
|
input_shape,
|
|
device=device)
|
|
training_configuration["batch_size"] = batch_size
|
|
return model, model_configuration, training_configuration, train_ids, val_ids, test_ids
|
|
|
|
|
|
def get_training(model_configuration: dict,
|
|
results_base_dir: str,
|
|
training_id: str = None) -> dict:
|
|
"""
|
|
Get the training configuration from the model directory.
|
|
:param model_configuration: model configuration
|
|
:param training_id: training name
|
|
:return: training configuration
|
|
"""
|
|
model_dir = model_configuration["model_dir"]
|
|
# get list of trainings
|
|
training_base_dir = os.path.join(model_dir, "trainings")
|
|
# if no training id is given, use the latest training
|
|
if training_id is None:
|
|
training_dirs = os.listdir(training_base_dir)
|
|
training_dirs = [os.path.join(training_base_dir, dir) for dir in training_dirs if
|
|
os.path.isdir(os.path.join(training_base_dir, dir))]
|
|
training_dirs = sorted(training_dirs, key=os.path.getmtime, reverse=True)
|
|
training_dir = training_dirs[0]
|
|
else:
|
|
training_dir = os.path.join(training_base_dir, training_id)
|
|
|
|
training_configuration = get_training_config_from_file(training_dir, results_base_dir, model_configuration)
|
|
return training_configuration
|
|
|
|
|
|
def find_models(base_dir: str, current_dir: str = None) -> dict:
|
|
"""
|
|
Find all models in the base directory.
|
|
:param base_dir: base directory
|
|
:return: list of model directories
|
|
"""
|
|
model_dirs = dict()
|
|
is_root = base_dir == current_dir
|
|
if current_dir is None:
|
|
current_dir = base_dir
|
|
for sub_dir_name in os.listdir(current_dir):
|
|
sub_dir = os.path.join(current_dir, sub_dir_name)
|
|
if not os.path.isdir(sub_dir):
|
|
continue
|
|
model_config_path = os.path.join(sub_dir, "model_configuration.pickle")
|
|
|
|
if current_dir not in model_dirs:
|
|
model_dirs[current_dir] = list()
|
|
if os.path.exists(model_config_path) and is_root:
|
|
model_dirs[current_dir].append(sub_dir)
|
|
elif os.path.exists(model_config_path) and not is_root:
|
|
model_dirs[current_dir].append(sub_dir)
|
|
else:
|
|
model_dirs = model_dirs | find_models(base_dir, sub_dir)
|
|
|
|
return model_dirs
|
|
|
|
|
|
def apply_sigmoid_if_necessary(model_outputs: np.ndarray | torch.Tensor,
|
|
model_configuration: dict) -> np.ndarray:
|
|
"""
|
|
Applies a sigmoid to the outputs of a model, if they use a BCE loss.
|
|
Args:
|
|
model_outputs: models outputs
|
|
model_configuration: configuration of the model that produced the outputs
|
|
|
|
Returns:
|
|
processed_model_outputs: processed model outputs with sigmoid applied where necessary
|
|
"""
|
|
|
|
target_features = model_configuration["feature_config"]["target_features"]
|
|
ignored_features = model_configuration["feature_config"]["ignored_features"]
|
|
used_targets = [feature for feature in target_features if feature["name"] not in ignored_features]
|
|
|
|
is_3d = len(model_outputs.shape) == 3
|
|
|
|
processed_outputs = np.empty_like(model_outputs)
|
|
for i, used_target in enumerate(used_targets):
|
|
if used_target["loss_fn"] == nn.BCEWithLogitsLoss:
|
|
if is_3d:
|
|
target_values = model_outputs[:, 0, i]
|
|
if isinstance(target_values, np.ndarray):
|
|
target_values = torch.from_numpy(target_values)
|
|
processed_outputs[:, 0, i] = torch.sigmoid(target_values).numpy()
|
|
else:
|
|
target_values = model_outputs[:, i]
|
|
if isinstance(target_values, np.ndarray):
|
|
target_values = torch.from_numpy(target_values)
|
|
processed_outputs[:, i] = torch.sigmoid(target_values).numpy()
|
|
else:
|
|
if is_3d:
|
|
processed_outputs[:, 0, i] = model_outputs[:, 0, i].numpy() if isinstance(model_outputs,
|
|
torch.Tensor) else model_outputs[
|
|
:, 0, i]
|
|
else:
|
|
processed_outputs[:, i] = model_outputs[:, i].numpy() if isinstance(model_outputs,
|
|
torch.Tensor) else model_outputs[:,
|
|
i]
|
|
|
|
return processed_outputs
|
|
|
|
|
|
def scale_features(features: np.ndarray | torch.Tensor,
|
|
feature_names: list | None,
|
|
model_configuration: dict,
|
|
scalers: dict = None,
|
|
subset_name: str = "train") -> np.ndarray:
|
|
"""
|
|
Scales the given features based on the scalers associated with the given model configuration
|
|
Args:
|
|
features: feature arrays, may be up to 3D (but may only have 2 usable dimensions, (B, 1, F) is allowed)
|
|
feature_names: names of the features in the same order as in the input array, may be None, then no scaling is performed
|
|
model_configuration: model configuration of the model the features belong to
|
|
scalers: dict of scales associated by name, will be fetched from model configuration, if None
|
|
subset_name: subset name, either "train", "test" or "val", used for determining the correct scalers
|
|
|
|
Returns:
|
|
features: scaled features
|
|
"""
|
|
|
|
if features.shape[-1] != len(feature_names):
|
|
raise ValueError("Number of features does not match number of given feature names")
|
|
|
|
# remove singleton dimension, if 3D
|
|
if len(features.shape) == 3:
|
|
if features.shape[1] != 1:
|
|
raise ValueError(f"Input has 3 dimensions, expected (B, 1, F), but got {features.shape}")
|
|
features = features[:, 0, :]
|
|
|
|
if scalers is None:
|
|
scalers = get_scalers_for_model(model_configuration)
|
|
|
|
if isinstance(features, torch.Tensor):
|
|
features = features.detach().cpu().numpy()
|
|
|
|
scaled_features = np.empty_like(features)
|
|
for i in range(len(feature_names)):
|
|
# skip scaling, if feature name is None
|
|
if feature_names[i] is not None:
|
|
scaled_output = inverse_scale_feature(features[:, i],
|
|
feature_names[i],
|
|
scalers,
|
|
subset_name=subset_name)
|
|
else:
|
|
scaled_output = features[:, i]
|
|
|
|
# make sure to make into 1D array when assigning
|
|
scaled_features[:, i] = scaled_output.ravel()
|
|
|
|
return scaled_features
|