added code

This commit is contained in:
2025-09-10 10:37:55 +02:00
parent 36901c736d
commit c78a68de80
199 changed files with 3561 additions and 22579 deletions
+63 -22
View File
@@ -2,11 +2,12 @@ import copy
import logging
import os
import random
from datetime import datetime
import lmdb
import numpy as np
from bson import ObjectId
from torch.utils.data import IterableDataset
from torch.utils.data import IterableDataset, get_worker_info
from tqdm import tqdm
from utils.dataset_creation import get_features, load_scalers, scale_item, combine_features
@@ -52,20 +53,24 @@ def get_number_of_windows(base_length: int, model_configuration: dict) -> int:
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
return (base_length - input_window_length) // max(window_shift, 1) + 1
else:
return (base_length - output_window_length - output_window_offset) // window_shift + 1
return (base_length - output_window_length - output_window_offset) // max(window_shift, 1) + 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:
lmdb_env=None,
item_subset_source: str = "train",
for_inference: bool = False) -> 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)
lmdb_env=lmdb_env,
item_subset_source=item_subset_source,
for_inference=for_inference)
collated = model_configuration["collate_fn"](sample_batch_for_key)
return collated
@@ -74,7 +79,9 @@ def get_batch_for_key(key,
model_configuration: dict,
start_cutoff: int = None,
end_cutoff: int = None,
lmdb_env=None) -> np.ndarray:
lmdb_env=None,
item_subset_source: str = "train",
for_inference: bool = False) -> np.ndarray:
"""
Get the batch for a given key from the lmdb database or compute it directly.
@@ -85,6 +92,8 @@ def get_batch_for_key(key,
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
item_subset_source: where the item "came" from, test, train or val, used to determine the scaler to use.
for_inference: whether to compute batched features directly from the database and ignore cycle filters
Returns:
batch: batch for the given key, as returned by the batch_fn in the model configuration
@@ -92,7 +101,10 @@ def get_batch_for_key(key,
"""
if lmdb_env is None:
# compute and scale features, ignored features are handled internally
features = get_scaled_feature_for_key(key, model_configuration)
features = get_scaled_feature_for_key(key,
model_configuration,
item_subset_source=item_subset_source,
for_inference=for_inference) # TODO: implement proper cutoff usage, otherwise the whole user history will be loaded every time
else:
# load features from lmdb
features = load_from_lmdb(lmdb_env, str(key))
@@ -141,12 +153,17 @@ def get_batch_for_key(key,
def get_scaled_feature_for_key(key: str,
model_configuration: dict) -> tuple:
model_configuration: dict,
item_subset_source: str = "train",
for_inference: bool = False) -> 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))
if for_inference:
user_cycles = list(
get_cycles_collection().find({"user_id": ObjectId(key)}).sort("starts_at", 1))
else:
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)
@@ -159,7 +176,7 @@ def get_scaled_feature_for_key(key: str,
scaler_dir = os.path.join(feature_config["dataset_dir"], "scalers")
scalers = load_scalers(scaler_dir)
# scale features
scaled_features = scale_item(features, scalers)
scaled_features = scale_item(features, subset_name=item_subset_source, scalers=scalers)
return scaled_features
@@ -476,7 +493,7 @@ class LMDBIterableDataset(IterableDataset):
Args:
key_subset: list of keys to use
"""
self.key_subset = key_subset
self.key_subset = self.get_worker_keyset(key_subset)
# reset length
self.len = None
# reset lmdb env
@@ -484,6 +501,26 @@ class LMDBIterableDataset(IterableDataset):
self.lmdb_env.close()
self.lmdb_env = None
def get_worker_keyset(self, key_set: list[str]):
"""
Adjust the key set to the current worker.
Args:
key_set: list of keys to use
"""
worker_info = get_worker_info()
if worker_info is not None:
# in a worker process
num_workers = worker_info.num_workers
worker_id = worker_info.id
# split the key set into chunks for each worker
chunk_size = len(key_set) // num_workers
start = worker_id * chunk_size
end = (worker_id + 1) * chunk_size if worker_id != num_workers - 1 else len(key_set)
return key_set[start:end]
else:
# use the whole key set
return key_set
def get_length_of_data_subset(self, key_set: list[str]):
"""
Get the length of the data subset.
@@ -495,6 +532,7 @@ class LMDBIterableDataset(IterableDataset):
num_steps = 0
i = 0
num_batch_items = 0
padding_length = get_padding_length(self.model_configuration, resampled=False)
while True:
if num_batch_items >= self.batch_size:
num_steps += 1
@@ -512,9 +550,8 @@ class LMDBIterableDataset(IterableDataset):
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)
item_length = len(range(0, raw_length, take_every_nth))
except:
item = load_from_lmdb(self.lmdb_env, key)
item_length = get_prepared_sequence_length(
@@ -529,7 +566,7 @@ class LMDBIterableDataset(IterableDataset):
# if no key subset is set, use all keys
if self.key_subset is None:
self.key_subset = self.lmdb_keys
self.set_key_subset(self.lmdb_keys)
self.init_lmdb_env()
@@ -548,14 +585,18 @@ class LMDBIterableDataset(IterableDataset):
else:
if counter >= len(self.key_subset):
if len(batch) > 0:
yield collate_fn(batch)
if collate_fn is None:
yield batch
else:
yield collate_fn(batch)
break
key = self.lmdb_keys[counter]
key = self.key_subset[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
except ValueError as e:
# if batch is empty, try next
continue
batch.extend(current_batch)
@@ -570,7 +611,7 @@ class LMDBIterableDataset(IterableDataset):
def __len__(self):
# if no key subset is set, use all keys
if self.key_subset is None:
self.key_subset = self.lmdb_keys
self.set_key_subset(self.lmdb_keys)
self.init_lmdb_env()