Alex Blank
577 lines
23 KiB
Python
577 lines
23 KiB
Python
import copy
|
|
import logging
|
|
import os
|
|
import random
|
|
|
|
import lmdb
|
|
import numpy as np
|
|
from bson import ObjectId
|
|
from torch.utils.data import IterableDataset
|
|
from tqdm import tqdm
|
|
|
|
from utils.dataset_creation import get_features, load_scalers, scale_item, combine_features
|
|
from utils.dataset_utils import load_key_stats
|
|
from utils.lmdb_utils import load_from_lmdb
|
|
from vsm_datascience_common.cycle_database_connection.cycle_data import get_cycle_by_id
|
|
from vsm_datascience_common.cycle_database_connection.db_utils import get_cycles_collection
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def get_prepared_sequence_length(input_sequence: np.ndarray, model_configuration: dict) -> int:
|
|
"""
|
|
Returns the length of the input sequence after padding and resampling
|
|
"""
|
|
|
|
initial_length = len(input_sequence)
|
|
take_every_nth = model_configuration["preprocessing"]["take_every_nth"]
|
|
padding_length = get_padding_length(model_configuration, resampled=False)
|
|
raw_length = initial_length + padding_length
|
|
adjusted = int(raw_length // take_every_nth)
|
|
return adjusted
|
|
|
|
|
|
def get_padding_length(model_configuration: dict, resampled: bool = True) -> int:
|
|
"""
|
|
Returns the length of the input sequence after padding and resampling
|
|
"""
|
|
take_every_nth = model_configuration["preprocessing"]["take_every_nth"]
|
|
padding_length = int(model_configuration["input_window_length"] * (
|
|
1 - model_configuration["preprocessing"]["min_input_length_fraction_for_padding"])) * (
|
|
take_every_nth if not resampled else 1)
|
|
return padding_length
|
|
|
|
|
|
def get_number_of_windows(base_length: int, model_configuration: dict) -> int:
|
|
"""
|
|
Returns the number of windows for a given sequence length
|
|
"""
|
|
window_shift = model_configuration["preprocessing"]["window_shift"]
|
|
input_window_length = model_configuration["input_window_length"]
|
|
output_window_length = model_configuration["output_window_length"]
|
|
output_window_offset = model_configuration["output_window_offset"]
|
|
|
|
if input_window_length > output_window_length + output_window_offset:
|
|
return (base_length - input_window_length) // window_shift + 1
|
|
else:
|
|
return (base_length - output_window_length - output_window_offset) // window_shift + 1
|
|
|
|
|
|
def get_collated_batch_for_key(sample_key: ObjectId | str,
|
|
model_configuration: dict,
|
|
start_cutoff: int = None,
|
|
end_cutoff: int = None,
|
|
lmdb_env=None) -> dict:
|
|
sample_batch_for_key = get_batch_for_key(sample_key, model_configuration,
|
|
start_cutoff=start_cutoff,
|
|
end_cutoff=end_cutoff,
|
|
lmdb_env=lmdb_env)
|
|
collated = model_configuration["collate_fn"](sample_batch_for_key)
|
|
return collated
|
|
|
|
|
|
def get_batch_for_key(key,
|
|
model_configuration: dict,
|
|
start_cutoff: int = None,
|
|
end_cutoff: int = None,
|
|
lmdb_env=None) -> np.ndarray:
|
|
"""
|
|
Get the batch for a given key from the lmdb database or compute it directly.
|
|
|
|
The return has the same format as the batch produced by the batch_fn in the model configuration.
|
|
Args:
|
|
key: key to get the batch for, can also be object id from database
|
|
model_configuration: model configuration to use
|
|
start_cutoff: start cutoff for the batch, if None, the whole batch is used, CAUTION: cutoff should not be normalized -> in raw data points
|
|
end_cutoff: end cutoff for the batch, if None, the whole batch is used, CAUTION: cutoff should not be normalized -> in raw data points
|
|
lmdb_env: the lmdb environment to use, if None, the features are computed directly from the database
|
|
|
|
Returns:
|
|
batch: batch for the given key, as returned by the batch_fn in the model configuration
|
|
|
|
"""
|
|
if lmdb_env is None:
|
|
# compute and scale features, ignored features are handled internally
|
|
features = get_scaled_feature_for_key(key, model_configuration)
|
|
else:
|
|
# load features from lmdb
|
|
features = load_from_lmdb(lmdb_env, str(key))
|
|
# ignored features are handled internally, thus we need to remove them here
|
|
ignored_features = model_configuration["feature_config"]["ignored_features"] if "ignored_features" in \
|
|
model_configuration[
|
|
"feature_config"] else []
|
|
# add scaled versions of features
|
|
ignored_features = ignored_features + [feat + "_scaled" for feat in ignored_features]
|
|
for feature_set in features:
|
|
for feature_name in list(features[feature_set].keys()):
|
|
if feature_name in ignored_features:
|
|
del features[feature_set][feature_name]
|
|
|
|
processed_chunk = process_chunk([key], [features], model_configuration)
|
|
|
|
# use cutoff, if provided
|
|
if start_cutoff is not None or end_cutoff is not None:
|
|
|
|
item_length = len(processed_chunk[0]["target_features"])
|
|
|
|
if start_cutoff is None:
|
|
start_cutoff = 0
|
|
if end_cutoff is None:
|
|
end_cutoff = item_length
|
|
|
|
# "normalize" cutoff to adjust for added padding and resampling
|
|
padding_length = get_padding_length(model_configuration, resampled=False)
|
|
take_every_nth = model_configuration["preprocessing"]["take_every_nth"]
|
|
# make sure that there is no padding added to the start cutoff, so that it includes the padding added to the sequence
|
|
start_cutoff_normalized = max(int(start_cutoff // take_every_nth), 0)
|
|
# end cutoff must be adjusted to include the padding added to the sequence
|
|
end_cutoff_normalized = min(int((end_cutoff + padding_length) // take_every_nth), item_length)
|
|
for feature_set in processed_chunk[0]:
|
|
if feature_set in model_configuration["feature_config"]["feature_sets"]:
|
|
processed_chunk[0][feature_set] = processed_chunk[0][feature_set][
|
|
start_cutoff_normalized:end_cutoff_normalized + 1]
|
|
|
|
if "batch_fn" in model_configuration and model_configuration["batch_fn"] is not None:
|
|
batch = model_configuration["batch_fn"](processed_chunk, model_configuration)
|
|
else:
|
|
batch = processed_chunk
|
|
if batch is None or len(batch) == 0:
|
|
raise ValueError(f"Batch is empty for key {key}")
|
|
return batch
|
|
|
|
|
|
def get_scaled_feature_for_key(key: str,
|
|
model_configuration: dict) -> tuple:
|
|
feature_config = model_configuration["feature_config"]
|
|
|
|
user_cycles = list(
|
|
get_cycles_collection().find({"user_id": ObjectId(key)} | feature_config["filter_criteria"]).sort("starts_at",
|
|
1))
|
|
cycle_features = list()
|
|
for cycle in user_cycles:
|
|
features = get_features(cycle, feature_config)
|
|
cycle_features.append(features)
|
|
|
|
# combine all cycles for the user
|
|
if len(cycle_features) == 0:
|
|
raise ValueError(f"No features found for key {key}")
|
|
features = combine_features(cycle_features, feature_config)
|
|
scaler_dir = os.path.join(feature_config["dataset_dir"], "scalers")
|
|
scalers = load_scalers(scaler_dir)
|
|
# scale features
|
|
scaled_features = scale_item(features, scalers)
|
|
return scaled_features
|
|
|
|
|
|
def augment_items(identifiers: list,
|
|
items: list[dict],
|
|
feature_config: dict,
|
|
lmdb_env: lmdb.Environment = None) -> list:
|
|
"""
|
|
Augments the given sequences by attaching previous sequences
|
|
Args:
|
|
identifiers: identifiers of items for finding previous sequences
|
|
items: actual items
|
|
feature_config: feature config of the dataset
|
|
lmdb_env: lmdb environment for loading previous sequences, can be None, in this case the items are computed directly
|
|
|
|
Returns:
|
|
list of augmented items
|
|
"""
|
|
|
|
if "augmentation" not in feature_config or \
|
|
feature_config["augmentation"]["use_augmentation"] is False:
|
|
return items
|
|
|
|
augmented_items = list()
|
|
for i, item in enumerate(items):
|
|
item_id = identifiers[i]
|
|
try:
|
|
item_data = get_cycles_collection().find_one({"_id": ObjectId(item_id)}, {"user_id": 1, "starts_at": 1})
|
|
item_user_id = item_data["user_id"]
|
|
item_starts_at = item_data["starts_at"]
|
|
except KeyError:
|
|
print(f"User ID not found for item {item_id}")
|
|
continue
|
|
previous_item_identifiers = list(get_cycles_collection().aggregate(
|
|
# make sure to filter by user first to significantly reduce the number of items
|
|
[
|
|
{
|
|
"$match": {
|
|
"user_id": item_user_id,
|
|
"starts_at": {"$lt": item_starts_at},
|
|
}
|
|
}
|
|
] + feature_config["filter_criteria_pipeline"] + [
|
|
{
|
|
"$sort": {
|
|
"starts_at": -1
|
|
}
|
|
},
|
|
{
|
|
"$project": {
|
|
"_id": 1,
|
|
"starts_at": 1
|
|
}
|
|
}
|
|
]
|
|
))
|
|
max_lookback = feature_config["augmentation"]["max_lookback"]
|
|
previous_items = list()
|
|
for previous_item in previous_item_identifiers:
|
|
if len(previous_items) >= max_lookback:
|
|
break
|
|
previous_item_id = previous_item["_id"]
|
|
if lmdb_env is None:
|
|
# compute item
|
|
cycle = get_cycle_by_id(previous_item_id)
|
|
features = get_features(cycle, feature_config)
|
|
scaler_dir = os.path.join(feature_config["dataset_dir"], "scalers")
|
|
scalers = load_scalers(scaler_dir)
|
|
previous_item = scale_item(features, scalers)
|
|
previous_items.append(previous_item)
|
|
else:
|
|
# take item from lmdb
|
|
try:
|
|
# make sure to parse object id
|
|
previous_item = load_from_lmdb(lmdb_env, str(previous_item_id))
|
|
except KeyError:
|
|
continue
|
|
previous_items.append(previous_item)
|
|
|
|
# merge items into one
|
|
feature_sets = feature_config["feature_sets"]
|
|
augmented_item = copy.deepcopy(item)
|
|
for feature_set in feature_sets:
|
|
if feature_set not in augmented_item:
|
|
continue
|
|
if "static" in feature_set:
|
|
continue
|
|
for feature_name in augmented_item[feature_set]:
|
|
if feature_name not in augmented_item[feature_set]:
|
|
continue
|
|
augmented_item[feature_set][feature_name] = np.concatenate(
|
|
[previous_item[feature_set][feature_name] for previous_item in previous_items]
|
|
+ [augmented_item[feature_set][feature_name]])
|
|
|
|
augmented_items.append(augmented_item)
|
|
|
|
return augmented_items
|
|
|
|
|
|
def process_chunk(ids: list,
|
|
chunk: list,
|
|
model_configuration: dict,
|
|
ignored_features: list = None,
|
|
pad_sequences: bool = True,
|
|
statics_as_list: bool = True):
|
|
"""
|
|
Creates a chunk of data as dataframe for training and inference of a tft model.
|
|
:param ids: list of ids of the data points in the chunk
|
|
:param chunk: list of data points
|
|
|
|
:return: dataframe with the data points in the chunk
|
|
"""
|
|
|
|
if ignored_features is None:
|
|
ignored_features = list()
|
|
|
|
input_features = list(chunk[0]["target_features"].keys())
|
|
feature_config = model_configuration["feature_config"]
|
|
|
|
# variables for padding
|
|
padding_value = 0
|
|
|
|
# create dataframe from items in chunk
|
|
data = []
|
|
for i, item in enumerate(chunk):
|
|
item_length = item["target_features"][input_features[0]].shape[0]
|
|
take_every_nth = model_configuration["preprocessing"]["take_every_nth"]
|
|
|
|
data_item = dict()
|
|
for feature_set in feature_config["feature_sets"]:
|
|
if feature_set in item and len(item[feature_set]) > 0:
|
|
# fill in the data
|
|
source = item[feature_set]
|
|
features_to_use = [feature for feature in source.keys() if "_scaled" in feature]
|
|
feature_set_data = [source[feature] for feature in features_to_use]
|
|
|
|
padding_length = get_padding_length(model_configuration, resampled=False)
|
|
|
|
if len(feature_set_data) == 0:
|
|
raise ValueError(f"No features found for key {ids[i]} in feature set {feature_set}")
|
|
|
|
if len(feature_set_data[0]) == 1:
|
|
# if the feature is constant, we need to repeat it for all time points
|
|
if statics_as_list:
|
|
if pad_sequences:
|
|
desired_length = item_length + padding_length
|
|
feature_set_data = np.array(
|
|
[np.full((desired_length,), val) for val in feature_set_data]).T[::take_every_nth]
|
|
else:
|
|
feature_set_data = np.array(
|
|
[np.full((item_length,), val) for val in feature_set_data]).T[::take_every_nth]
|
|
else:
|
|
feature_set_data = np.array(feature_set_data).flatten()
|
|
else:
|
|
if pad_sequences:
|
|
padding_values = np.full((padding_length, len(feature_set_data)),
|
|
[padding_value for vals in feature_set_data]).T
|
|
feature_set_data = np.concatenate([padding_values, np.array(feature_set_data)], axis=1)
|
|
# create bins of n values and apply mean
|
|
|
|
# check, if special accumulation function has been specified
|
|
individual_feature_data = list()
|
|
for i, feature_to_use in enumerate(features_to_use):
|
|
try:
|
|
# get config for feature, make sure to replace scaled to get actual config
|
|
current_feature_config = \
|
|
[x for x in feature_config[feature_set] if
|
|
x["name"] == feature_to_use.replace("_scaled", "")][0]
|
|
current_accumulation_fn = current_feature_config["accumulation_fn"]
|
|
except (KeyError, IndexError):
|
|
# default to mean
|
|
current_accumulation_fn = np.mean
|
|
|
|
current_feature_data = apply_fn_to_bins(feature_set_data[i], take_every_nth,
|
|
current_accumulation_fn)
|
|
individual_feature_data.append(current_feature_data)
|
|
|
|
# stack individual features
|
|
feature_set_data = np.stack(individual_feature_data).T
|
|
|
|
if feature_set not in data_item:
|
|
data_item[feature_set] = feature_set_data
|
|
else:
|
|
data_item[feature_set] = np.concatenate([data_item[feature_set], feature_set_data], axis=1)
|
|
|
|
data.append(data_item)
|
|
|
|
return data
|
|
|
|
|
|
def apply_fn_to_bins(input_sequence: np.ndarray,
|
|
bin_size: int,
|
|
fn: callable) -> np.ndarray:
|
|
"""
|
|
Applies a function to the bins of the input sequence along the last dimension and returns the results.
|
|
Args:
|
|
input_sequence: sequence to apply the function to (1D or 2D)
|
|
bin_size: size of the bins
|
|
fn: function to apply to the bins
|
|
|
|
Returns:
|
|
np.ndarray: result of the function applied to the bins
|
|
|
|
"""
|
|
if input_sequence.ndim == 1:
|
|
input_sequence = input_sequence.reshape(1, -1)
|
|
dims = 1
|
|
else:
|
|
dims = input_sequence.shape[-1]
|
|
|
|
result = []
|
|
for row in input_sequence:
|
|
row_result = []
|
|
for start in range(0, len(row), bin_size):
|
|
end = min(start + bin_size, len(row))
|
|
bin_slice = row[start:end]
|
|
row_result.append(fn(bin_slice))
|
|
result.append(row_result)
|
|
|
|
if dims > 1:
|
|
return np.array(result)
|
|
else:
|
|
return np.array(result[0])
|
|
|
|
|
|
def produce_window_batches(data_chunk: list,
|
|
model_configuration: dict,
|
|
offsets: list | np.ndarray = None) -> list:
|
|
"""
|
|
Produces batches of windows from the data chunk.
|
|
Args:
|
|
data_chunk: data chunk
|
|
model_configuration: configuration to use
|
|
offsets: offsets used to determine the length of data to use, uses item[-offset:] of data
|
|
|
|
Returns:
|
|
data: list of windows
|
|
"""
|
|
input_window_length = model_configuration["input_window_length"]
|
|
output_window_length = model_configuration["output_window_length"]
|
|
output_window_offset = model_configuration["output_window_offset"]
|
|
window_shift = model_configuration["preprocessing"]["window_shift"]
|
|
|
|
if offsets is None:
|
|
offsets = np.zeros(len(data_chunk), dtype=int)
|
|
|
|
data = list()
|
|
for i, item in enumerate(data_chunk):
|
|
item_length = item["target_features"][-offsets[i]:].shape[0]
|
|
total_window_length = input_window_length if input_window_length > output_window_length + output_window_offset \
|
|
else output_window_length + output_window_offset
|
|
|
|
if item_length < total_window_length:
|
|
continue
|
|
|
|
# number of windows is defined by the input window length and output window length with offset
|
|
num_windows = get_number_of_windows(item_length, model_configuration)
|
|
for j in range(num_windows):
|
|
window = dict()
|
|
for key, value in item.items():
|
|
if value is None:
|
|
window[key] = None
|
|
else:
|
|
if key == "target_features":
|
|
window[key] = value[-offsets[i]:][j * window_shift + output_window_offset:
|
|
j * window_shift + output_window_length + output_window_offset]
|
|
else:
|
|
window[key] = value[-offsets[i]:][j * window_shift:j * window_shift + input_window_length]
|
|
|
|
data.append(window)
|
|
|
|
return data
|
|
|
|
|
|
def produce_simple_batches(data_chunk: list,
|
|
model_configuration: dict,
|
|
offsets: list | np.ndarray = None) -> list:
|
|
"""
|
|
Produces batches without windowing.
|
|
Args:
|
|
data_chunk: data chunk
|
|
model_configuration: configuration to use
|
|
offsets: offsets used to determine the length of data to use, uses item[-offset:] of data
|
|
|
|
Returns:
|
|
data: list of items
|
|
"""
|
|
pass
|
|
|
|
|
|
class LMDBIterableDataset(IterableDataset):
|
|
def __init__(self,
|
|
lmdb_env_path: str,
|
|
lmdb_keys: list[str],
|
|
model_configuration: dict,
|
|
batch_size: int = 32):
|
|
self.lmdb_path = lmdb_env_path
|
|
self.lmdb_env = None
|
|
self.lmdb_keys = lmdb_keys
|
|
self.key_subset = None
|
|
self.model_configuration = model_configuration
|
|
self.batch_size = batch_size
|
|
self.keys_stats = load_key_stats(model_configuration["feature_config"]["dataset_dir"])
|
|
self.len = None
|
|
|
|
try:
|
|
self.model_configuration["batch_fn"]
|
|
except KeyError:
|
|
raise KeyError("Batch function not found in model configuration")
|
|
|
|
def set_key_subset(self, key_subset: list[str]):
|
|
"""
|
|
Set the key subset to use for the dataset.
|
|
Args:
|
|
key_subset: list of keys to use
|
|
"""
|
|
self.key_subset = key_subset
|
|
# reset length
|
|
self.len = None
|
|
|
|
def get_length_of_data_subset(self, key_set: list[str]):
|
|
"""
|
|
Get the length of the data subset.
|
|
Args:
|
|
key_set: list of keys to use
|
|
"""
|
|
# run simplified version of __iter__ to get the length
|
|
|
|
num_steps = 0
|
|
i = 0
|
|
num_batch_items = 0
|
|
while True:
|
|
if num_batch_items >= self.batch_size:
|
|
num_steps += 1
|
|
num_batch_items -= self.batch_size
|
|
else:
|
|
if i >= len(key_set):
|
|
if num_batch_items > 0:
|
|
num_steps += 1
|
|
break
|
|
key = key_set[i]
|
|
i += 1
|
|
|
|
# try to fetch stats from key_stats
|
|
try:
|
|
current_key_stats = self.keys_stats["by_key"][key]
|
|
base_length = current_key_stats["item_length"]
|
|
take_every_nth = self.model_configuration["preprocessing"]["take_every_nth"]
|
|
padding_length = get_padding_length(self.model_configuration, resampled=False)
|
|
raw_length = base_length + padding_length
|
|
item_length = int(raw_length // take_every_nth)
|
|
except:
|
|
item = load_from_lmdb(self.lmdb_env, key)
|
|
item_length = get_prepared_sequence_length(
|
|
item["target_features"][list(item["target_features"])[0]],
|
|
self.model_configuration)
|
|
num_windows = get_number_of_windows(item_length, self.model_configuration)
|
|
num_batch_items += num_windows
|
|
|
|
return num_steps
|
|
|
|
def __iter__(self):
|
|
self.init_lmdb_env()
|
|
|
|
# if no key subset is set, use all keys
|
|
if self.key_subset is None:
|
|
self.key_subset = self.lmdb_keys
|
|
|
|
random.shuffle(self.key_subset)
|
|
|
|
batch = list()
|
|
counter = 0
|
|
collate_fn = self.model_configuration["collate_fn"]
|
|
while True:
|
|
if len(batch) >= self.batch_size:
|
|
if collate_fn is None:
|
|
yield batch[:self.batch_size]
|
|
else:
|
|
yield collate_fn(batch[:self.batch_size])
|
|
batch = batch[self.batch_size:]
|
|
else:
|
|
if counter >= len(self.key_subset):
|
|
if len(batch) > 0:
|
|
yield collate_fn(batch)
|
|
break
|
|
key = self.lmdb_keys[counter]
|
|
counter += 1
|
|
try:
|
|
current_batch = get_batch_for_key(key, self.model_configuration, lmdb_env=self.lmdb_env)
|
|
except ValueError:
|
|
# if betch is empty, try next
|
|
continue
|
|
batch.extend(current_batch)
|
|
|
|
def init_lmdb_env(self):
|
|
if self.lmdb_env is None:
|
|
self.lmdb_env = lmdb.open(self.lmdb_path,
|
|
readonly=True,
|
|
lock=False,
|
|
readahead=False,
|
|
meminit=False)
|
|
|
|
def __len__(self):
|
|
self.init_lmdb_env()
|
|
|
|
# if no key subset is set, use all keys
|
|
if self.key_subset is None:
|
|
self.key_subset = self.lmdb_keys
|
|
|
|
if self.len is None:
|
|
num_steps = self.get_length_of_data_subset(self.key_subset)
|
|
self.len = num_steps
|
|
|
|
return self.len
|