added code
This commit is contained in:
+63
-22
@@ -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()
|
||||
|
||||
|
||||
Reference in New Issue
Block a user