added code
This commit is contained in:
Alex Blank
@@ -0,0 +1,576 @@
|
||||
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
|
||||
Reference in New Issue
Block a user