diff --git a/.github/workflows/python.yml b/.github/workflows/python.yml index d55f9adc6c0..d679aa74a92 100644 --- a/.github/workflows/python.yml +++ b/.github/workflows/python.yml @@ -167,7 +167,7 @@ jobs: pip install --upgrade pip wheel setuptools # Use CUDA 12.1 wheels to avoid slow/source builds pip install --extra-index-url https://download.pytorch.org/whl/cu121 \ - torch==2.4.1 torchvision==0.19.1 + torch==2.4.1 torchvision==0.19.1 torchaudio==2.4.1 pip install \ transformers \ opencv-python \ @@ -179,7 +179,10 @@ jobs: fvcore \ scikit-optimize \ flair \ - optuna + optuna \ + openface-test \ + imagebind \ + "pytorchvideo @ git+https://github.com/facebookresearch/pytorchvideo.git@eb04d1b" kill $KA cd src/main/python python -m unittest discover -s tests/scuro -p 'test_*.py' -v diff --git a/src/main/python/systemds/scuro/__init__.py b/src/main/python/systemds/scuro/__init__.py index 168f036b1e3..c8ba06e72d9 100644 --- a/src/main/python/systemds/scuro/__init__.py +++ b/src/main/python/systemds/scuro/__init__.py @@ -24,6 +24,7 @@ from systemds.scuro.dataloader.video_loader import VideoLoader from systemds.scuro.dataloader.text_loader import TextLoader from systemds.scuro.dataloader.json_loader import JSONLoader +from systemds.scuro.dataloader.tabular_loader import TabularLoader from systemds.scuro.representations.representation import Representation from systemds.scuro.representations.aggregate import Aggregation from systemds.scuro.representations.aggregated_representation import ( @@ -71,9 +72,23 @@ from systemds.scuro.representations.representation_dataloader import JSON from systemds.scuro.representations.representation_dataloader import Pickle from systemds.scuro.representations.resnet import ResNet +from systemds.scuro.representations.openface import OpenFace + +try: + from systemds.scuro.representations.image_bind import ImageBind +except ImportError as _imagebind_import_error: # pragma: no cover + import warnings as _warnings + + _warnings.warn( + f"ImageBind representation unavailable, it is excluded from the search " + f"({_imagebind_import_error})." + ) + ImageBind = None + from systemds.scuro.representations.spectrogram import Spectrogram from systemds.scuro.representations.sum import Sum from systemds.scuro.representations.swin_video_transformer import SwinVideoTransformer +from systemds.scuro.representations.tabular_features import TabularFeatures from systemds.scuro.representations.tfidf import TfIdf from systemds.scuro.representations.unimodal import UnimodalRepresentation from systemds.scuro.representations.wav2vec import Wav2Vec @@ -82,6 +97,10 @@ DynamicWindow, StaticWindow, ) +from systemds.scuro.representations.physiological_window import ( + AdaptiveWindow, + PhysiologicalEventWindow, +) from systemds.scuro.representations.word2vec import W2V from systemds.scuro.representations.x3d import X3D from systemds.scuro.representations.color_histogram import ColorHistogram @@ -124,6 +143,28 @@ MLPLearnedDimReduction, ) + +from systemds.scuro.representations.physiological_representations import ( + SDNN, + RMSSD, + pNN, + RRPerMinute, + HRVBandPower, + HRVVLF, + HRVLF, + HRVHF, + HRVLFHF, + PoincareSD1, + PoincareSD2, + # SampleEntropy, + # DFAAlpha, + SCLSlope, + SCLDynamicRange, + SCRPeaksPerMinute, + SCRAverageAmplitude, + SCRAverageDuration, +) + __all__ = [ "BaseLoader", "ImageLoader", @@ -156,6 +197,8 @@ "JSON", "Pickle", "ResNet", + "OpenFace", + "ImageBind", "Spectrogram", "Sum", "BoW", @@ -187,6 +230,8 @@ "AttentionFusion", "DynamicWindow", "StaticWindow", + "AdaptiveWindow", + "PhysiologicalEventWindow", "Min", "Max", "Mean", @@ -211,4 +256,22 @@ "MLPAveraging", "MLPLearnedDimReduction", "DimensionalityReduction", + "MultimodalGAPymooOptimizer", + "SDNN", + "RMSSD", + "pNN", + "RRPerMinute", + "HRVBandPower", + "HRVVLF", + "HRVLF", + "HRVHF", + "HRVLFHF", + "PoincareSD1", + "PoincareSD2", + # "DFAAlpha", + "SCLSlope", + "SCLDynamicRange", + "SCRPeaksPerMinute", + "SCRAverageAmplitude", + "SCRAverageDuration", ] diff --git a/src/main/python/systemds/scuro/dataloader/base_loader.py b/src/main/python/systemds/scuro/dataloader/base_loader.py index 9b89c773942..e58a755838c 100644 --- a/src/main/python/systemds/scuro/dataloader/base_loader.py +++ b/src/main/python/systemds/scuro/dataloader/base_loader.py @@ -20,12 +20,33 @@ # ------------------------------------------------------------- import os from abc import ABC, abstractmethod -from typing import Iterator, List, Optional, Tuple, Union +from collections.abc import Sequence +from typing import Callable, Iterator, List, Optional, Tuple, Union import math +from numbers import Integral import numpy as np +class LazyFileSequence(Sequence): + """List-like file references decoded only when a sample is requested.""" + + def __init__(self, file_names: List[str], decoder: Callable[[str], object]): + self.file_names = tuple(file_names) + self.decoder = decoder + + def __getitem__(self, index): + if isinstance(index, slice): + return [self[i] for i in range(*index.indices(len(self)))] + return self.decoder(self.file_names[index]) + + def __len__(self): + return len(self.file_names) + + def subset(self, indices): + return type(self)([self.file_names[i] for i in indices], self.decoder) + + class BaseLoader(ABC): def __init__( self, @@ -54,8 +75,7 @@ def __init__( self._data_type = data_type self._ext = ext self.stats = None - if chunk_size: - self.chunk_size = chunk_size + self.chunk_size = chunk_size @property def chunk_size(self): @@ -63,8 +83,28 @@ def chunk_size(self): @chunk_size.setter def chunk_size(self, value): - self._chunk_size = value - self._num_chunks = int(math.ceil(len(self.indices) / self._chunk_size)) + if value is None: + self._chunk_size = None + self._num_chunks = 1 + else: + if isinstance(value, bool) or not isinstance(value, Integral) or value <= 0: + raise ValueError("chunk_size must be None or a positive integer") + self._chunk_size = int(value) + self._num_chunks = int(math.ceil(len(self.indices) / self._chunk_size)) + + stats = getattr(self, "stats", None) + if stats is not None and hasattr(stats, "num_instances"): + stats.num_instances = ( + len(self.indices) + if self._chunk_size is None + else min(len(self.indices), self._chunk_size) + ) + if stats is not None and hasattr(stats, "num_total_instances"): + stats.num_total_instances = len(self.indices) + + @property + def is_chunked(self): + return self._chunk_size is not None @property def num_chunks(self): @@ -91,7 +131,7 @@ def load(self): """ Takes care of loading the raw data either chunk wise (if chunk size is defined) or all at once """ - if self._chunk_size: + if self.is_chunked: return self._load_next_chunk() return self._load(self.indices) @@ -102,8 +142,8 @@ def iter_loaded_chunks( if reset: self.reset() - if not self._chunk_size: - data, metadata = self._load(self.indices) + if not self.is_chunked: + data, metadata = self.load() yield data, metadata, self.indices return @@ -115,24 +155,21 @@ def iter_loaded_chunks( yield data, metadata, chunk_indices def update_chunk_sizes(self, other): - if not self._chunk_size and not other.chunk_size: + sizes = [ + size for size in (self.chunk_size, other.chunk_size) if size is not None + ] + if not sizes: return - - if ( - self._chunk_size - and not other.chunk_size - or self._chunk_size < other.chunk_size - ): - other.chunk_size = self.chunk_size - else: - self.chunk_size = other.chunk_size + shared_size = min(sizes) + self.chunk_size = shared_size + other.chunk_size = shared_size def _load_next_chunk(self): """ Loads the next chunk of data """ self.data = [] - # TODO: Handle metadata correctly + self.metadata = [] next_chunk_indices = self.indices[ self._next_chunk * self._chunk_size : (self._next_chunk + 1) @@ -161,7 +198,7 @@ def get_file_names(self, indices=None): if self._ext is None: _, self._ext = os.path.splitext(os.listdir(self.source_path)[0]) for index in self.indices if indices is None else indices: - file_names.append(self.source_path + index + self._ext) + file_names.append(os.path.join(self.source_path, index + self._ext)) return file_names else: return self.source_path diff --git a/src/main/python/systemds/scuro/dataloader/image_loader.py b/src/main/python/systemds/scuro/dataloader/image_loader.py index 25e8690cf5a..8da151a113e 100644 --- a/src/main/python/systemds/scuro/dataloader/image_loader.py +++ b/src/main/python/systemds/scuro/dataloader/image_loader.py @@ -24,7 +24,7 @@ import numpy as np -from systemds.scuro.dataloader.base_loader import BaseLoader +from systemds.scuro.dataloader.base_loader import BaseLoader, LazyFileSequence import cv2 from systemds.scuro.modality.type import ModalityType @@ -55,13 +55,28 @@ def __init__( source_path, indices, data_type, chunk_size, ModalityType.IMAGE, ext ) self.load_data_from_file = load + self._all_metadata = [] self.stats = self.get_stats(source_path) - def extract(self, file: str, index: Optional[Union[str, List[str]]] = None): - self.file_sanity_check(file) + def load(self): + if self.chunk_size: + return super().load() + self.data = LazyFileSequence( + self.get_file_names(self.indices), self._decode_file + ) + self.metadata = self._all_metadata.copy() + return self.data, self.metadata + + def _decode_file(self, file: str): + self.file_sanity_check(file) image = cv2.imread(file, cv2.IMREAD_COLOR) - image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) + if image is None: + raise ValueError(f"Could not read image at path: {file}") + return cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.uint8, copy=False) + + def extract(self, file: str, index: Optional[Union[str, List[str]]] = None): + image = self._decode_file(file) if image.ndim == 2: height, width = image.shape @@ -69,8 +84,6 @@ def extract(self, file: str, index: Optional[Union[str, List[str]]] = None): else: height, width, channels = image.shape - image = image.astype(np.uint8, copy=False) - self.metadata.append( self.modality_type.create_metadata(width, height, channels) ) @@ -78,6 +91,7 @@ def extract(self, file: str, index: Optional[Union[str, List[str]]] = None): self.data.append(image) def get_stats(self, source_path: str): + self._all_metadata = [] max_width = 0 max_height = 0 max_channels = 0 @@ -87,18 +101,15 @@ def get_stats(self, source_path: str): average_channels = 0 for file in self.indices: path = os.path.join(source_path, f"{file}{self._ext}") - # if self.chunk_size is None: - # self.extract(path) - # md = self.metadata[path] - # max_width = max(max_width, md["width"]) - # max_height = max(max_height, md["height"]) - # max_channels = max(max_channels, md["num_channels"]) - # num_instances += 1 - # else: self.file_sanity_check(path) image = cv2.imread(path, cv2.IMREAD_COLOR) + if image is None: + raise ValueError(f"Could not read image at path: {path}") image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) height, width, channels = image.shape + self._all_metadata.append( + self.modality_type.create_metadata(width, height, channels) + ) max_width = max(max_width, width) max_height = max(max_height, height) max_channels = max(max_channels, channels) @@ -119,3 +130,15 @@ def get_stats(self, source_path: str): average_height, average_channels, ) + + def estimate_peak_memory_bytes(self) -> dict: + n = self.chunk_size if self.chunk_size is not None else 1 + per_instance = ( + self.stats.average_width + * self.stats.average_height + * self.stats.average_channels + ) + return { + "cpu_peak_bytes": int(n * per_instance * np.dtype(np.uint8).itemsize), + "gpu_peak_bytes": 0, + } diff --git a/src/main/python/systemds/scuro/dataloader/video_loader.py b/src/main/python/systemds/scuro/dataloader/video_loader.py index bf7bdd846c7..e2ee7b72d24 100644 --- a/src/main/python/systemds/scuro/dataloader/video_loader.py +++ b/src/main/python/systemds/scuro/dataloader/video_loader.py @@ -19,12 +19,12 @@ # # ------------------------------------------------------------- from dataclasses import dataclass -import os -from typing import List, Optional, Union +import math +from typing import List, Optional, Tuple, Union import numpy as np -from systemds.scuro.dataloader.base_loader import BaseLoader +from systemds.scuro.dataloader.base_loader import BaseLoader, LazyFileSequence import cv2 from systemds.scuro.modality.type import ModalityType @@ -39,6 +39,7 @@ class VideoStats: max_channels: int num_instances: int num_total_instances: int + shape_variance: float = 0.0 @property def output_shape(self): @@ -52,6 +53,15 @@ def output_shape(self): """ return (self.max_length, self.max_height, self.max_width, self.max_channels) + @property + def avg_output_shape(self): + return ( + max(1, int(round(self.avg_length))), + self.max_height, + self.max_width, + self.max_channels, + ) + class VideoLoader(BaseLoader): def __init__( @@ -62,109 +72,181 @@ def __init__( chunk_size: Optional[int] = None, load=True, fps=None, + target_size: Optional[Tuple[int, int]] = None, ): super().__init__( source_path, indices, data_type, chunk_size, ModalityType.VIDEO ) self.load_data_from_file = load self.fps = fps + self.target_size = tuple(int(v) for v in target_size) if target_size else None + self._all_metadata = [] self.stats = self.get_stats(source_path) - def extract(self, file: str, index: Optional[Union[str, List[str]]] = None): + def load(self): + if self.chunk_size: + return super().load() + + self.data = LazyFileSequence( + self.get_file_names(self.indices), self._decode_data + ) + self.metadata = self._all_metadata.copy() + return self.data, self.metadata + + def _decode_data(self, file: str): + return self._decode_file(file)[0] + + def _frame_interval(self, source_fps: float) -> int: + if self.fps and source_fps and self.fps < source_fps: + return max(1, int(round(source_fps / self.fps))) + return 1 + + def _stored_length(self, source_length: int, source_fps: float) -> int: + interval = self._frame_interval(source_fps) + return int(math.ceil(source_length / interval)) if source_length > 0 else 0 + + def _stored_frame_size(self, width: int, height: int) -> Tuple[int, int]: + return self.target_size if self.target_size else (width, height) + + def _fit_frame(self, frame: np.ndarray) -> np.ndarray: + if self.target_size is None: + return frame + + target_w, target_h = self.target_size + height, width = frame.shape[:2] + if (width, height) == (target_w, target_h): + return frame + + scale = max(target_w / width, target_h / height) + new_w = max(target_w, int(round(width * scale))) + new_h = max(target_h, int(round(height * scale))) + interpolation = cv2.INTER_AREA if scale < 1.0 else cv2.INTER_LINEAR + frame = cv2.resize(frame, (new_w, new_h), interpolation=interpolation) + + left = (new_w - target_w) // 2 + top = (new_h - target_h) // 2 + return frame[top : top + target_h, left : left + target_w] + + def _decode_file(self, file: str): self.file_sanity_check(file) cap = cv2.VideoCapture(file) if not cap.isOpened(): - raise f"Could not read video at path: {file}" - - orig_fps = cap.get(cv2.CAP_PROP_FPS) - frame_interval = 1 - if self.fps is not None and self.fps < orig_fps: - frame_interval = int(round(orig_fps / self.fps)) - else: - self.fps = orig_fps - - length = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) - width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) - height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) - num_channels = 3 - - self.metadata.append( - self.modality_type.create_metadata( - self.fps, length, width, height, num_channels - ) + raise ValueError(f"Could not read video at path: {file}") + + try: + source_fps = cap.get(cv2.CAP_PROP_FPS) + frame_interval = self._frame_interval(source_fps) + stored_fps = source_fps / frame_interval if source_fps else source_fps + + scale_denominator = np.dtype(self._data_type).type(255.0) + frames = [] + idx = 0 + while True: + ret, frame = cap.read() + if not ret: + break + if idx % frame_interval == 0: + frame = self._fit_frame(frame) + frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + frame = frame.astype(self._data_type) / scale_denominator + frames.append(frame) + idx += 1 + finally: + cap.release() + + if not frames: + raise ValueError(f"No frames could be decoded from {file}") + + data = np.stack(frames) + + num_frames, height, width = data.shape[0], data.shape[1], data.shape[2] + metadata = self.modality_type.create_metadata( + stored_fps, num_frames, width, height, data.shape[3] ) + return data, metadata - frames = [] - idx = 0 - while cap.isOpened(): - ret, frame = cap.read() - - if not ret: - break - if idx % frame_interval == 0: - frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) - frame = frame.astype(self._data_type, copy=False) / 255.0 - frames.append(frame) - idx += 1 - - self.data.append(np.stack(frames)) + def extract(self, file: str, index: Optional[Union[str, List[str]]] = None): + data, metadata = self._decode_file(file) + self.metadata.append(metadata) + self.data.append(data) def get_stats(self, source_path: str): + self._all_metadata = [] self.file_sanity_check(source_path) - fps = 0 max_length = 0 max_width = 0 max_height = 0 max_num_channels = 0 num_instances = 0 - avg_length = 0 - for file in os.listdir(source_path): - file_name = file.split(".")[0] - if file_name not in self.indices: - continue - self.file_sanity_check(source_path + file) - cap = cv2.VideoCapture(source_path + file) - - length = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) - width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) - height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + stored_lengths = [] + stored_fps = [] + + for file in self.get_file_names(self.indices): + self.file_sanity_check(file) + cap = cv2.VideoCapture(file) + if not cap.isOpened(): + raise ValueError(f"Could not read video at path: {file}") + try: + source_length = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + source_fps = cap.get(cv2.CAP_PROP_FPS) + width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + finally: + cap.release() + + length = self._stored_length(source_length, source_fps) + width, height = self._stored_frame_size(width, height) num_channels = 3 + stored_frequency = ( + source_fps / self._frame_interval(source_fps) if source_fps else 0 + ) + self._all_metadata.append( + self.modality_type.create_metadata( + stored_frequency, length, width, height, num_channels + ) + ) + max_length = max(max_length, length) - avg_length += length max_width = max(max_width, width) max_height = max(max_height, height) max_num_channels = max(max_num_channels, num_channels) + stored_lengths.append(length) + if source_fps: + stored_fps.append(stored_frequency) num_instances += 1 + num_total_instances = num_instances + avg_length = float(np.mean(stored_lengths)) if stored_lengths else 0.0 + shape_variance = ( + float(np.std(stored_lengths) / avg_length) + if stored_lengths and avg_length > 0 + else 0.0 + ) num_instances = ( min(num_instances, self.chunk_size) if self.chunk_size is not None else num_instances ) return VideoStats( - fps, + float(np.mean(stored_fps)) if stored_fps else 0, max_length, - avg_length / num_instances, + avg_length, max_width, max_height, max_num_channels, num_instances, num_total_instances, + shape_variance, ) def estimate_peak_memory_bytes(self) -> dict: - s = self.stats - if self.chunk_size is not None: - n = self.chunk_size - else: - n = s.num_instances + stats = self.stats + n = self.chunk_size if self.chunk_size is not None else 1 + n = min(n, stats.num_total_instances) + per_instance = int(np.prod(stats.avg_output_shape)) + itemsize = np.dtype(self._data_type).itemsize return { - "cpu_peak_bytes": n - * s.output_shape[0] - * s.output_shape[1] - * s.output_shape[2] - * s.output_shape[3] - * 4, + "cpu_peak_bytes": int(n * per_instance * itemsize), "gpu_peak_bytes": 0, } diff --git a/src/main/python/systemds/scuro/drsearch/dag_group_executor.py b/src/main/python/systemds/scuro/drsearch/dag_group_executor.py index 57007cb06be..5181f934f4e 100644 --- a/src/main/python/systemds/scuro/drsearch/dag_group_executor.py +++ b/src/main/python/systemds/scuro/drsearch/dag_group_executor.py @@ -25,7 +25,7 @@ import time from typing import Any, Dict, List, Optional -from systemds.scuro import Modality +from systemds.scuro.modality.modality import Modality from systemds.scuro.drsearch.representation_dag import ( LRUCache, RepresentationDag, diff --git a/src/main/python/systemds/scuro/drsearch/dag_group_scheduler.py b/src/main/python/systemds/scuro/drsearch/dag_group_scheduler.py index def12219f54..23799f78a0c 100644 --- a/src/main/python/systemds/scuro/drsearch/dag_group_scheduler.py +++ b/src/main/python/systemds/scuro/drsearch/dag_group_scheduler.py @@ -28,10 +28,10 @@ def get_peak_memory_from_dag_group( dag_group: List[RepresentationDag], modality: Modality -) -> tuple[float, float]: +) -> Tuple[float, float]: peak_memory_cpu = 0.0 peak_memory_gpu = 0.0 - leaf_memory_bytes = modality.estimate_memory_bytes() + leaf_memory_bytes = modality.estimate_peak_memory_bytes()["cpu_peak_bytes"] for dag in dag_group: prev_stats = modality.get_stats() for node in dag.nodes[1:]: diff --git a/src/main/python/systemds/scuro/drsearch/modality_shared_memory.py b/src/main/python/systemds/scuro/drsearch/modality_shared_memory.py index e68d7195a2d..57a5048cda3 100644 --- a/src/main/python/systemds/scuro/drsearch/modality_shared_memory.py +++ b/src/main/python/systemds/scuro/drsearch/modality_shared_memory.py @@ -275,6 +275,12 @@ def __getstate__(self): state["_buffer"] = None return state + def __reduce_ex__(self, protocol): + return ( + type(self), + (self.shm_name, self.dtype_str, self.offsets, self.total_elems), + ) + def _is_string_list_shared_memory_candidate(data: Any) -> bool: if not isinstance(data, list) or not data: diff --git a/src/main/python/systemds/scuro/drsearch/multimodal_optimizer.py b/src/main/python/systemds/scuro/drsearch/multimodal_optimizer.py index 596cab0237b..42369e423be 100644 --- a/src/main/python/systemds/scuro/drsearch/multimodal_optimizer.py +++ b/src/main/python/systemds/scuro/drsearch/multimodal_optimizer.py @@ -33,7 +33,6 @@ from systemds.scuro.representations.aggregated_representation import ( AggregatedRepresentation, ) -from systemds.scuro.representations.aggregate import Aggregation from systemds.scuro.drsearch.operator_registry import Registry from systemds.scuro.utils.schema_helpers import get_shape @@ -69,10 +68,9 @@ def _evaluate_dag_worker(dag_pickle, task_pickle, modalities_pickle, debug=False from systemds.scuro.representations.aggregated_representation import ( AggregatedRepresentation, ) - from systemds.scuro.representations.aggregate import Aggregation if task.expected_dim == 1 and get_shape(final_representation.metadata) > 1: - agg_operator = AggregatedRepresentation(Aggregation()) + agg_operator = AggregatedRepresentation() final_representation = agg_operator.transform(final_representation) eval_start = time.time() @@ -424,7 +422,7 @@ def _evaluate_dag(self, dag: RepresentationDag, task: Task) -> "OptimizationResu return None if task.expected_dim == 1 and get_shape(fused_representation.metadata) > 1: - agg_operator = AggregatedRepresentation(Aggregation()) + agg_operator = AggregatedRepresentation() fused_representation = agg_operator.transform(fused_representation) eval_start = time.time() diff --git a/src/main/python/systemds/scuro/drsearch/node_executor.py b/src/main/python/systemds/scuro/drsearch/node_executor.py index 3400e6c6c50..c56a6f4ca1f 100644 --- a/src/main/python/systemds/scuro/drsearch/node_executor.py +++ b/src/main/python/systemds/scuro/drsearch/node_executor.py @@ -26,7 +26,7 @@ import torch -from systemds.scuro import Modality +from systemds.scuro.modality.modality import Modality from systemds.scuro.drsearch.modality_result_cache import RefCountResultCache from systemds.scuro.drsearch.modality_shared_memory import ( add_shared_memory_candidate, @@ -66,6 +66,24 @@ _MAX_NODE_RETRIES = int(os.environ.get("SCURO_MAX_NODE_RETRIES", "3")) +def _place_on_device(obj: Any, gpu_id: Optional[int]) -> None: + """Apply scheduler placement even when an object selected CUDA itself.""" + device = torch.device("cpu" if gpu_id is None else f"cuda:{gpu_id}") + if hasattr(obj, "gpu_id"): + try: + obj.gpu_id = gpu_id + except (AttributeError, RuntimeError): + pass + if hasattr(obj, "device"): + try: + obj.device = device + except (AttributeError, RuntimeError): + pass + model = getattr(obj, "model", None) + if model is not None and hasattr(model, "to"): + obj.model = model.to(device) + + def _run_gpu_op(fn, gpu_id: Optional[int]): if gpu_id is None or not torch.cuda.is_available(): return fn() @@ -143,13 +161,11 @@ def _execute_node_worker(node, input_mods: List[Any], gpu_id: Optional[int]): torch.cuda.reset_peak_memory_stats(device) node_operation = _instantiate_operation(node) + _place_on_device(node_operation, gpu_id) operation_name = node_operation.name if DEBUG: print(f"Executing node {node.node_id} {operation_name} on GPU {gpu_id}") - if gpu_id is not None and hasattr(node_operation, "gpu_id"): - node_operation.gpu_id = gpu_id - def _run_node_op(): if len(input_mods) == 1: if isinstance(node_operation, Context): @@ -227,8 +243,7 @@ def _execute_task_worker( torch.cuda.set_device(device) torch.cuda.reset_peak_memory_stats(device) - if gpu_id is not None and hasattr(task, "model") and hasattr(task.model, "device"): - task.model.device = torch.device(f"cuda:{gpu_id}") + _place_on_device(getattr(task, "model", task), gpu_id) def _run_task(): start = time.perf_counter() @@ -271,17 +286,27 @@ def _run_task(): def _execute_leaf_batch_worker(nodes: List[Any], modality: Any, gpu_id: Optional[int]): - node_id_by_representation = {} - def _run(): representations = [] + representation_keys = [] + aggregations = [] for node in nodes: operation = node.operation(params=node.parameters) - if hasattr(operation, "gpu_id"): - operation.gpu_id = gpu_id + _place_on_device(operation, gpu_id) representations.append(operation) - node_id_by_representation[operation.name] = node.node_id - return modality.apply_representations(representations, parallel=True) + representation_keys.append(node.node_id) + pushdown_config = node.parameters.get("_pushdown_aggregation") + aggregations.append( + AggregatedRepresentation(params=pushdown_config) + if pushdown_config is not None + else None + ) + return modality.apply_representations( + representations, + parallel=True, + representation_keys=representation_keys, + aggregations=aggregations, + ) modality_results = _run_gpu_op(_run, gpu_id) shm_info = {} @@ -297,8 +322,13 @@ def _run(): } return { "results": modality_results, - "node_id_by_representation": node_id_by_representation, "shm_info": shm_info, + "failed_nodes": { + node_id: f"{type(error).__name__}: {error}" + for node_id, error in getattr( + modality, "failed_representations", {} + ).items() + }, } @@ -309,7 +339,7 @@ def _load_leaf_worker(modality: Any) -> Dict[str, Any]: data = modality.data resident_bytes = 0 try: - resident_bytes = modality.estimate_memory_bytes() + resident_bytes = modality.estimate_peak_memory_bytes()["cpu_peak_bytes"] except Exception: resident_bytes = 0 @@ -324,7 +354,14 @@ def _load_leaf_worker(modality: Any) -> Dict[str, Any]: def _dispatch_node(payload, gpu_id): node, input_mods = payload - return _execute_node_worker(node, input_mods, gpu_id) + try: + return _execute_node_worker(node, input_mods, gpu_id) + finally: + # CSE executes each node once; retaining its model after completion + # makes scheduler reservations lie about persistent worker memory. + _WORKER_OP_CACHE.clear() + if gpu_id is not None: + cleanup_gpu(gpu_id) def _dispatch_task(payload, gpu_id): @@ -434,8 +471,19 @@ def __init__( _WORKER_DISPATCH, ctx=create_mp_context(), threads_per_worker=threads_per_worker, + gpu_devices=list(getattr(self.scheduler, "gpu_devices", [])), + gpu_slots_per_device=int( + os.environ.get("SCURO_GPU_SLOTS_PER_DEVICE", "1") + ), + gpu_demand_fraction=self.scheduler.gpu_demand_fraction(), ) self._pool = worker_pool + restrict_devices = getattr(self.scheduler, "restrict_gpu_devices", None) + if callable(restrict_devices): + restrict_devices( + getattr(worker_pool, "gpu_worker_devices", []), + getattr(worker_pool, "gpu_slots_per_device", 1), + ) def _requeue_or_give_up(self, node_id: str, reason: str) -> bool: attempts = self._node_attempts.get(node_id, 0) + 1 @@ -474,6 +522,8 @@ def _retain_for_submit(self, parent_ids: List[str], payload: Any) -> List[str]: def _load_leaf_modalities(self) -> None: for modality in self._modalities: + if self._loads_in_chunks(modality): + continue if getattr(modality, "has_data", None) and modality.has_data(): continue attempts = 0 @@ -494,12 +544,19 @@ def _load_leaf_modalities(self) -> None: if shm_name is not None: self._leaf_shm_names.append(shm_name) + @staticmethod + def _loads_in_chunks(modality: Any) -> bool: + data_loader = getattr(modality, "data_loader", None) + return getattr(data_loader, "chunk_size", None) is not None + def _cleanup_leaf_shared_memory(self) -> None: for shm_name in self._leaf_shm_names: unlink_shm(shm_name) self._leaf_shm_names = [] - def _submit_node(self, node_id: str) -> None: + def _submit_node( + self, node_id: str, allow_gpu_worker_for_cpu: bool = False + ) -> None: node = self.scheduler.mapping[node_id] gpu_id = node.gpu_id parent_ids = self.scheduler.get_valid_parents(node_id) @@ -523,18 +580,26 @@ def _submit_node(self, node_id: str) -> None: "task", (node_id, self._tasks[task_idx], payload, node.aggregation), gpu_id=gpu_id, + allow_gpu_worker_for_cpu=allow_gpu_worker_for_cpu, ) else: payload = self._modalities if parent_results is None else parent_results retained = self._retain_for_submit(parent_ids, payload) self.scheduler.begin_execution(node_id) self.scheduler.move_to_running(node_id) - job_id = self._pool.submit("node", (node, payload), gpu_id=gpu_id) + job_id = self._pool.submit( + "node", + (node, payload), + gpu_id=gpu_id, + allow_gpu_worker_for_cpu=allow_gpu_worker_for_cpu, + ) self._job_units[job_id] = _NodeUnit(node_id) self._job_retained_shm[job_id] = retained - def _submit_leaf_batch(self, node_ids: List[str]) -> None: + def _submit_leaf_batch( + self, node_ids: List[str], allow_gpu_worker_for_cpu: bool = False + ) -> None: nodes = [self.scheduler.mapping[nid] for nid in node_ids] gpu_id = nodes[0].gpu_id retained = self._retain_for_submit([], self._modalities[0].data) @@ -542,22 +607,42 @@ def _submit_leaf_batch(self, node_ids: List[str]) -> None: self.scheduler.begin_execution(nid) self.scheduler.move_to_running(node_ids) job_id = self._pool.submit( - "leaf_batch", (nodes, self._modalities[0]), gpu_id=gpu_id + "leaf_batch", + (nodes, self._modalities[0]), + gpu_id=gpu_id, + allow_gpu_worker_for_cpu=allow_gpu_worker_for_cpu, ) self._job_units[job_id] = _BatchUnit(node_ids) self._job_retained_shm[job_id] = retained def _fill_pipeline(self) -> None: ready = self.scheduler.get_runnable().copy() + + def _gpu_id(entry): + node_id = entry[0] if isinstance(entry, list) else entry + return self.scheduler.mapping[node_id].gpu_id + + ready.sort(key=lambda entry: _gpu_id(entry) is None) for entry in ready: - if not self._pool.has_idle_worker: - break + gpu_id = _gpu_id(entry) + allow_gpu_worker_for_cpu = gpu_id is None + if not self._pool.has_idle_worker_for( + gpu_id, + allow_gpu_worker_for_cpu=allow_gpu_worker_for_cpu, + ): + continue if isinstance(entry, list): - self._submit_leaf_batch(entry) + self._submit_leaf_batch( + entry, + allow_gpu_worker_for_cpu=allow_gpu_worker_for_cpu, + ) else: if not self.scheduler.can_start_now(entry): continue - self._submit_node(entry) + self._submit_node( + entry, + allow_gpu_worker_for_cpu=allow_gpu_worker_for_cpu, + ) def _record_stats(self, node_id: str, pid: int, start_time: float, end_time: float): node_stats = self.statistics["node_stats"] @@ -648,23 +733,25 @@ def _handle_node_success(self, node_id: str, value: Dict[str, Any]) -> None: def _handle_batch_success(self, value: Dict[str, Any]) -> None: results = value["results"] - node_id_by_representation = value["node_id_by_representation"] shm_info = value.get("shm_info", {}) - for representation, transformed_modality in results.items(): - node_id = node_id_by_representation[representation] - info = shm_info.get(representation, {}) + for node_id, transformed_modality in results.items(): + info = shm_info.get(node_id, {}) self._handle_modality_result( transformed_modality, node_id, None, None, - representation, + self.scheduler.mapping[node_id].operation.__name__, actual_stats=info.get("actual_stats"), shm_name=info.get("shm_name"), resident_bytes=info.get("resident_bytes"), shm_bytes=info.get("shm_bytes", 0), ) + for node_id, reason in value.get("failed_nodes", {}).items(): + if not self._requeue_or_give_up(node_id, reason): + self._release_parents(node_id) + def _handle_modality_result( self, transformed_modality: Any, diff --git a/src/main/python/systemds/scuro/drsearch/node_scheduler.py b/src/main/python/systemds/scuro/drsearch/node_scheduler.py index 1ca681e88ad..8398c82a81b 100644 --- a/src/main/python/systemds/scuro/drsearch/node_scheduler.py +++ b/src/main/python/systemds/scuro/drsearch/node_scheduler.py @@ -70,6 +70,14 @@ def __init__( self.nodes = self._get_nodes_from_dags(nodes) self.unresolved_parents = self._get_unresolved_parents() self.node_resources = self._estimate_node_resources() + self.node_costs = self._estimate_node_costs() + self.node_priorities = self._compute_upward_ranks() + self.gpu_devices = list(self.memory_budget["gpu"]) + self.gpu_slots_per_device = max( + 1, int(os.environ.get("SCURO_GPU_SLOTS_PER_DEVICE", "1")) + ) + self.gpu_slots_in_use = {gpu_id: 0 for gpu_id in self.gpu_devices} + self._gpu_slot_nodes: Dict[str, int] = {} self.success = False self.deadlock = False self.ready_nodes = [] @@ -91,9 +99,7 @@ def __init__( for node_id in self.topo_order if node_id not in self.leaves and self.unresolved_parents[node_id] == 0 } - self.n_gpu = ( - torch.cuda.device_count() if torch and torch.cuda.is_available() else 0 - ) + self.n_gpu = len(self.gpu_devices) leaf_cached = sum(self.node_resources[node][0] for node in self.leaves) self.memory_stats = { "cpu_cached": leaf_cached, @@ -134,18 +140,25 @@ def get_runnable(self) -> List[RepresentationNode]: runnable_nodes = self._get_runnable_nodes() admitted_bytes = self._pending_admitted_cpu_bytes() + pending_gpu = {gpu_id: 0 for gpu_id in self.gpu_devices} + pending_slots = {gpu_id: 0 for gpu_id in self.gpu_devices} for node in runnable_nodes: if node in self._ready_set: continue - ok, gpu_id = self._check_memory_constraints(node, admitted_bytes) + ok, gpu_id = self._check_memory_constraints( + node, admitted_bytes, pending_gpu, pending_slots + ) if ok: admitted_bytes += self.node_resources[node][0] + if gpu_id is not None: + pending_gpu[gpu_id] += self.node_resources[node][1] + pending_slots[gpu_id] += 1 self.mapping[node].gpu_id = gpu_id self._candidates.discard(node) self.ready_nodes.append(node) self._ready_set.add(node) - contains_leaf = [] + chunked_leaf_by_device = defaultdict(list) for node in self.ready_nodes: if isinstance(node, list): continue @@ -156,14 +169,16 @@ def get_runnable(self) -> List[RepresentationNode]: == self.mapping[self.mapping[node].inputs[0]].modality_id ): if mod.data_loader.chunk_size is not None: - contains_leaf.append(node) + chunked_leaf_by_device[self.mapping[node].gpu_id].append( + node + ) break - for node in contains_leaf: - self.ready_nodes.remove(node) + for nodes_for_device in chunked_leaf_by_device.values(): + for node in nodes_for_device: + self.ready_nodes.remove(node) - if len(contains_leaf) > 0: - self.ready_nodes.append(contains_leaf) + self.ready_nodes.extend(chunked_leaf_by_device.values()) return self.ready_nodes def _get_runnable_nodes(self) -> List[str]: @@ -177,11 +192,64 @@ def _score(node_id: str): and self.remaining_children.get(parent_id, 0) == 1 ): release_bytes += self.node_resources[parent_id][0] - return (-release_bytes, node_id not in self.roots, node_id) + return ( + -self.node_priorities.get(node_id, 0.0), + -release_bytes, + node_id not in self.roots, + node_id, + ) runnable_nodes.sort(key=_score) return runnable_nodes + def _estimate_node_costs(self) -> Dict[str, float]: + costs: Dict[str, float] = {} + for node_id in self.topo_order: + if node_id in self.leaves: + costs[node_id] = 1.0 + continue + if node_id in self.roots: + parent_ids = list(self.parents.get(node_id, set())) + parent_stats = [ + self.node_stats.get(parent_id) for parent_id in parent_ids + ] + input_stats = ( + parent_stats[0] if len(parent_stats) == 1 else parent_stats + ) + task_idx = self.mapping[node_id].parameters.get("_task_idx", 0) + try: + costs[node_id] = max( + 1.0, + float(self.tasks[task_idx].estimate_relative_cost(input_stats)), + ) + except Exception: + costs[node_id] = 1.0 + else: + cpu_bytes, gpu_bytes = self.node_resources.get(node_id, (0, 0)) + costs[node_id] = max(1.0, float(cpu_bytes + gpu_bytes) / (1024**2)) + return costs + + def _compute_upward_ranks(self) -> Dict[str, float]: + ranks: Dict[str, float] = {} + for node_id in reversed(self.topo_order): + downstream = [ranks[child] for child in self.children.get(node_id, set())] + ranks[node_id] = float(self.node_costs.get(node_id, 1.0)) + ( + max(downstream) if downstream else 0.0 + ) + return ranks + + def gpu_demand_fraction(self) -> float: + total = 0.0 + gpu_total = 0.0 + for node_id, resources in self.node_resources.items(): + weight = float(self.node_costs.get(node_id, 1.0)) + if weight <= 0: + weight = 1.0 + total += weight + if resources[1] > 0: + gpu_total += weight + return gpu_total / total if total else 0.0 + def add_failed_node(self, node_id: str, reason: str = "unknown failure"): self.failed_nodes.append(node_id) self.failed_node_reasons[node_id] = reason @@ -196,8 +264,10 @@ def requeue_node(self, node_id: str) -> None: def begin_execution(self, node_id: str) -> None: gpu_id = self.mapping[node_id].gpu_id cpu_mem, gpu_mem = self.node_resources[node_id] - if gpu_id is not None and gpu_mem > 0: + if gpu_id is not None and gpu_mem > 0 and node_id not in self._gpu_slot_nodes: self.memory_stats["gpu_in_use"][gpu_id] += gpu_mem + self.gpu_slots_in_use[gpu_id] += 1 + self._gpu_slot_nodes[node_id] = gpu_id if cpu_mem > 0 and node_id not in self._cpu_reserved_nodes: self._cpu_reserved_nodes[node_id] = int(cpu_mem) self.memory_stats["cpu_in_flight"] += int(cpu_mem) @@ -382,7 +452,13 @@ def not_enough_memory(self) -> bool: return True return False - def _check_memory_constraints(self, node_id: str, pending_bytes: int = 0) -> bool: + def _check_memory_constraints( + self, + node_id: str, + pending_bytes: int = 0, + pending_gpu: Optional[Dict[int, int]] = None, + pending_slots: Optional[Dict[int, int]] = None, + ) -> bool: cpu_mem, gpu_mem = self.node_resources[node_id] gpu_id = None if ( @@ -400,9 +476,13 @@ def _check_memory_constraints(self, node_id: str, pending_bytes: int = 0) -> boo return False, None if gpu_mem > 0.0 and self.n_gpu > 0: - gpu_id = self._gpu_with_most_free_memory(gpu_mem) + gpu_id = self._gpu_with_most_free_memory( + gpu_mem, pending_gpu, pending_slots + ) if gpu_id is None: + if self._waiting_for_gpu_slot(gpu_mem, pending_gpu, pending_slots): + return False, None attempts = self.gpu_wait_attempts.get(node_id, 0) + 1 self.gpu_wait_attempts[node_id] = attempts if attempts > _MAX_GPU_SCHEDULE_ATTEMPTS: @@ -421,25 +501,68 @@ def _check_memory_constraints(self, node_id: str, pending_bytes: int = 0) -> boo return True, gpu_id - def _gpu_with_most_free_memory(self, memory_needed): - free_memory = [] - for i in range(self.n_gpu): - free_memory.append( - self.memory_budget["gpu"][i] - self.memory_stats["gpu_in_use"][i] - ) - - if max(free_memory) < memory_needed: + def _gpu_with_most_free_memory( + self, + memory_needed: int, + pending_gpu: Optional[Dict[int, int]] = None, + pending_slots: Optional[Dict[int, int]] = None, + ) -> Optional[int]: + pending_gpu = pending_gpu or {} + pending_slots = pending_slots or {} + candidates = [] + for gpu_id in self.gpu_devices: + used_slots = self.gpu_slots_in_use.get(gpu_id, 0) + used_slots += pending_slots.get(gpu_id, 0) + if used_slots >= self.gpu_slots_per_device: + continue + free = self.memory_budget["gpu"][gpu_id] + free -= self.memory_stats["gpu_in_use"].get(gpu_id, 0) + free -= pending_gpu.get(gpu_id, 0) + if free >= memory_needed: + candidates.append((free, -used_slots, -gpu_id, gpu_id)) + if not candidates: return None + return max(candidates)[-1] + + def _waiting_for_gpu_slot( + self, + memory_needed: int, + pending_gpu: Optional[Dict[int, int]] = None, + pending_slots: Optional[Dict[int, int]] = None, + ) -> bool: + pending_gpu = pending_gpu or {} + pending_slots = pending_slots or {} + for gpu_id in self.gpu_devices: + free = self.memory_budget["gpu"][gpu_id] + free -= self.memory_stats["gpu_in_use"].get(gpu_id, 0) + free -= pending_gpu.get(gpu_id, 0) + if free >= memory_needed: + return True + return False - return free_memory.index(max(free_memory)) + def restrict_gpu_devices( + self, devices: List[int], slots_per_device: int = 1 + ) -> None: + self.gpu_devices = [ + gpu_id for gpu_id in devices if gpu_id in self.memory_budget["gpu"] + ] + self.n_gpu = len(self.gpu_devices) + self.gpu_slots_per_device = max(1, int(slots_per_device)) + self.gpu_slots_in_use = {gpu_id: 0 for gpu_id in self.gpu_devices} def _get_pending_nodes(self) -> List[str]: return list(self._candidates) def _release_execution_memory(self, node_id: str, gpu_id: int) -> None: _, gpu_mem = self.node_resources[node_id] - if gpu_id is not None and gpu_mem > 0: - self.memory_stats["gpu_in_use"][gpu_id] -= gpu_mem + reserved_gpu_id = self._gpu_slot_nodes.pop(node_id, None) + if reserved_gpu_id is not None and gpu_mem > 0: + self.memory_stats["gpu_in_use"][reserved_gpu_id] = max( + 0, self.memory_stats["gpu_in_use"][reserved_gpu_id] - gpu_mem + ) + self.gpu_slots_in_use[reserved_gpu_id] = max( + 0, self.gpu_slots_in_use[reserved_gpu_id] - 1 + ) reserved = self._cpu_reserved_nodes.pop(node_id, 0) if reserved: self.memory_stats["cpu_in_flight"] = max( diff --git a/src/main/python/systemds/scuro/drsearch/representation_dag.py b/src/main/python/systemds/scuro/drsearch/representation_dag.py index ad7f174acfa..6315e853dda 100644 --- a/src/main/python/systemds/scuro/drsearch/representation_dag.py +++ b/src/main/python/systemds/scuro/drsearch/representation_dag.py @@ -577,15 +577,6 @@ def get_consumer_count(dags: List[RepresentationDag]) -> Dict[str, int]: def pushdown_aggregation(dag_group: List[RepresentationDag]) -> List[RepresentationDag]: - consumer_count: Dict[str, int] = defaultdict(int) - - for dag in dag_group: - for node in dag.nodes: - for inp in node.inputs: - consumer_count[inp] += 1 - - processed_agg_ids: Set[str] = set() - for dag in dag_group: agg_nodes = [ n @@ -593,10 +584,6 @@ def pushdown_aggregation(dag_group: List[RepresentationDag]) -> List[Representat if n.operation and issubclass(n.operation, AggregatedRepresentation) ] for agg_node in agg_nodes: - if agg_node.node_id in processed_agg_ids: - continue - processed_agg_ids.add(agg_node.node_id) - if len(agg_node.inputs) != 1: print( f"Aggregation node {agg_node.node_id} has {len(agg_node.inputs)} inputs, skipping (SHOULD NOT HAPPEN)" @@ -604,37 +591,22 @@ def pushdown_aggregation(dag_group: List[RepresentationDag]) -> List[Representat continue input_id = agg_node.inputs[0] - - processed_agg_ids.add(input_id) - if consumer_count[input_id] != 1: - continue - - input_node = None - for d in dag_group: - input_node = d.get_node_by_id(input_id) - if input_node is not None: - break + input_node = dag.get_node_by_id(input_id) if not input_node or not input_node.operation: continue - op_instance = input_node.operation(params=input_node.parameters) - if op_instance.__class__.__bases__[0].__name__ != "BertFamily": + if not getattr( + input_node.operation, "supports_aggregation_pushdown", False + ): continue - input_node.parameters["_pushdown_aggregation"] = agg_node.parameters - - for d in dag_group: - for node in d.nodes: - node.inputs = [ - input_id if inp == agg_node.node_id else inp - for inp in node.inputs - ] - - if d.root_node_id == agg_node.node_id: - d.root_node_id = input_id - - d.nodes = [n for n in d.nodes if n.node_id != agg_node.node_id] + aggregation_parameters = copy.deepcopy(agg_node.parameters) + agg_node.operation = input_node.operation + agg_node.inputs = list(input_node.inputs) + agg_node.parameters = copy.deepcopy(input_node.parameters) + agg_node.parameters["_pushdown_aggregation"] = aggregation_parameters + dag.nodes = dag.filter_connected_nodes(dag.nodes) return dag_group @@ -758,7 +730,7 @@ def group_dags_by_dependencies( return [] unique_dags: List[RepresentationDag] = [] - seen_signatures: set[Hashable] = set() + seen_signatures: Set[Hashable] = set() for dag in dags: dag_sig = dag.compute_full_node_signature(dag.root_node_id) diff --git a/src/main/python/systemds/scuro/drsearch/task.py b/src/main/python/systemds/scuro/drsearch/task.py index e74e791794c..7977c628c12 100644 --- a/src/main/python/systemds/scuro/drsearch/task.py +++ b/src/main/python/systemds/scuro/drsearch/task.py @@ -19,6 +19,7 @@ # # ------------------------------------------------------------- import copy +import os import time from typing import List from systemds.scuro.models.model import Model @@ -26,6 +27,10 @@ from sklearn.model_selection import train_test_split from systemds.scuro.representations.representation import RepresentationStats +_GPU_CONTEXT_FLOOR_BYTES = ( + int(os.environ.get("SCURO_GPU_CONTEXT_FLOOR_MB", "512")) * 1024 * 1024 +) + class PerformanceMeasure: def __init__(self, name, metrics, higher_is_better=True): @@ -170,13 +175,32 @@ def estimate_peak_memory_bytes(self, input_stats): self.model.estimate_peak_memory_bytes(feature_dim, n_train) ) + uses_gpu = getattr(self.model, "uses_gpu", None) + if uses_gpu is None: + uses_gpu = model_peak_memory_gpu > 0 + + gpu_peak_bytes = int(model_peak_memory_gpu * 1.4) if uses_gpu else 0 + if uses_gpu and gpu_peak_bytes <= 0: + gpu_peak_bytes = _GPU_CONTEXT_FLOOR_BYTES + return { "cpu_peak_bytes": int( model_peak_memory_cpu * 1.5 + scheduler_side_cpu_bytes * 2 ), - "gpu_peak_bytes": int(model_peak_memory_gpu) * 1.4, + "gpu_peak_bytes": gpu_peak_bytes, } + def estimate_relative_cost(self, input_stats) -> float: + feature_dim = ( + int(np.prod(input_stats.output_shape)) if input_stats.output_shape else 0 + ) + fold_sizes = [len(fold) for fold in self.cv_train_indices] + n_fold_train = max(fold_sizes, default=len(self.train_indices or [])) + passes = max(1, int(getattr(self.model, "epochs", 1) or 1)) + return float( + max(1, self.kfold) * passes * max(1, n_fold_train) * max(1, feature_dim) + ) + def _create_cv_splits(self): train_labels = [self.labels[i] for i in self.train_indices] train_labels_array = np.array(train_labels) diff --git a/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py b/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py index 1b2227b773c..b185bd45dfb 100644 --- a/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py +++ b/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py @@ -69,6 +69,7 @@ def __init__( window_combination_chains: int = 1, ): self._node_stats: Dict[str, Any] = {} + self.pruned = [] self.window_combination_chains = window_combination_chains self.enable_checkpointing = enable_checkpointing self.modalities = modalities @@ -483,7 +484,11 @@ def _build_modality_dag( ) current_node_id = rep_node_id rep_dag = builder.build(current_node_id) - dags.append(rep_dag) + requires_dimensionality_reduction = getattr( + operator, "requires_dimensionality_reduction", False + ) + if not requires_dimensionality_reduction: + dags.append(rep_dag) dimensionality_reduction_dags = self.add_dimensionality_reduction_operators( builder, current_node_id @@ -515,7 +520,9 @@ def _build_modality_dag( operator.get_current_parameters(), ) - agg_operator = AggregatedRepresentation(target_dimensions=1) + agg_operator = AggregatedRepresentation( + target_dimensions=1, aggregate_leading=True + ) context_agg_node_id = builder.create_operation_node( agg_operator.__class__, [context_rep_node_id], @@ -560,9 +567,9 @@ def _build_modality_dag( ) ) - if rep_dag.nodes[-1].operation().output_modality_type in [ - ModalityType.EMBEDDING - ]: + if not requires_dimensionality_reduction and rep_dag.nodes[ + -1 + ].operation().output_modality_type in [ModalityType.EMBEDDING]: dags.extend( self.default_context_operators( modality, builder, leaf_id, rep_dag, True @@ -695,12 +702,57 @@ def temporal_context_operators(self, modality, builder, leaf_id): modality.modality_type, modality.stats ) ) + if not window_lengths: + for context_operator in context_operators: + self.pruned.append( + { + "operation": context_operator.__name__, + "reason": ("no configured window length fits the input signal"), + } + ) + return [] dags = [] for context_operator in context_operators: for window_size, num_window in zip(window_lengths, num_windows): window_node_ids = [] for agg in aggregators: - context_operator_instance = context_operator(agg()) + aggregation_instance = agg() + effective_length = self._effective_window_length( + context_operator(), + window_size, + num_window, + modality.stats.max_length, + ) + input_stats = self._window_input_stats(modality, effective_length) + for parameter, values in ( + aggregation_instance.parameters or {} + ).items(): + if not isinstance(values, list): + continue + accepted = aggregation_instance.filter_parameter_domain( + parameter, values, input_stats + ) + for value in values: + if value not in accepted: + self.pruned.append( + { + "operation": aggregation_instance.name, + "window_length": effective_length, + "parameters": {parameter: value}, + "reason": "outside the valid input domain", + } + ) + failure = aggregation_instance.check_preconditions(input_stats) + if failure is not None: + self.pruned.append( + { + "operation": aggregation_instance.name, + "window_length": effective_length, + "reason": failure, + } + ) + continue + context_operator_instance = context_operator(aggregation_instance) self._apply_granularity( context_operator_instance, window_size, num_window ) diff --git a/src/main/python/systemds/scuro/drsearch/worker_pool.py b/src/main/python/systemds/scuro/drsearch/worker_pool.py index 7e78862a104..752efdb6eed 100644 --- a/src/main/python/systemds/scuro/drsearch/worker_pool.py +++ b/src/main/python/systemds/scuro/drsearch/worker_pool.py @@ -24,7 +24,7 @@ import os import signal from dataclasses import dataclass -from typing import Any, Callable, Dict, List, Optional +from typing import Any, Callable, Dict, List, Optional, Tuple import torch @@ -86,15 +86,25 @@ class _JobResult: def _worker_main( - job_q, result_q, dispatch: Dict[str, Callable], num_threads: int + job_q, + result_q, + dispatch: Dict[str, Callable], + num_threads: int, + physical_gpu_id: Optional[int], ) -> None: + os.environ["CUDA_VISIBLE_DEVICES"] = ( + "" if physical_gpu_id is None else str(physical_gpu_id) + ) _worker_initializer(num_threads) while True: job = job_q.get() if job is None: return try: - value = dispatch[job.kind](job.payload, job.gpu_id) + local_gpu_id = ( + 0 if physical_gpu_id is not None and job.gpu_id is not None else None + ) + value = dispatch[job.kind](job.payload, local_gpu_id) result_q.put(_JobResult(job.job_id, True, os.getpid(), value=value)) except Exception as e: result_q.put( @@ -141,44 +151,143 @@ def __init__( dispatch: Dict[str, Callable], ctx=None, threads_per_worker: int = 1, + gpu_devices: Optional[List[int]] = None, + gpu_slots_per_device: int = 1, + gpu_demand_fraction: float = 1.0, ): self._ctx = ctx or create_mp_context() self._dispatch = dispatch self._threads_per_worker = max(1, int(threads_per_worker)) + self.gpu_devices = list(dict.fromkeys(gpu_devices or [])) + self.gpu_slots_per_device = max(1, int(gpu_slots_per_device)) self._result_q = self._ctx.Queue() self._job_counter = itertools.count() self._workers: Dict[int, Dict[str, Any]] = {} self._idle_pids: List[int] = [] + self._idle_gpu_pids: Dict[int, List[int]] = { + gpu_id: [] for gpu_id in self.gpu_devices + } self._running: Dict[int, tuple] = {} - for _ in range(max(1, n_workers)): - self._spawn_worker() - def _spawn_worker(self) -> None: + n_workers = max(1, int(n_workers)) + gpu_capacity = len(self.gpu_devices) * self.gpu_slots_per_device + gpu_workers = 0 + if gpu_capacity: + gpu_workers = min( + n_workers, + gpu_capacity, + max(1, int(round(n_workers * float(gpu_demand_fraction)))), + ) + self._cpu_worker_count = n_workers - gpu_workers + for worker_index in range(gpu_workers): + gpu_id = self.gpu_devices[worker_index % len(self.gpu_devices)] + self._spawn_worker(gpu_id) + for _ in range(self._cpu_worker_count): + self._spawn_worker(None) + + @property + def gpu_worker_devices(self) -> List[int]: + return list( + dict.fromkeys( + worker["gpu_id"] + for worker in self._workers.values() + if worker["gpu_id"] is not None + ) + ) + + def _spawn_worker(self, physical_gpu_id: Optional[int]) -> None: set_thread_env_before_spawn(self._threads_per_worker) job_q = self._ctx.Queue() + previous_visible = os.environ.get("CUDA_VISIBLE_DEVICES") + os.environ["CUDA_VISIBLE_DEVICES"] = ( + "" if physical_gpu_id is None else str(physical_gpu_id) + ) p = self._ctx.Process( target=_worker_main, - args=(job_q, self._result_q, self._dispatch, self._threads_per_worker), + args=( + job_q, + self._result_q, + self._dispatch, + self._threads_per_worker, + physical_gpu_id, + ), daemon=True, ) p.start() - self._workers[p.pid] = {"process": p, "job_q": job_q} - self._idle_pids.append(p.pid) + if previous_visible is None: + os.environ.pop("CUDA_VISIBLE_DEVICES", None) + else: + os.environ["CUDA_VISIBLE_DEVICES"] = previous_visible + self._workers[p.pid] = { + "process": p, + "job_q": job_q, + "gpu_id": physical_gpu_id, + } + self._mark_idle(p.pid) + + def _mark_idle(self, pid: int) -> None: + worker = self._workers.get(pid) + if worker is None: + return + gpu_id = worker["gpu_id"] + idle = self._idle_pids if gpu_id is None else self._idle_gpu_pids[gpu_id] + if pid not in idle: + idle.append(pid) @property def has_idle_worker(self) -> bool: - return len(self._idle_pids) > 0 + return bool(self._idle_pids) or any(self._idle_gpu_pids.values()) + + @property + def has_idle_cpu_worker(self) -> bool: + return bool(self._idle_pids) + + def has_idle_worker_for( + self, gpu_id: Optional[int], allow_gpu_worker_for_cpu: bool = False + ) -> bool: + if gpu_id is not None: + if self._idle_gpu_pids.get(gpu_id): + return True + return not self.gpu_devices and bool(self._idle_pids) + if self._idle_pids: + return True + return (allow_gpu_worker_for_cpu or self._cpu_worker_count == 0) and any( + self._idle_gpu_pids.values() + ) @property def num_in_flight(self) -> int: return len(self._running) - def submit(self, kind: str, payload: tuple, gpu_id: Optional[int] = None) -> int: - if not self._idle_pids: + def _take_worker( + self, gpu_id: Optional[int], allow_gpu_worker_for_cpu: bool + ) -> Tuple[int, Optional[int]]: + if gpu_id is not None: + gpu_idle = self._idle_gpu_pids.get(gpu_id, []) + if gpu_idle: + return gpu_idle.pop(), gpu_id + if not self.gpu_devices and self._idle_pids: + return self._idle_pids.pop(), None + elif self._idle_pids: + return self._idle_pids.pop(), None + elif allow_gpu_worker_for_cpu or self._cpu_worker_count == 0: + for lane_gpu_id in self.gpu_devices: + if self._idle_gpu_pids[lane_gpu_id]: + return self._idle_gpu_pids[lane_gpu_id].pop(), None + raise RuntimeError("submit() called with no compatible idle worker") + + def submit( + self, + kind: str, + payload: tuple, + gpu_id: Optional[int] = None, + allow_gpu_worker_for_cpu: bool = False, + ) -> int: + if not self.has_idle_worker_for(gpu_id, allow_gpu_worker_for_cpu): raise RuntimeError("submit() called with no idle worker available") job_id = next(self._job_counter) - job = _Job(job_id, kind, payload, gpu_id) - pid = self._idle_pids.pop() + pid, dispatched_gpu_id = self._take_worker(gpu_id, allow_gpu_worker_for_cpu) + job = _Job(job_id, kind, payload, dispatched_gpu_id) self._running[job_id] = (pid, job) self._workers[pid]["job_q"].put(job) return job_id @@ -197,7 +306,7 @@ def wait(self) -> _JobResult: if entry is not None: pid, _job = entry if pid in self._workers: - self._idle_pids.append(pid) + self._mark_idle(pid) return jr for r in ready: dead_pid = sentinel_to_pid.get(r) @@ -211,6 +320,12 @@ def _replace_dead_worker(self, pid: int) -> Optional[_JobResult]: w = self._workers.pop(pid, None) if w is None: return None + physical_gpu_id = w.get("gpu_id") + gpu_idle = self._idle_gpu_pids.get(physical_gpu_id, []) + try: + gpu_idle.remove(pid) + except ValueError: + pass try: if pid in self._idle_pids: self._idle_pids.remove(pid) @@ -235,7 +350,7 @@ def _replace_dead_worker(self, pid: int) -> Optional[_JobResult]: if failed_job_id is not None: self._running.pop(failed_job_id, None) - self._spawn_worker() + self._spawn_worker(physical_gpu_id) if failed_job_id is None: return None @@ -274,4 +389,6 @@ def shutdown(self) -> None: pass self._workers.clear() self._idle_pids.clear() + for idle in self._idle_gpu_pids.values(): + idle.clear() self._running.clear() diff --git a/src/main/python/systemds/scuro/modality/joined.py b/src/main/python/systemds/scuro/modality/joined.py index 124c7952fd4..9b4bdc4791b 100644 --- a/src/main/python/systemds/scuro/modality/joined.py +++ b/src/main/python/systemds/scuro/modality/joined.py @@ -18,6 +18,7 @@ # under the License. # # ------------------------------------------------------------- +import copy import importlib import sys @@ -256,7 +257,8 @@ def _apply_representation_chunked( ) def _apply_representation(self, modality, representation): - transformed = representation.transform(modality) + normalized = self._normalize_representation_input(modality) + transformed = representation.transform(normalized) # if self.aggregation: # aggregated_data_left = self.aggregation.execute(transformed) # transformed = Modality( @@ -266,3 +268,53 @@ def _apply_representation(self, modality, representation): # transformed.data = aggregated_data_left return transformed + + @staticmethod + def _normalize_representation_input(modality): + data = [] + changed = False + + def collect_samples(value): + try: + array = np.asarray(value) + except ValueError: + array = np.asarray(value, dtype=object) + + numeric = array.dtype != object and np.issubdtype(array.dtype, np.number) + if numeric and array.ndim == 3 and array.shape[-1] <= 4: + return [array], False + if numeric and array.ndim == 2: + return [np.repeat(array[..., np.newaxis], 3, axis=-1)], True + if array.ndim == 0: + return None + + samples = [] + for entry in value: + result = collect_samples(entry) + if result is None: + return None + entry_samples, _ = result + samples.extend(entry_samples) + return samples, True + + for instance in modality.data: + result = collect_samples(instance) + if result is None: + return modality + samples, instance_changed = result + data.append(samples) + changed = changed or instance_changed + + if not changed: + return modality + + normalized = Modality( + modality.modality_type, + modality.modality_id, + copy.deepcopy(modality.metadata), + modality.data_type, + modality.transform_time, + ) + normalized.data = data + normalized.stats = modality.stats + return normalized diff --git a/src/main/python/systemds/scuro/modality/modality.py b/src/main/python/systemds/scuro/modality/modality.py index a0d1e36377d..6d9af9394b7 100644 --- a/src/main/python/systemds/scuro/modality/modality.py +++ b/src/main/python/systemds/scuro/modality/modality.py @@ -257,6 +257,30 @@ def get_data_layout(self): return None + def subset(self, indices): + indices = list(indices) + subset_modality = self.copy_from_instance() + + if self.has_metadata(): + metadata = [selective_copy_metadata(self.metadata[i]) for i in indices] + else: + metadata = [] + + subset_modality.metadata = metadata + + if self.has_data(): + data = self.data + if hasattr(data, "subset"): + subset_modality._data = data.subset(indices) + elif isinstance(data, np.ndarray): + subset_modality.data = data[indices] + else: + subset_modality.data = [data[i] for i in indices] + + subset_modality.data_type = self.data_type + subset_modality.modality_id = self.modality_id + return subset_modality + def has_data(self): return self.data is not None and len(self.data) != 0 diff --git a/src/main/python/systemds/scuro/modality/unimodal_modality.py b/src/main/python/systemds/scuro/modality/unimodal_modality.py index e1bba7df9da..c23faa06e09 100644 --- a/src/main/python/systemds/scuro/modality/unimodal_modality.py +++ b/src/main/python/systemds/scuro/modality/unimodal_modality.py @@ -18,13 +18,14 @@ # under the License. # # ------------------------------------------------------------- +import copy from concurrent.futures import ThreadPoolExecutor, as_completed import gc import time import numpy as np -from systemds.scuro import ModalityType from systemds.scuro.dataloader.base_loader import BaseLoader from systemds.scuro.modality.modality import Modality +from systemds.scuro.modality.type import ModalityType from systemds.scuro.modality.joined import JoinedModality from systemds.scuro.modality.transformed import TransformedModality from systemds.scuro.representations.representation import ( @@ -58,13 +59,27 @@ def copy_from_instance(self): new_instance.metadata = self.metadata.copy() return new_instance - def get_metadata_at_position(self, position: int): - if self.data_loader.chunk_size: - return self.metadata[ - (self.data_loader.next_chunk - 1) * self.data_loader.chunk_size - + position - ] + def subset(self, indices): + if self.data_loader.chunk_size is None: + return super().subset(indices) + + indices = list(indices) + subset_loader = copy.copy(self.data_loader) + subset_loader.indices = [self.data_loader.indices[i] for i in indices] + subset_loader.stats = copy.copy(self.data_loader.stats) + if hasattr(subset_loader.stats, "num_instances"): + subset_loader.stats.num_instances = len(indices) + subset_loader.reset() + # Re-run the setter after replacing indices so num_chunks reflects the + # subset rather than the original dataset. + subset_loader.chunk_size = self.data_loader.chunk_size + subset_modality = type(self)(subset_loader) + subset_modality.modality_id = self.modality_id + subset_modality.transform_time = self.transform_time + return subset_modality + + def get_metadata_at_position(self, position: int): return self.metadata[position] def get_stats(self): @@ -105,13 +120,15 @@ def extract_raw_data(self): Uses the data loader to read the raw data from a specified location and stores the data in the data location. """ - self.data, self.metadata = self.data_loader.load() + data, metadata = self.data_loader.load() + self._data = data + self.metadata = metadata def iter_raw_data_chunks(self, reset: bool = True): for data, metadata, chunk_indices in self.data_loader.iter_loaded_chunks( reset=reset ): - self.data = data + self._data = data self.metadata = metadata yield chunk_indices @@ -155,74 +172,117 @@ def aggregate(self, aggregation_function): if self.data is None: raise Exception("Data is None") - def apply_representations(self, representations, aggregation=None, parallel=False): + def apply_representations( + self, + representations, + aggregation=None, + parallel=False, + representation_keys=None, + aggregations=None, + ): """ Applies a list of representations to the modality. Specifically, it applies the representations to the modality in a chunked manner. :param representations: List of representations to apply :return: List of transformed modalities """ + if representation_keys is None: + representation_keys = [ + representation.name for representation in representations + ] + if len(representation_keys) != len(representations): + raise ValueError("representation_keys must match representations") + if len(set(representation_keys)) != len(representation_keys): + raise ValueError("representation_keys must be unique") + + if aggregations is None: + aggregations = [aggregation] * len(representations) + if len(aggregations) != len(representations): + raise ValueError("aggregations must match representations") + + representation_specs = list( + zip(representation_keys, representations, aggregations) + ) transformed_modalities_per_representation = {} padding_per_representation = {} original_lengths_per_representation = {} + failed_representations = {} - for representation in representations: - transformed_modality = TransformedModality(self, representation.name) + for representation_key, representation, _ in representation_specs: + transformed_modality = TransformedModality( + self, representation.name, representation.output_modality_type + ) transformed_modality.data = [] transformed_modality.metadata = [] - transformed_modalities_per_representation[representation.name] = ( + transformed_modalities_per_representation[representation_key] = ( transformed_modality ) - padding_per_representation[representation.name] = False - original_lengths_per_representation[representation.name] = [] + padding_per_representation[representation_key] = False + original_lengths_per_representation[representation_key] = [] start = ( time.time() ) # TODO: should be repalced in unimodal_representation.transform - if self.data_loader.chunk_size: - with ThreadPoolExecutor( - max_workers=len(representations) if parallel else 1 - ) as executor: - time_s = time.time() - for _ in self.iter_raw_data_chunks(reset=True): - representations_futures = {} - for representation in representations: - future = executor.submit(representation.transform, self) - representations_futures[future] = representation.name - for future in as_completed(representations_futures.keys()): - representation_name = representations_futures.get(future) + with ThreadPoolExecutor( + max_workers=len(representations) if parallel else 1 + ) as executor: + time_s = time.time() + for _ in self.iter_raw_data_chunks(reset=True): + representations_futures = {} + for ( + representation_key, + representation, + rep_aggregation, + ) in representation_specs: + if representation_key in failed_representations: + continue + future = executor.submit( + representation.transform, self, rep_aggregation + ) + representations_futures[future] = representation_key + for future in as_completed(representations_futures.keys()): + representation_key = representations_futures.get(future) + try: transformed_chunk = future.result() - transformed_modalities_per_representation[ - representation_name - ].data.extend(transformed_chunk.data) - transformed_modalities_per_representation[ - representation_name - ].metadata.extend(transformed_chunk.metadata) - for d in transformed_chunk.data: + except Exception as e: + failed_representations[representation_key] = e + continue + transformed_modalities_per_representation[ + representation_key + ].data.extend(transformed_chunk.data) + transformed_modalities_per_representation[ + representation_key + ].metadata.extend(transformed_chunk.metadata) + for d in transformed_chunk.data: + shape = getattr(d, "shape", ()) + if shape: original_lengths_per_representation[ - representation_name - ].append(d.shape[0]) + representation_key + ].append(int(shape[0])) + if self.data_loader.is_chunked: print(f"Time for transforming data chunks: {time.time() - time_s}") - else: - if not self.has_data(): - self.extract_raw_data() - new_modality = representation.transform(self) - transformed_modalities_per_representation[representation.name] = ( - new_modality - ) - for representation in representations: + for representation_name in failed_representations: + transformed_modalities_per_representation.pop(representation_name, None) + + if representations and not transformed_modalities_per_representation: + raise next(iter(failed_representations.values())) + + for representation_key, representation, _ in representation_specs: + if representation_key in failed_representations: + continue self._apply_padding( - transformed_modalities_per_representation[representation.name], - original_lengths_per_representation[representation.name], - padding_per_representation[representation.name], + transformed_modalities_per_representation[representation_key], + original_lengths_per_representation[representation_key], + padding_per_representation[representation_key], ) transformed_modalities_per_representation[ - representation.name + representation_key ].transform_time += (time.time() - start) transformed_modalities_per_representation[ - representation.name + representation_key ].self_contained = representation.self_contained gc.collect() + self.failed_representations = failed_representations return transformed_modalities_per_representation def apply_representation(self, representation, aggregation=None): diff --git a/src/main/python/systemds/scuro/models/model.py b/src/main/python/systemds/scuro/models/model.py index 22d1bbeccfd..fcd4150cab2 100644 --- a/src/main/python/systemds/scuro/models/model.py +++ b/src/main/python/systemds/scuro/models/model.py @@ -21,6 +21,8 @@ class Model: + uses_gpu = None + def __init__(self, name: str): """ Parent class for models used to perform a given task diff --git a/src/main/python/systemds/scuro/representations/bert.py b/src/main/python/systemds/scuro/representations/bert.py index 9e60f843416..394da666ccd 100644 --- a/src/main/python/systemds/scuro/representations/bert.py +++ b/src/main/python/systemds/scuro/representations/bert.py @@ -25,7 +25,17 @@ from systemds.scuro.representations.unimodal import UnimodalRepresentation import torch from transformers import AutoTokenizer, AutoModel -from systemds.scuro.representations.utils import save_embeddings +from systemds.scuro.representations.utils import ( + LengthBucketBatchSampler, + OwnerAccumulator, + OwnedSequenceDataset, + flatten_owned_sequences, + move_batch_to_device, + pin_memory_for, + pool_transformer_output, + save_embeddings, + transformer_inference_context, +) from systemds.scuro.modality.type import ModalityType from systemds.scuro.drsearch.operator_registry import register_representation from systemds.scuro.utils.memory_utility import ( @@ -39,6 +49,9 @@ class BertFamily(UnimodalRepresentation): + supports_aggregation_pushdown = True + cache_in_worker = True + def __init__( self, representation_name, @@ -69,6 +82,10 @@ def __init__( self.data_type = torch.float32 self.aggregation = aggregation self.params = params + self.model = None + self.tokenizer = None + self.bert_output = None + self._activation_hook = None if params is not None: self.layer = params.get("layer", self.layer) self.batch_size = int(params.get("batch_size", self.batch_size)) @@ -100,18 +117,15 @@ def get_output_stats(self, input_stats) -> RepresentationStats: if not isinstance(input_stats, RepresentationStats): self.stats = RepresentationStats( input_stats.num_instances, - (self.max_seq_length, 768), - aggregate_dim=(0,), + (768,), + aggregate_dim=None, dtype=self.data_type, ) else: self.stats = RepresentationStats( input_stats.num_instances, - (input_stats.output_shape[0], self.max_seq_length, 768), - aggregate_dim=( - 0, - 1, - ), + (input_stats.output_shape[0], 768), + aggregate_dim=(0,), dtype=self.data_type, ) if self.params and "_pushdown_aggregation" in self.params: @@ -181,12 +195,15 @@ def estimate_peak_memory_bytes(self, input_stats): def transform(self, modality, aggregation=None): transformed_modality = TransformedModality(modality, self) - tokenizer = AutoTokenizer.from_pretrained( - self.model_name, clean_up_tokenization_spaces=True - ) - self.model = AutoModel.from_pretrained(self.model_name) + if self.tokenizer is None: + self.tokenizer = AutoTokenizer.from_pretrained( + self.model_name, clean_up_tokenization_spaces=True + ) + if self.model is None: + self.model = AutoModel.from_pretrained(self.model_name) self.model = self.model.to(self.device) + self.model.eval() self.bert_output = None def get_activation(name): @@ -202,26 +219,29 @@ def hook(model, input, output): aggregate_dim = (0,) if self.layer != "cls": - for name, layer in self.model.named_modules(): - if name == self.layer: - layer.register_forward_hook(get_activation(name)) - break + if self._activation_hook is None: + for name, layer in self.model.named_modules(): + if name == self.layer: + self._activation_hook = layer.register_forward_hook( + get_activation(name) + ) + break if ModalityType.TEXT.has_field(modality.metadata, "text_spans"): - dataset = TextSpanDataset(modality.data, modality.metadata) - embeddings = [] - aggregate_dim = (0, 1) - for text in dataset: - embedding = self.create_embeddings( - text, self.model, tokenizer, aggregation - ) - embeddings.append( - aggregation.execute(embedding) - if aggregation is not None - else embedding - ) + chunk_groups = list(TextSpanDataset(modality.data, modality.metadata)) + chunks, owner_ids = flatten_owned_sequences(chunk_groups) + aggregate_dim = None if aggregation is not None else (0,) + embeddings = self.create_embeddings( + chunks, + self.model, + self.tokenizer, + aggregation, + owner_ids=owner_ids, + num_owners=len(chunk_groups), + grouped=aggregation is None, + ) else: embeddings = self.create_embeddings( - modality.data, self.model, tokenizer, aggregation + modality.data, self.model, self.tokenizer ) if self.output_file is not None: save_embeddings(embeddings, self.output_file) @@ -235,69 +255,97 @@ def hook(model, input, output): def assert_output_stats(self, transformed_modality): if self.stats: assert len(transformed_modality.data) == self.stats.num_instances - if len(self.stats.output_shape) == 3: - assert ( - transformed_modality.data[0].shape[0] <= self.stats.output_shape[0] - ), f"Output shape: {transformed_modality.data[0].shape}, Expected shape: {self.stats.output_shape}" + actual_shape = np.asarray(transformed_modality.data[0]).shape + if len(self.stats.output_shape) == 2: assert ( - transformed_modality.data[0].shape[1] == self.stats.output_shape[1] - ), f"Output shape: {transformed_modality.data[0].shape}, Expected shape: {self.stats.output_shape}" + actual_shape[0] <= self.stats.output_shape[0] + ), f"Output shape: {actual_shape}, Expected shape: {self.stats.output_shape}" assert ( - transformed_modality.data[0].shape[2] == self.stats.output_shape[2] - ), f"Output shape: {transformed_modality.data[0].shape}, Expected shape: {self.stats.output_shape}" + actual_shape[1] == self.stats.output_shape[1] + ), f"Output shape: {actual_shape}, Expected shape: {self.stats.output_shape}" else: assert ( - transformed_modality.data[0].shape[0] == self.stats.output_shape[0] - ), f"Output shape: {transformed_modality.data[0].shape}, Expected shape: {self.stats.output_shape}" - assert ( - transformed_modality.data[0].shape[1] == self.stats.output_shape[1] - ), f"Output shape: {transformed_modality.data[0].shape}, Expected shape: {self.stats.output_shape}" + actual_shape == self.stats.output_shape + ), f"Output shape: {actual_shape}, Expected shape: {self.stats.output_shape}" - def create_embeddings(self, data, model, tokenizer, aggregation=None): - dataset = TextDataset(data) - dataloader = DataLoader( - dataset, batch_size=self.batch_size, shuffle=False, collate_fn=None + def create_embeddings( + self, + data, + model, + tokenizer, + aggregation=None, + owner_ids=None, + num_owners=None, + grouped=False, + ): + texts = list(TextDataset(data)) + single_owner = owner_ids is None and aggregation is not None + if owner_ids is None: + owner_ids = [0] * len(texts) if single_owner else range(len(texts)) + if num_owners is None: + num_owners = 1 if single_owner else len(texts) + dataset = OwnedSequenceDataset(texts, owner_ids) + + length_encoding = tokenizer( + texts, + padding=False, + truncation=True, + max_length=self.max_seq_length, + return_attention_mask=True, ) - cls_embeddings = [] - for batch in dataloader: + attention_mask = length_encoding.get("attention_mask") + if isinstance(attention_mask, torch.Tensor): + lengths = attention_mask.sum(dim=1).tolist() + else: + lengths = [sum(mask) for mask in attention_mask] + + def collate(samples): + batch_texts, batch_owner_ids, chunk_ids = zip(*samples) inputs = tokenizer( - batch, + list(batch_texts), return_offsets_mapping=True, return_tensors="pt", - padding="max_length", + padding=True, return_attention_mask=True, truncation=True, - max_length=self.max_seq_length, # TODO: make this dynamic with parameter to tune + max_length=self.max_seq_length, + ) + inputs = dict(inputs) + inputs.pop("offset_mapping", None) + return ( + inputs, + torch.tensor(batch_owner_ids, dtype=torch.long), + torch.tensor(chunk_ids, dtype=torch.long), ) - inputs.to(self.device) - # ModalityType.TEXT.add_field_for_instances( - # modality.metadata, - # "token_to_character_mapping", - # inputs.data["offset_mapping"].tolist(), - # ) - # - # ModalityType.TEXT.add_field_for_instances( - # modality.metadata, - # "attention_masks", - # inputs.data["attention_mask"].tolist(), - # ) - del inputs.data["offset_mapping"] - - with torch.no_grad(): + + dataloader = DataLoader( + dataset, + batch_sampler=LengthBucketBatchSampler(lengths, self.batch_size), + collate_fn=collate, + pin_memory=pin_memory_for(self.device), + ) + accumulator = OwnerAccumulator(num_owners, len(dataset), aggregation) + with transformer_inference_context(self.device): + for inputs, batch_owner_ids, chunk_ids in dataloader: + inputs = move_batch_to_device(inputs, self.device) outputs = model(**inputs) if self.layer == "cls": - cls_embedding = outputs.last_hidden_state.detach().cpu().numpy() + hidden_state = outputs.last_hidden_state else: - cls_embedding = self.bert_output.cpu().numpy() - if ( - aggregation is not None - and self.layer != "pooler" - and self.layer != "pooler.activation" - ): - cls_embedding = aggregation.execute(cls_embedding) - cls_embeddings.extend(cls_embedding) - - return cls_embeddings + hidden_state = self.bert_output + pooled = pool_transformer_output( + hidden_state, + inputs["attention_mask"], + use_cls=self.layer == "cls", + ) + accumulator.update(pooled, batch_owner_ids, chunk_ids) + + embeddings = accumulator.finalize(grouped=grouped) + if single_owner: + return embeddings[0] + if aggregation is None and not grouped: + return list(embeddings) + return embeddings @register_representation(ModalityType.TEXT) diff --git a/src/main/python/systemds/scuro/representations/bow.py b/src/main/python/systemds/scuro/representations/bow.py index 9a7766106c1..12fc2e4eb69 100644 --- a/src/main/python/systemds/scuro/representations/bow.py +++ b/src/main/python/systemds/scuro/representations/bow.py @@ -44,6 +44,7 @@ def __init__(self, ngram_range=2, min_df=2, output_file=None, params=None): self.min_df = int(min_df) self.output_file = output_file self.data_type = np.float32 + self.requires_dimensionality_reduction = True def get_output_stats(self, input_stats: TextStats) -> RepresentationStats: vocab_estimate = min( diff --git a/src/main/python/systemds/scuro/representations/clip.py b/src/main/python/systemds/scuro/representations/clip.py index c4e28404466..adffc0f39cc 100644 --- a/src/main/python/systemds/scuro/representations/clip.py +++ b/src/main/python/systemds/scuro/representations/clip.py @@ -19,14 +19,25 @@ # # ------------------------------------------------------------- import numpy as np +import torch from torchvision import transforms from systemds.scuro.dataloader.video_loader import VideoStats from systemds.scuro.modality.transformed import TransformedModality from systemds.scuro.representations.representation import RepresentationStats from systemds.scuro.representations.unimodal import UnimodalRepresentation -import torch -from systemds.scuro.representations.utils import save_embeddings +from systemds.scuro.representations.utils import ( + LengthBucketBatchSampler, + OwnerAccumulator, + OwnedSequenceDataset, + flatten_owned_sequences, + get_sequence_lengths, + move_batch_to_device, + pin_memory_for, + pool_transformer_output, + save_embeddings, + transformer_inference_context, +) from systemds.scuro.modality.type import ModalityType from systemds.scuro.drsearch.operator_registry import register_representation from transformers import CLIPProcessor, CLIPModel @@ -49,9 +60,14 @@ @register_representation([ModalityType.VIDEO, ModalityType.IMAGE]) class CLIPVisual(UnimodalRepresentation): + supports_aggregation_pushdown = True + cache_in_worker = True + def __init__(self, output_file=None, batch_size=32, layer_name="", params=None): parameters = self._get_parameters() super().__init__("CLIPVisual", ModalityType.EMBEDDING, parameters) + self.params = params + self._activation_hook = None self.model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32") self.processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32") if params is not None: @@ -99,9 +115,17 @@ def _get_parameters(self): return parameters def estimate_output_memory_bytes(self, input_stats) -> int: - return input_stats.num_instances * 512 * self.data_type.itemsize + shape = self.get_output_stats(input_stats).output_shape + return int(input_stats.num_instances * np.prod(shape) * self.data_type.itemsize) def get_output_stats(self, input_stats) -> RepresentationStats: + if self.params and "_pushdown_aggregation" in self.params: + return RepresentationStats( + input_stats.num_instances, + (512,), + aggregate_dim=None, + dtype=self.data_type, + ) if isinstance(input_stats, VideoStats): return RepresentationStats( input_stats.num_instances, @@ -219,32 +243,38 @@ def transform(self, modality, aggregation=None): self.model = self.model.to(self.data_type) self.model = self.model.to(self.device) + self.model.eval() self.clip_output = None def get_activation(name): def hook(model, input, output): - self.clip_output = ( - output[0].detach() if isinstance(output, tuple) else output.detach() - ) + self.clip_output = output[0] if isinstance(output, tuple) else output return hook - if self.layer_name != "": + if self.layer_name != "" and self._activation_hook is None: for name, layer in self.model.vision_model.named_modules(): if name == self.layer_name: - layer.register_forward_hook(get_activation(name)) + self._activation_hook = layer.register_forward_hook( + get_activation(name) + ) break - embeddings = self.create_visual_embeddings(modality) + embeddings = self.create_visual_embeddings(modality, aggregation) if self.output_file is not None: save_embeddings(embeddings, self.output_file) + transformed_modality.data_type = np.float32 + transformed_modality.aggregate_dim = ( + None + if aggregation is not None or modality.modality_type == ModalityType.IMAGE + else (0,) + ) transformed_modality.data = embeddings return transformed_modality - def create_visual_embeddings(self, modality): - + def create_visual_embeddings(self, modality, aggregation=None): clip_transform = transforms.Compose( [ transforms.ToPILImage(), @@ -254,74 +284,48 @@ def create_visual_embeddings(self, modality): transforms.ConvertImageDtype(dtype=self.data_type), ] ) - dataset = CustomDataset(modality.data, self.data_type, "cpu", tf=clip_transform) - - embeddings = {} - if modality.modality_type == ModalityType.IMAGE: - embeddings = [] - for batch in torch.utils.data.DataLoader( - dataset, batch_size=self.batch_size - ): - images = batch["data"] - inputs = self.processor( - images=images, return_tensors="pt", do_rescale=False - ) - inputs.to(self.device) - - with torch.no_grad(): - if self.layer_name != "": - _ = self.model.vision_model(**inputs) - output = self.clip_output - else: - output = self.model.get_image_features(**inputs) - - output = self._pool_visual_output(output) - - embeddings.extend( - torch.flatten(output, 1) - .detach() - .cpu() - .float() - .numpy() - .astype(np.float32) - ) - return embeddings - - for instance in torch.utils.data.DataLoader(dataset): - id = int(instance["id"][0]) - frames = instance["data"][0] - embeddings[id] = [] - batch_size = self.batch_size + is_image = modality.modality_type == ModalityType.IMAGE + if is_image: + samples = modality.data + owner_ids = list(range(len(samples))) + else: + lengths = get_sequence_lengths(modality.data, modality.metadata) + samples, owner_ids = flatten_owned_sequences(modality.data, lengths) - for start_index in range(0, len(frames), batch_size): - end_index = min(start_index + batch_size, len(frames)) - frame_ids_range = range(start_index, end_index) - frame_batch = frames[frame_ids_range] + dataset = CustomDataset(samples, self.data_type, "cpu", tf=clip_transform) + dataloader = DataLoader( + dataset, + batch_size=self.batch_size, + shuffle=False, + pin_memory=pin_memory_for(self.device), + ) + owner_by_chunk = torch.tensor(owner_ids, dtype=torch.long) + accumulator = OwnerAccumulator(len(modality.data), len(dataset), aggregation) + with transformer_inference_context(self.device): + for batch in dataloader: + chunk_ids = batch["id"].long() inputs = self.processor( - images=frame_batch, return_tensors="pt", do_rescale=False + images=batch["data"], return_tensors="pt", do_rescale=False ) - inputs.to(self.device) - with torch.no_grad(): - if self.layer_name != "": - _ = self.model.vision_model(**inputs) - output = self.clip_output - else: - output = self.model.get_image_features(**inputs) + inputs = move_batch_to_device(dict(inputs), self.device) + if self.layer_name != "": + _ = self.model.vision_model(**inputs) + output = self.clip_output + else: + output = self.model.get_image_features(**inputs) output = self._pool_visual_output(output) - - embeddings[id].extend( - torch.flatten(output, 1) - .detach() - .cpu() - .float() - .numpy() - .astype(np.float32) + accumulator.update( + torch.flatten(output, 1), + owner_by_chunk.index_select(0, chunk_ids), + chunk_ids, ) - embeddings[id] = np.array(embeddings[id]) - return list(embeddings.values()) + embeddings = accumulator.finalize(grouped=not is_image and aggregation is None) + if is_image and aggregation is None: + return list(embeddings) + return embeddings def _pool_visual_output(self, output: torch.Tensor) -> torch.Tensor: if output.ndim == 4: @@ -336,6 +340,9 @@ def _pool_visual_output(self, output: torch.Tensor) -> torch.Tensor: @register_representation(ModalityType.TEXT) class CLIPText(UnimodalRepresentation): + supports_aggregation_pushdown = True + cache_in_worker = True + def __init__(self, output_file=None, batch_size=32, layer_name="", params=None): if params is not None: self.batch_size = int(params.get("batch_size", batch_size)) @@ -356,6 +363,7 @@ def __init__(self, output_file=None, batch_size=32, layer_name="", params=None): self.gpu_id = None self.device = get_device() self.params = params + self._activation_hook = None @property def gpu_id(self): @@ -400,17 +408,14 @@ def get_output_stats(self, input_stats) -> RepresentationStats: self.stats = RepresentationStats( input_stats.num_instances, (512,), - aggregate_dim=(0,), + aggregate_dim=None, dtype=self.data_type, ) else: self.stats = RepresentationStats( input_stats.num_instances, (input_stats.output_shape[0], 512), - aggregate_dim=( - 0, - 1, - ), + aggregate_dim=(0,), dtype=self.data_type, ) if self.params and "_pushdown_aggregation" in self.params: @@ -480,72 +485,120 @@ def transform(self, modality, aggregation=None): transformed_modality = TransformedModality( modality, self, self.output_modality_type ) - self.processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32") - self.model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32") + if self.processor is None: + self.processor = CLIPProcessor.from_pretrained( + "openai/clip-vit-base-patch32" + ) + if self.model is None: + self.model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32") self.model = self.model.to(self.device) + self.model.eval() self.clip_output = None def get_activation(name): def hook(model, input, output): - self.clip_output = ( - output[0].detach() if isinstance(output, tuple) else output.detach() - ) + self.clip_output = output[0] if isinstance(output, tuple) else output return hook - if self.layer_name != "": + if self.layer_name != "" and self._activation_hook is None: for name, layer in self.model.text_model.named_modules(): if name == self.layer_name: - layer.register_forward_hook(get_activation(name)) + self._activation_hook = layer.register_forward_hook( + get_activation(name) + ) break + aggregate_dim = None if ModalityType.TEXT.has_field(modality.metadata, "text_spans"): - dataset = TextSpanDataset(modality.data, modality.metadata) - embeddings = [] - for text_chunks in dataset: - embedding = self.create_text_embeddings( - text_chunks, self.model, aggregation - ) - embeddings.append(embedding) - else: + chunk_groups = list(TextSpanDataset(modality.data, modality.metadata)) + chunks, owner_ids = flatten_owned_sequences(chunk_groups) + aggregate_dim = None if aggregation is not None else (0,) embeddings = self.create_text_embeddings( - modality.data, self.model, aggregation + chunks, + self.model, + aggregation, + owner_ids=owner_ids, + num_owners=len(chunk_groups), + grouped=aggregation is None, ) + else: + embeddings = self.create_text_embeddings(modality.data, self.model) if self.output_file is not None: save_embeddings(embeddings, self.output_file) + transformed_modality.data_type = np.float32 + transformed_modality.aggregate_dim = aggregate_dim transformed_modality.data = embeddings return transformed_modality - def create_text_embeddings(self, data, model, aggregation=None): - dataset = TextDataset(data) - dataloader = DataLoader( - dataset, batch_size=self.batch_size, shuffle=False, collate_fn=None + def create_text_embeddings( + self, + data, + model, + aggregation=None, + owner_ids=None, + num_owners=None, + grouped=False, + ): + texts = list(TextDataset(data)) + single_owner = owner_ids is None and aggregation is not None + if owner_ids is None: + owner_ids = [0] * len(texts) if single_owner else range(len(texts)) + if num_owners is None: + num_owners = 1 if single_owner else len(texts) + dataset = OwnedSequenceDataset(texts, owner_ids) + + length_encoding = self.processor( + text=texts, + padding=False, + truncation=True, + max_length=self.max_seq_length, ) - embeddings = [] - for batch in dataloader: + attention_mask = length_encoding.get("attention_mask") + if isinstance(attention_mask, torch.Tensor): + lengths = attention_mask.sum(dim=1).tolist() + else: + lengths = [sum(mask) for mask in attention_mask] + + def collate(samples): + batch_texts, batch_owner_ids, chunk_ids = zip(*samples) inputs = self.processor( - text=batch, + text=list(batch_texts), return_tensors="pt", padding=True, truncation=True, - max_length=77, + max_length=self.max_seq_length, + ) + return ( + dict(inputs), + torch.tensor(batch_owner_ids, dtype=torch.long), + torch.tensor(chunk_ids, dtype=torch.long), ) - inputs.to(self.device) - with torch.no_grad(): + + dataloader = DataLoader( + dataset, + batch_sampler=LengthBucketBatchSampler(lengths, self.batch_size), + collate_fn=collate, + pin_memory=pin_memory_for(self.device), + ) + accumulator = OwnerAccumulator(num_owners, len(dataset), aggregation) + with transformer_inference_context(self.device): + for inputs, batch_owner_ids, chunk_ids in dataloader: + inputs = move_batch_to_device(inputs, self.device) if self.layer_name != "": _ = model.text_model(**inputs) - - batch_np = self.clip_output.cpu().float().numpy() - if batch_np.ndim == 3: - batch_np = batch_np.mean(axis=1) + pooled = pool_transformer_output( + self.clip_output, inputs["attention_mask"] + ) else: - batch_np = model.get_text_features(**inputs).cpu().float().numpy() - - if aggregation is not None: - batch_np = aggregation.execute(batch_np) - - embeddings.extend(batch_np) - + pooled = model.get_text_features(**inputs) + accumulator.update(pooled, batch_owner_ids, chunk_ids) + + embeddings = accumulator.finalize(grouped=grouped) + if single_owner: + return embeddings[0] + if aggregation is None and not grouped: + return list(embeddings) return embeddings diff --git a/src/main/python/systemds/scuro/representations/color_histogram.py b/src/main/python/systemds/scuro/representations/color_histogram.py index 993b179fb8b..c159b285893 100644 --- a/src/main/python/systemds/scuro/representations/color_histogram.py +++ b/src/main/python/systemds/scuro/representations/color_histogram.py @@ -35,7 +35,7 @@ ) -@register_representation(ModalityType.IMAGE) +@register_representation([ModalityType.IMAGE, ModalityType.VIDEO]) class ColorHistogram(UnimodalRepresentation): def __init__( self, @@ -70,6 +70,11 @@ def _get_parameters(self): } def compute_histogram(self, image): + if np.issubdtype(image.dtype, np.floating): + if image.size and image.min() >= 0 and image.max() <= 1: + image = image * 255 + image = np.clip(image, 0, 255).astype(np.uint8) + if self.color_space == "HSV": img = cv2.cvtColor(image, cv2.COLOR_RGB2HSV) channels = [0, 1, 2] diff --git a/src/main/python/systemds/scuro/representations/image_bind.py b/src/main/python/systemds/scuro/representations/image_bind.py index a96cffeb6f8..0cd30af1a70 100644 --- a/src/main/python/systemds/scuro/representations/image_bind.py +++ b/src/main/python/systemds/scuro/representations/image_bind.py @@ -18,83 +18,323 @@ # under the License. # # ------------------------------------------------------------- -import torch -import imagebind.data as data +import math -from imagebind.models.imagebind_model import ModalityType as IBModalityType +import numpy as np +import torch +from pytorchvideo import transforms as pv_transforms +from pytorchvideo.data.clip_sampling import ConstantClipsPerVideoSampler +import torchaudio +from torchvision import transforms -from imagebind.models import imagebind_model +from systemds.scuro.drsearch.operator_registry import register_representation from systemds.scuro.modality.transformed import TransformedModality -from systemds.scuro.representations.unimodal import UnimodalRepresentation -from systemds.scuro.representations.utils import save_embeddings - from systemds.scuro.modality.type import ModalityType -from systemds.scuro.drsearch.operator_registry import register_representation - -if torch.backends.mps.is_available(): - DEVICE = torch.device("mps") -# elif torch.cuda.is_available(): -# DEVICE = torch.device("cuda") -else: - DEVICE = torch.device("cpu") +from systemds.scuro.representations.representation import RepresentationStats +from systemds.scuro.representations.unimodal import UnimodalRepresentation +from systemds.scuro.representations.utils import ( + OwnerAccumulator, + flatten_owned_sequences, + inference_context, + save_embeddings, +) +from systemds.scuro.utils.memory_utility import get_device +from systemds.scuro.utils.torch_dataset import TextDataset, TextSpanDataset -# @register_representation([ModalityType.TEXT, ModalityType.AUDIO, ModalityType.VIDEO]) +@register_representation([ModalityType.VIDEO, ModalityType.AUDIO, ModalityType.TEXT]) class ImageBind(UnimodalRepresentation): - def __init__(self): - parameters = {} + _EMBEDDING_DIM = 1024 + _MODEL_PARAMETER_COUNT = 1_200_000_000 + _CLIPS_PER_VIDEO = 5 + _SPATIAL_CROPS = 3 + _FRAMES_PER_CLIP = 2 + _CROP_SIZE = 224 + supports_aggregation_pushdown = True + cache_in_worker = True + + def __init__(self, output_file=None, batch_size=8, params=None): + parameters = {"batch_size": [1, 2, 4, 8, 16, 32]} super().__init__("ImageBind", ModalityType.EMBEDDING, parameters) - self.model = imagebind_model.imagebind_huge(pretrained=True) - for param in self.model.parameters(): - param.requires_grad = False + self.params = params + self.output_file = output_file + self.batch_size = batch_size + if params is not None: + self.batch_size = int(params.get("batch_size", batch_size)) + self.output_file = params.get("output_file", output_file) + self.data_type = torch.float32 + self.model = None + self.device = get_device() + self._gpu_id = self.device.index + + @property + def gpu_id(self): + return self._gpu_id + + @gpu_id.setter + def gpu_id(self, gpu_id): + self._gpu_id = gpu_id + self.device = get_device(gpu_id) + if self.model is not None: + self.model = self.model.to(self.device) + + def get_output_stats(self, input_stats) -> RepresentationStats: + return RepresentationStats( + input_stats.num_instances, + (self._EMBEDDING_DIM,), + aggregate_dim=None, + dtype=self.data_type, + ) + + def estimate_output_memory_bytes(self, input_stats) -> int: + return input_stats.num_instances * self._EMBEDDING_DIM * self.data_type.itemsize + + def estimate_peak_memory_bytes(self, input_stats) -> dict: + num_instances = max(getattr(input_stats, "num_instances", 1), 1) + + model_bytes = self._MODEL_PARAMETER_COUNT * self.data_type.itemsize + + clip_bytes = ( + self._CLIPS_PER_VIDEO + * self._SPATIAL_CROPS + * self._FRAMES_PER_CLIP + * 3 + * self._CROP_SIZE + * self._CROP_SIZE + * self.data_type.itemsize + ) + preprocessed_bytes = num_instances * clip_bytes + + batch_activation_bytes = self.batch_size * clip_bytes * 4 + + output_bytes = self.estimate_output_memory_bytes(input_stats) + + decoded_bytes = ( + getattr(input_stats, "max_length", 0) + * getattr(input_stats, "max_channels", 3) + * getattr(input_stats, "max_height", self._CROP_SIZE) + * getattr(input_stats, "max_width", self._CROP_SIZE) + * self.data_type.itemsize + ) + + safety_margin_bytes = 512 * 1024 * 1024 + + gpu_peak = ( + model_bytes + preprocessed_bytes + batch_activation_bytes + output_bytes + ) + cpu_peak = ( + model_bytes + decoded_bytes + preprocessed_bytes + output_bytes + ) + safety_margin_bytes + return {"cpu_peak_bytes": int(cpu_peak), "gpu_peak_bytes": int(gpu_peak)} + + def _ensure_model(self): + global data, imagebind_model, IBModalityType + try: + import imagebind.data as data + from imagebind.models import imagebind_model + from imagebind.models.imagebind_model import ModalityType as IBModalityType + except ImportError as error: + raise ImportError( + "ImageBind requires the optional 'imagebind' package" + ) from error + if self.model is None: + self.model = imagebind_model.imagebind_huge(pretrained=True) + for param in self.model.parameters(): + param.requires_grad = False + self.model = self.model.to(self.device) self.model.eval() - self.model.to(DEVICE) - def transform(self, modality, aggregation=None): - transformed_modality = TransformedModality( - modality, self, ModalityType.EMBEDDING + @staticmethod + def _metadata_for_sample(modality, index): + # Scuro loaders retain metadata across chunks, so resolve the current + # chunk through the modality instead of indexing the raw list directly. + get_metadata = getattr(modality, "get_metadata_at_position", None) + if callable(get_metadata): + return get_metadata(index) + return modality.metadata[index] + + @staticmethod + def _sampling_rate(metadata, modality_type): + sampling_rate = metadata.get("frequency") + if sampling_rate is None or sampling_rate <= 0: + raise ValueError( + f"ImageBind requires a positive sampling frequency for {modality_type}" + ) + return sampling_rate + + def _transform_audio_data(self, samples, metadata): + """Adapt ImageBind audio preprocessing to Scuro-loaded waveforms. + + The clip sampling, mel conversion, and normalization match ImageBind's + load_and_transform_audio_data; only path-based loading is replaced. + """ + sample_rate = 16000 + clip_duration = 2 + clip_sampler = ConstantClipsPerVideoSampler( + clip_duration=clip_duration, clips_per_video=3 ) + normalize = transforms.Normalize(mean=-4.268, std=9.138) + audio_outputs = [] - result = [] - if modality.modality_type == ModalityType.TEXT: - for i, instance in enumerate(modality.data): - text_inputs = data.load_and_transform_text(instance, DEVICE) - text_embeddings = self.model({IBModalityType.TEXT: text_inputs})[ - IBModalityType.TEXT - ] - result.append(text_embeddings.mean(axis=0).cpu().detach().numpy()) - if modality.modality_type == ModalityType.AUDIO: - audio_inputs = data.load_and_transform_audio_data( - list(modality.metadata)[ - (modality.data_loader.next_chunk - 1) - * (modality.data_loader.chunk_size) : ( - modality.data_loader.next_chunk - 1 - ) - * (modality.data_loader.chunk_size) - + (modality.data_loader.chunk_size) - ], - DEVICE, + for sample, sample_metadata in zip(samples, metadata): + # ImageBind uses torchaudio.load(path). Scuro already supplies the + # decoded waveform, and cloning prevents in-place mean centering in + # waveform2melspec from modifying the modality data. + waveform = torch.as_tensor(np.asarray(sample)).clone().float() + if waveform.ndim == 1: + waveform = waveform.unsqueeze(0) + elif waveform.ndim != 2: + raise ValueError( + "ImageBind audio samples must have shape (samples,) or " + "(channels, samples)" + ) + + original_rate = self._sampling_rate(sample_metadata, ModalityType.AUDIO) + if sample_metadata.get("length") == waveform.shape[0]: + waveform = waveform.transpose(0, 1) + if original_rate != sample_rate: + waveform = torchaudio.functional.resample( + waveform, orig_freq=original_rate, new_freq=sample_rate + ) + + timepoints = data.get_clip_timepoints( + clip_sampler, waveform.size(1) / sample_rate ) - audio_embeddings = self.model({IBModalityType.AUDIO: audio_inputs})[ - IBModalityType.AUDIO + clips = [] + for start, end in timepoints: + waveform_clip = waveform[ + :, int(start * sample_rate) : int(end * sample_rate) + ] + mel_spectrogram = data.waveform2melspec( + waveform_clip, + sample_rate, + num_mel_bins=128, + target_length=204, + ) + clips.append(normalize(mel_spectrogram)) + audio_outputs.append(torch.stack(clips)) + + return torch.stack(audio_outputs).to(self.device) + + def _transform_video_data(self, samples, metadata): + """Adapt ImageBind video preprocessing to Scuro-loaded frame arrays. + + ImageBind's temporal sampling, spatial transforms, normalization, and + crop expansion are retained; file decoding is replaced by array slicing. + """ + clip_duration = 2 + clip_sampler = ConstantClipsPerVideoSampler( + clip_duration=clip_duration, clips_per_video=5 + ) + frame_sampler = pv_transforms.UniformTemporalSubsample( + num_samples=clip_duration + ) + video_transform = transforms.Compose( + [ + pv_transforms.ShortSideScale(224), + data.NormalizeVideo( + mean=(0.48145466, 0.4578275, 0.40821073), + std=(0.26862954, 0.26130258, 0.27577711), + ), ] - result.extend(audio_embeddings.cpu().detach().numpy()) - if modality.modality_type == ModalityType.VIDEO: - video_inputs = data.load_and_transform_video_data( - list(modality.metadata)[ - (modality.data_loader.next_chunk - 1) - * (modality.data_loader.chunk_size) : ( - modality.data_loader.next_chunk - 1 - ) - * (modality.data_loader.chunk_size) - + (modality.data_loader.chunk_size) - ], - DEVICE, + ) + video_outputs = [] + + for sample, sample_metadata in zip(samples, metadata): + # ImageBind's decoder returns (C, T, H, W). Scuro stores decoded + # video as (T, H, W, C), already scaled when it has a float dtype. + frames = np.asarray(sample) + if frames.ndim != 4 or frames.shape[-1] != 3: + raise ValueError( + "ImageBind video samples must have shape " + "(frames, height, width, 3)" + ) + + video = torch.as_tensor(frames).permute(3, 0, 1, 2).float() + if np.issubdtype(frames.dtype, np.integer): + video = video / 255.0 + + sampling_rate = self._sampling_rate(sample_metadata, ModalityType.VIDEO) + timepoints = data.get_clip_timepoints( + clip_sampler, video.shape[1] / sampling_rate ) - video_embeddings = self.model({IBModalityType.VISION: video_inputs})[ - IBModalityType.VISION + clips = [] + for start, end in timepoints: + # Match the ceil-based [start, end) indexing used by the + # Decord-backed EncodedVideo loader in ImageBind. + start_frame = math.ceil(sampling_rate * start) + end_frame = min(math.ceil(sampling_rate * end), video.shape[1]) + video_clip = frame_sampler(video[:, start_frame:end_frame]) + clips.append(video_transform(video_clip)) + + clips = data.SpatialCrop(224, num_crops=3)(clips) + video_outputs.append(torch.stack(clips)) + + return torch.stack(video_outputs).to(self.device) + + def _prepare_inputs(self, samples, metadata, modality_type): + if modality_type == ModalityType.TEXT: + # ImageBind's text loader already accepts decoded strings, so its + # tokenizer can be reused without a Scuro-specific adapter. + return IBModalityType.TEXT, data.load_and_transform_text( + samples, self.device + ) + if modality_type == ModalityType.AUDIO: + return IBModalityType.AUDIO, self._transform_audio_data(samples, metadata) + if modality_type == ModalityType.VIDEO: + return IBModalityType.VISION, self._transform_video_data(samples, metadata) + raise ValueError(f"ImageBind does not support {modality_type}") + + def transform(self, modality, aggregation=None): + self._ensure_model() + grouped = False + if modality.modality_type == ModalityType.TEXT and ModalityType.TEXT.has_field( + modality.metadata, "text_spans" + ): + groups = list(TextSpanDataset(modality.data, modality.metadata)) + samples, owner_ids = flatten_owned_sequences(groups) + num_owners = len(groups) + grouped = aggregation is None + else: + if modality.modality_type == ModalityType.TEXT: + samples = list(TextDataset(modality.data)) + else: + samples = modality.data + owner_ids = list(range(len(samples))) + num_owners = len(samples) + + if modality.modality_type == ModalityType.TEXT: + metadata = [None] * len(samples) + else: + metadata = [ + self._metadata_for_sample(modality, index) + for index in range(len(samples)) ] - result.extend(video_embeddings.cpu().detach().numpy()) - transformed_modality.data = result + accumulator = OwnerAccumulator(num_owners, len(samples), aggregation) + with inference_context(self.device): + for start in range(0, len(samples), self.batch_size): + end = min(start + self.batch_size, len(samples)) + modality_key, inputs = self._prepare_inputs( + samples[start:end], + metadata[start:end], + modality.modality_type, + ) + output = self.model({modality_key: inputs})[modality_key] + chunk_ids = torch.arange(start, end, device=output.device) + batch_owner_ids = torch.as_tensor( + owner_ids[start:end], device=output.device, dtype=torch.long + ) + accumulator.update(torch.flatten(output, 1), batch_owner_ids, chunk_ids) + + embeddings = accumulator.finalize(grouped=grouped) + if self.output_file is not None: + save_embeddings(embeddings, self.output_file) + + transformed_modality = TransformedModality( + modality, self, self.output_modality_type + ) + transformed_modality.data_type = np.float32 + transformed_modality.aggregate_dim = (0,) if grouped else None + transformed_modality.data = embeddings return transformed_modality diff --git a/src/main/python/systemds/scuro/representations/mel_spectrogram.py b/src/main/python/systemds/scuro/representations/mel_spectrogram.py index 3cf25b44c4a..5683b8f2117 100644 --- a/src/main/python/systemds/scuro/representations/mel_spectrogram.py +++ b/src/main/python/systemds/scuro/representations/mel_spectrogram.py @@ -93,7 +93,10 @@ def compute_feature(self, instance, sr=None): if instance.ndim == 1: return S.T - return S.transpose(0, 2, 1) + return np.swapaxes(S, -2, -1) + + def compute_features_batched(self, data, sr=None): + return self.compute_feature(np.asarray(data), sr=sr) def get_output_stats(self, input_stats) -> RepresentationStats: num_instances = getattr(input_stats, "num_instances", 0) diff --git a/src/main/python/systemds/scuro/representations/mfcc.py b/src/main/python/systemds/scuro/representations/mfcc.py index 804dc04f1f2..2994a37efeb 100644 --- a/src/main/python/systemds/scuro/representations/mfcc.py +++ b/src/main/python/systemds/scuro/representations/mfcc.py @@ -102,15 +102,18 @@ def compute_feature(self, instance, sr=None): mean = mfcc.mean(keepdims=True) std = mfcc.std(keepdims=True) else: - mean = mfcc.mean(axis=(1, 2), keepdims=True) - std = mfcc.std(axis=(1, 2), keepdims=True) + mean = mfcc.mean(axis=(-2, -1), keepdims=True) + std = mfcc.std(axis=(-2, -1), keepdims=True) mfcc -= mean mfcc /= np.maximum(std, 1e-8) if instance.ndim == 1: return mfcc.T - return mfcc.transpose(0, 2, 1) + return np.swapaxes(mfcc, -2, -1) + + def compute_features_batched(self, data, sr=None): + return self.compute_feature(np.asarray(data), sr=sr) def get_output_stats(self, input_stats) -> RepresentationStats: num_instances = getattr(input_stats, "num_instances", 0) diff --git a/src/main/python/systemds/scuro/representations/openface.py b/src/main/python/systemds/scuro/representations/openface.py new file mode 100644 index 00000000000..7715e85951f --- /dev/null +++ b/src/main/python/systemds/scuro/representations/openface.py @@ -0,0 +1,577 @@ +# ------------------------------------------------------------- +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# +# ------------------------------------------------------------- +import os +import tempfile +from contextlib import redirect_stdout +from io import StringIO +from pathlib import Path + +import cv2 +import numpy as np +import torch + +from systemds.scuro.dataloader.video_loader import VideoStats +from systemds.scuro.drsearch.operator_registry import register_representation +from systemds.scuro.modality.transformed import TransformedModality +from systemds.scuro.modality.type import ModalityType +from systemds.scuro.representations.representation import ( + CONTAINER_LIST, + RepresentationStats, +) +from systemds.scuro.representations.unimodal import UnimodalRepresentation +from systemds.scuro.representations.utils import get_sequence_lengths +from systemds.scuro.utils.static_variables import ( + NP_ARRAY_HEADER_BYTES, + PY_LIST_HEADER_BYTES, + PY_LIST_SLOT_BYTES, + get_device, +) + +_retinaface_pretrain_patched = False +_star_dirs_patched = False + + +def _patch_openface_package_defaults(needs_landmarks: bool) -> None: + global _retinaface_pretrain_patched, _star_dirs_patched + + if not _retinaface_pretrain_patched: + try: + from openface.Pytorch_Retinaface.data.config import cfg_mnet, cfg_re50 + except ImportError: + pass + else: + cfg_mnet["pretrain"] = False + cfg_re50["pretrain"] = False + _retinaface_pretrain_patched = True + + if needs_landmarks and not _star_dirs_patched: + import openface.STAR.conf.alignment as star_alignment + + star_work_dir = Path(tempfile.gettempdir()) / "scuro_openface_star" + original_init = star_alignment.Alignment.__init__ + + def patched_init(self, args): + original_init(self, args) + self.ckpt_dir = str(star_work_dir) + self.work_dir = os.path.join( + self.ckpt_dir, self.data_definition, self.folder + ) + self.model_dir = os.path.join(self.work_dir, "model") + self.log_dir = os.path.join(self.work_dir, "log") + + star_alignment.Alignment.__init__ = patched_init + _star_dirs_patched = True + + +@register_representation([ModalityType.IMAGE, ModalityType.VIDEO]) +class OpenFace(UnimodalRepresentation): + supports_aggregation_pushdown = True + cache_in_worker = True + + MODEL_REPOSITORY = "nutPace/openface_weights" + FACE_MODEL_FILENAME = "Alignment_RetinaFace.pth" + MULTITASK_MODEL_FILENAME = "MTL_backbone.pth" + LANDMARK_MODEL_FILENAME = "Landmark_98.pkl" + + FEATURE_SETS = ("landmarks", "behavioral", "multitask", "backbone", "all") + DEFAULT_FEATURE_SET = "landmarks" + NUM_LANDMARKS = 98 + BACKBONE_DIM = 1280 + + ACTION_UNIT_INTENSITIES = ("01", "06", "17", "25", "26", "02", "12", "15") + BEHAVIORAL_COLUMNS = ( + "gaze_yaw", + "gaze_pitch", + *(f"AU{action_unit}_r" for action_unit in ACTION_UNIT_INTENSITIES), + ) + EMOTION_COLUMNS = ( + "emotion_neutral", + "emotion_happy", + "emotion_sad", + "emotion_surprise", + "emotion_fear", + "emotion_disgust", + "emotion_anger", + "emotion_contempt", + ) + MULTITASK_COLUMNS = BEHAVIORAL_COLUMNS + EMOTION_COLUMNS + LANDMARK_COLUMNS = tuple( + coordinate + for landmark_id in range(NUM_LANDMARKS) + for coordinate in (f"landmark_{landmark_id}_x", f"landmark_{landmark_id}_y") + ) + DETECTION_COLUMNS = ( + "face_x1", + "face_y1", + "face_x2", + "face_y2", + "face_confidence", + *( + coordinate + for landmark_id in range(5) + for coordinate in ( + f"retinaface_landmark_{landmark_id}_x", + f"retinaface_landmark_{landmark_id}_y", + ) + ), + ) + BACKBONE_COLUMNS = tuple( + f"backbone_{feature_id}" for feature_id in range(BACKBONE_DIM) + ) + + FEATURE_SET_DIMS = { + "behavioral": len(BEHAVIORAL_COLUMNS), + "multitask": len(MULTITASK_COLUMNS), + "landmarks": len(MULTITASK_COLUMNS) + len(LANDMARK_COLUMNS), + "backbone": len(BACKBONE_COLUMNS), + "all": ( + len(MULTITASK_COLUMNS) + + len(DETECTION_COLUMNS) + + len(LANDMARK_COLUMNS) + + len(BACKBONE_COLUMNS) + ), + } + FEATURE_COLUMNS = MULTITASK_COLUMNS + LANDMARK_COLUMNS + FEATURE_DIM = len(FEATURE_COLUMNS) + + FACE_MODEL_MEMORY_BYTES = 8 * 1024 * 1024 + MULTITASK_MODEL_MEMORY_BYTES = 128 * 1024 * 1024 + LANDMARK_MODEL_MEMORY_BYTES = 192 * 1024 * 1024 + CPU_RUNTIME_OVERHEAD_BYTES = 128 * 1024 * 1024 + GPU_RUNTIME_OVERHEAD_BYTES = 64 * 1024 * 1024 + + def __init__( + self, + feature_set=DEFAULT_FEATURE_SET, + confidence_threshold=0.02, + nms_threshold=0.4, + vis_threshold=0.5, + params=None, + ): + if params is not None: + feature_set = params.get("feature_set", feature_set) + confidence_threshold = params.get( + "confidence_threshold", confidence_threshold + ) + nms_threshold = params.get("nms_threshold", nms_threshold) + vis_threshold = params.get("vis_threshold", vis_threshold) + + parameters = { + "feature_set": list(self.FEATURE_SETS), + "confidence_threshold": [0.01, 0.02, 0.05], + "nms_threshold": [0.3, 0.4, 0.5], + "vis_threshold": [0.4, 0.5, 0.6], + } + super().__init__("OpenFace", ModalityType.EMBEDDING, parameters) + self.feature_set = feature_set + self.confidence_threshold = float(confidence_threshold) + self.nms_threshold = float(nms_threshold) + self.vis_threshold = float(vis_threshold) + self.params = params + self.data_type = np.float32 + self._gpu_id = None + self.device = get_device() + self._face_detector = None + self._multitask_predictor = None + self._landmark_detector = None + self._backbone_hook = None + self._backbone_output = None + + @property + def feature_set(self): + return self._feature_set + + @feature_set.setter + def feature_set(self, feature_set): + if feature_set not in self.FEATURE_SETS: + raise ValueError( + f"Unknown OpenFace feature set '{feature_set}'. " + f"Expected one of: {', '.join(self.FEATURE_SETS)}" + ) + self._feature_set = feature_set + + @property + def feature_dim(self): + return self.FEATURE_SET_DIMS[self.feature_set] + + @property + def gpu_id(self): + return self._gpu_id + + @gpu_id.setter + def gpu_id(self, gpu_id): + self._gpu_id = gpu_id + self.device = get_device(gpu_id) + if self._backbone_hook is not None: + self._backbone_hook.remove() + self._face_detector = None + self._multitask_predictor = None + self._landmark_detector = None + self._backbone_hook = None + self._backbone_output = None + + def get_output_stats(self, input_stats) -> RepresentationStats: + if self.params and "_pushdown_aggregation" in self.params: + return RepresentationStats( + input_stats.num_instances, + (self.feature_dim,), + aggregate_dim=None, + dtype=self.data_type, + ) + + if isinstance(input_stats, VideoStats): + return RepresentationStats( + input_stats.num_instances, + (input_stats.max_length, self.feature_dim), + dtype=self.data_type, + container=CONTAINER_LIST, + ) + + if isinstance(input_stats, RepresentationStats): + return RepresentationStats( + input_stats.num_instances, + (input_stats.output_shape[0], self.feature_dim), + dtype=self.data_type, + container=CONTAINER_LIST, + ) + + return RepresentationStats( + input_stats.num_instances, + (self.feature_dim,), + aggregate_dim=None, + dtype=self.data_type, + container=CONTAINER_LIST, + ) + + def estimate_output_memory_bytes(self, input_stats) -> int: + stats = self.get_output_stats(input_stats) + payload = int( + input_stats.num_instances + * np.prod(stats.output_shape) + * np.dtype(self.data_type).itemsize + ) + return int( + PY_LIST_HEADER_BYTES + + input_stats.num_instances * (NP_ARRAY_HEADER_BYTES + PY_LIST_SLOT_BYTES) + + payload + ) + + def estimate_peak_memory_bytes(self, input_stats) -> dict: + max_height = int(getattr(input_stats, "max_height", 224)) + max_width = int(getattr(input_stats, "max_width", 224)) + max_channels = int(getattr(input_stats, "max_channels", 3)) + frame_bytes = ( + max_height * max_width * max_channels * np.dtype(np.float32).itemsize + ) + model_bytes = self.FACE_MODEL_MEMORY_BYTES + self.MULTITASK_MODEL_MEMORY_BYTES + if self.feature_set in ("landmarks", "all"): + model_bytes += self.LANDMARK_MODEL_MEMORY_BYTES + + output_bytes = self.estimate_output_memory_bytes(input_stats) + return { + "cpu_peak_bytes": int( + output_bytes + + frame_bytes + + model_bytes + + self.CPU_RUNTIME_OVERHEAD_BYTES + ), + "gpu_peak_bytes": int( + model_bytes + 2 * frame_bytes + self.GPU_RUNTIME_OVERHEAD_BYTES + ), + } + + def transform(self, modality, aggregation=None): + if modality.modality_type not in (ModalityType.IMAGE, ModalityType.VIDEO): + raise ValueError("OpenFace supports only image and video modalities") + + is_image = modality.modality_type == ModalityType.IMAGE + if is_image: + embeddings = self._extract_image_features(modality.data) + else: + embeddings = [] + lengths = get_sequence_lengths(modality.data, modality.metadata) + for owner_id, (frames, length) in enumerate(zip(modality.data, lengths)): + features = self._extract_video_features(frames[:length], owner_id) + embeddings.append( + self._aggregate(features, aggregation) + if aggregation is not None + else features + ) + + if aggregation is not None: + embeddings = np.stack(embeddings).astype(np.float32, copy=False) + + if is_image and aggregation is not None: + embeddings = np.stack( + [ + self._aggregate(feature[None, :], aggregation) + for feature in embeddings + ] + ).astype(np.float32, copy=False) + + transformed_modality = TransformedModality( + modality, self, self.output_modality_type + ) + transformed_modality.data_type = np.float32 + transformed_modality.aggregate_dim = ( + None if aggregation is not None or is_image else (0,) + ) + transformed_modality.data = embeddings + return transformed_modality + + def _extract_image_features(self, images): + return [self._extract_features(image) for image in images] + + def _extract_video_features(self, frames, owner_id): + if len(frames) == 0: + raise ValueError(f"Video instance {owner_id} contains no frames") + return np.stack([self._extract_features(frame) for frame in frames]) + + def _extract_features(self, image): + self._load_models() + image_bgr = self._to_bgr(image) + face, detections = self._face_detector.get_face(image_bgr) + if face is None or face.size == 0: + return np.zeros(self.feature_dim, dtype=np.float32) + + needs_backbone = self.feature_set in ("backbone", "all") + if needs_backbone: + self._ensure_backbone_hook() + self._backbone_output = None + + emotion_output, gaze_output, action_unit_output = ( + self._multitask_predictor.predict(face) + ) + behavioral = self._behavioral_features(gaze_output, action_unit_output) + emotion = self._checked_vector( + emotion_output, len(self.EMOTION_COLUMNS), "emotion" + ) + multitask = np.concatenate((behavioral, emotion)) + + if self.feature_set == "behavioral": + return behavioral + if self.feature_set == "multitask": + return multitask.astype(np.float32, copy=False) + + backbone = None + if needs_backbone: + backbone = self._checked_vector( + self._backbone_output, self.BACKBONE_DIM, "backbone" + ) + if self.feature_set == "backbone": + return backbone + + landmarks = self._landmark_features(image_bgr, detections) + if self.feature_set == "landmarks": + return np.concatenate((multitask, landmarks)).astype(np.float32, copy=False) + + detection = self._detection_features(detections, image_bgr.shape) + return np.concatenate((multitask, detection, landmarks, backbone)).astype( + np.float32, copy=False + ) + + def _behavioral_features(self, gaze_output, action_unit_output): + gaze = self._checked_vector(gaze_output, 2, "gaze") + action_units = self._checked_vector( + action_unit_output, len(self.ACTION_UNIT_INTENSITIES), "action-unit" + ) + return np.concatenate((gaze, action_units)).astype(np.float32, copy=False) + + def _landmark_features(self, image_bgr, detections): + with redirect_stdout(StringIO()): + landmarks = self._landmark_detector.detect_landmarks( + image_bgr, + detections[:1], + confidence_threshold=self.vis_threshold, + ) + if not landmarks: + return np.zeros(len(self.LANDMARK_COLUMNS), dtype=np.float32) + + points = np.asarray(landmarks[0], dtype=np.float32) + if points.shape != (self.NUM_LANDMARKS, 2): + raise RuntimeError( + "OpenFace 3.0 returned unexpected landmark dimensions: " + f"{points.shape}" + ) + height, width = image_bgr.shape[:2] + points = points.copy() + points[:, 0] /= width + points[:, 1] /= height + return points.reshape(-1) + + @classmethod + def _detection_features(cls, detections, image_shape): + detection = np.asarray(detections[0], dtype=np.float32) + if detection.size < len(cls.DETECTION_COLUMNS): + raise RuntimeError( + "OpenFace 3.0 returned unexpected face-detection dimensions: " + f"{detection.size}" + ) + + detection = detection[: len(cls.DETECTION_COLUMNS)].copy() + height, width = image_shape[:2] + detection[[0, 2, 5, 7, 9, 11, 13]] /= width + detection[[1, 3, 6, 8, 10, 12, 14]] /= height + return detection + + def _ensure_backbone_hook(self): + if self._backbone_hook is not None: + return + + def capture_backbone(_module, _inputs, output): + self._backbone_output = output + + self._backbone_hook = ( + self._multitask_predictor.model.base_model.register_forward_hook( + capture_backbone + ) + ) + + def _load_models(self): + needs_landmarks = self.feature_set in ("landmarks", "all") + models_ready = ( + self._face_detector is not None + and self._multitask_predictor is not None + and (not needs_landmarks or self._landmark_detector is not None) + ) + if models_ready: + return + + try: + from huggingface_hub import snapshot_download + from openface.face_detection import FaceDetector + from openface.multitask_model import MultitaskPredictor + + if needs_landmarks: + from openface.landmark_detection import LandmarkDetector + except ImportError as error: + raise ImportError( + "OpenFace 3.0 is required for this representation. " + "Install it with 'pip install openface-test'." + ) from error + + _patch_openface_package_defaults(needs_landmarks) + + weight_files = [self.FACE_MODEL_FILENAME, self.MULTITASK_MODEL_FILENAME] + if needs_landmarks: + weight_files.append(self.LANDMARK_MODEL_FILENAME) + weights_directory = Path( + snapshot_download( + repo_id=self.MODEL_REPOSITORY, + allow_patterns=weight_files, + ) + ) + + class ArrayFaceDetector(FaceDetector): + def preprocess_image(self, image, resize=1.0): + image_raw = np.ascontiguousarray(image) + detector_input = np.float32(image_raw) + if resize != 1: + detector_input = cv2.resize( + detector_input, + None, + fx=resize, + fy=resize, + interpolation=cv2.INTER_LINEAR, + ) + detector_input -= (104, 117, 123) + detector_input = detector_input.transpose(2, 0, 1) + detector_input = ( + torch.from_numpy(detector_input).unsqueeze(0).to(self.device) + ) + return detector_input, image_raw + + device = str(self.device) + if self._face_detector is None: + self._face_detector = ArrayFaceDetector( + model_path=str(weights_directory / self.FACE_MODEL_FILENAME), + device=device, + confidence_threshold=self.confidence_threshold, + nms_threshold=self.nms_threshold, + vis_threshold=self.vis_threshold, + ) + if self._multitask_predictor is None: + self._multitask_predictor = MultitaskPredictor( + model_path=str(weights_directory / self.MULTITASK_MODEL_FILENAME), + device=device, + ) + if needs_landmarks and self._landmark_detector is None: + landmark_device = self.device.type + device_ids = ( + [-1] + if landmark_device == "cpu" + else [self.device.index if self.device.index is not None else 0] + ) + self._landmark_detector = LandmarkDetector( + model_path=str(weights_directory / self.LANDMARK_MODEL_FILENAME), + device=landmark_device, + device_ids=device_ids, + ) + + @staticmethod + def _checked_vector(output, expected_size, output_name): + if output is None: + size = 0 + else: + if hasattr(output, "detach"): + output = output.detach().cpu().numpy() + output = np.asarray(output, dtype=np.float32).reshape(-1) + size = output.size + if size != expected_size: + raise RuntimeError( + f"OpenFace 3.0 returned unexpected {output_name} dimensions: {size}" + ) + return output + + @staticmethod + def _aggregate(features, aggregation): + return np.asarray(aggregation.execute(features), dtype=np.float32) + + @staticmethod + def _to_bgr(image): + if hasattr(image, "detach"): + image = image.detach().cpu().numpy() + image = np.asarray(image) + if image.ndim not in (2, 3): + raise ValueError( + f"Expected an image tensor with 2 or 3 dimensions, got {image.shape}" + ) + + if np.issubdtype(image.dtype, np.floating): + finite = image[np.isfinite(image)] + if finite.size and finite.min() >= 0.0 and finite.max() <= 1.0: + image = image * 255.0 + image = np.nan_to_num(image, nan=0.0, posinf=255.0, neginf=0.0) + image = np.clip(image, 0, 255).astype(np.uint8, copy=False) + + if image.ndim == 2: + image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR) + elif image.shape[2] == 1: + image = cv2.cvtColor(image[:, :, 0], cv2.COLOR_GRAY2BGR) + elif image.shape[2] == 3: + image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR) + elif image.shape[2] == 4: + image = cv2.cvtColor(image, cv2.COLOR_RGBA2BGR) + else: + raise ValueError( + f"Expected 1, 3, or 4 image channels, got {image.shape[2]}" + ) + return np.ascontiguousarray(image) diff --git a/src/main/python/systemds/scuro/representations/optical_flow.py b/src/main/python/systemds/scuro/representations/optical_flow.py index 943cd2ab499..e4f147dec56 100644 --- a/src/main/python/systemds/scuro/representations/optical_flow.py +++ b/src/main/python/systemds/scuro/representations/optical_flow.py @@ -52,11 +52,14 @@ def transform(self, modality, aggregation=None): ) for video_id, instance in enumerate(modality.data): + if instance.dtype == np.float16: + instance = instance.astype(np.float32) + transformed_modality.data.append([]) - previous_gray = cv2.cvtColor(instance[0], cv2.COLOR_BGR2GRAY) + previous_gray = cv2.cvtColor(instance[0], cv2.COLOR_RGB2GRAY) for frame_id in range(1, len(instance)): - gray = cv2.cvtColor(instance[frame_id], cv2.COLOR_BGR2GRAY) + gray = cv2.cvtColor(instance[frame_id], cv2.COLOR_RGB2GRAY) flow = cv2.calcOpticalFlowFarneback( previous_gray, @@ -72,5 +75,6 @@ def transform(self, modality, aggregation=None): ) transformed_modality.data[video_id].append(flow) + previous_gray = gray transformed_modality.update_metadata() return transformed_modality diff --git a/src/main/python/systemds/scuro/representations/resnet.py b/src/main/python/systemds/scuro/representations/resnet.py index 26a3350e9c8..ccefd4e11b4 100644 --- a/src/main/python/systemds/scuro/representations/resnet.py +++ b/src/main/python/systemds/scuro/representations/resnet.py @@ -21,6 +21,14 @@ from systemds.scuro.dataloader.image_loader import ImageStats from systemds.scuro.dataloader.video_loader import VideoStats from systemds.scuro.representations.representation import RepresentationStats +from systemds.scuro.representations.utils import ( + OwnerAccumulator, + flatten_owned_sequences, + get_sequence_lengths, + inference_context, + move_batch_to_device, + pin_memory_for, +) from systemds.scuro.utils.torch_dataset import CustomDataset from systemds.scuro.modality.transformed import TransformedModality from systemds.scuro.representations.unimodal import UnimodalRepresentation @@ -41,6 +49,9 @@ def forward(self, input_: torch.Tensor) -> torch.Tensor: @register_representation([ModalityType.IMAGE, ModalityType.VIDEO]) class ResNet(UnimodalRepresentation): + supports_aggregation_pushdown = True + cache_in_worker = True + def __init__( self, model_name="ResNet18", @@ -51,6 +62,8 @@ def __init__( ): self.data_type = torch.float32 self.model = None + self._activation_hook = None + self.activation = None self.gpu_id = None self.device = get_device() if params is not None: @@ -63,6 +76,7 @@ def __init__( self.model_name = model_name parameters = self._get_parameters() super().__init__("ResNet", ModalityType.EMBEDDING, parameters) + self.params = params self.output_file = output_file self.model.eval() @@ -116,16 +130,15 @@ def model_name(self, model_name): raise NotImplementedError def estimate_output_memory_bytes(self, input_stats: ImageStats) -> int: - if isinstance(input_stats, VideoStats): - return ( - input_stats.num_instances - * input_stats.max_length - * 512 - * self.data_type.itemsize - ) - return input_stats.num_instances * 512 * self.data_type.itemsize + shape = self.get_output_stats(input_stats).output_shape + return int(input_stats.num_instances * np.prod(shape) * self.data_type.itemsize) def get_output_stats(self, input_stats) -> RepresentationStats: + if self.params and "_pushdown_aggregation" in self.params: + return RepresentationStats( + input_stats.num_instances, (512,), aggregate_dim=None + ) + if isinstance(input_stats, VideoStats): return RepresentationStats( input_stats.num_instances, @@ -178,7 +191,11 @@ def estimate_peak_memory_bytes(self, input_stats: ImageStats) -> dict: return {"cpu_peak_bytes": cpu_peak, "gpu_peak_bytes": gpu_peak} def _get_parameters(self, high_level=True): - parameters = {"model_name": [], "layer_name": []} + parameters = { + "batch_size": [1, 2, 4, 8, 16, 32, 64, 128], + "model_name": [], + "layer_name": [], + } for m in ["ResNet18", "ResNet34", "ResNet50", "ResNet101", "ResNet152"]: parameters["model_name"].append(m) @@ -199,76 +216,68 @@ def _get_parameters(self, high_level=True): def transform(self, modality, aggregation=None): if next(self.model.parameters()).dtype != self.data_type: self.model = self.model.to(self.data_type) - - embeddings = {} - dataset = CustomDataset(modality.data, self.data_type, self.device) - res5c_output = None + self.model = self.model.to(self.device) + self.model.eval() + self.activation = None def get_features(name_): def hook( _module: torch.nn.Module, input_: Tuple[torch.Tensor], output: Any ): - nonlocal res5c_output - res5c_output = output + self.activation = output return hook - if self.layer_name: + if self.layer_name and self._activation_hook is None: for name, layer in self.model.named_modules(): if name == self.layer_name: - layer.register_forward_hook(get_features(name)) + self._activation_hook = layer.register_forward_hook( + get_features(name) + ) break - if modality.modality_type == ModalityType.IMAGE: - embeddings = [] - for batch in torch.utils.data.DataLoader( - dataset, batch_size=self.batch_size - ): - image_batch = batch["data"] - _ = self.model(image_batch) - output = res5c_output - embeddings.extend( - output.squeeze().detach().cpu().numpy().astype(modality.data_type) - ) - torch.cuda.empty_cache() + is_image = modality.modality_type == ModalityType.IMAGE + if is_image: + samples = modality.data + owner_ids = list(range(len(samples))) else: - for instance in torch.utils.data.DataLoader(dataset): - video_id = instance["id"][0] - frames = instance["data"][0] - embeddings[video_id] = [] - batch_size = 64 - - if modality.modality_type == ModalityType.IMAGE: - frames = frames.unsqueeze(0) - - for start_index in range(0, len(frames), batch_size): - end_index = min(start_index + batch_size, len(frames)) - frame_ids_range = range(start_index, end_index) - frame_batch = frames[frame_ids_range] - - _ = self.model(frame_batch) - output = res5c_output - if len(output.shape) > 2: - output = torch.nn.functional.adaptive_avg_pool2d(output, (1, 1)) - # TODO: check if the dimensions are correct here - embeddings[video_id].extend( - torch.flatten(output, 1) - .detach() - .cpu() - .float() - .numpy() - .astype(np.float32) - ) + lengths = get_sequence_lengths(modality.data, modality.metadata) + samples, owner_ids = flatten_owned_sequences(modality.data, lengths) + + dataset = CustomDataset(samples, self.data_type, "cpu") + dataloader = torch.utils.data.DataLoader( + dataset, + batch_size=self.batch_size, + shuffle=False, + pin_memory=pin_memory_for(self.device), + ) + owner_by_chunk = torch.tensor(owner_ids, dtype=torch.long) + accumulator = OwnerAccumulator(len(modality.data), len(dataset), aggregation) + + with inference_context(self.device): + for batch in dataloader: + chunk_ids = batch["id"].long() + batch = move_batch_to_device(batch, self.device) + _ = self.model(batch["data"]) + output = self.activation + if output.ndim > 2: + output = torch.nn.functional.adaptive_avg_pool2d(output, (1, 1)) + accumulator.update( + torch.flatten(output, 1), + owner_by_chunk.index_select(0, chunk_ids), + chunk_ids, + ) - embeddings[video_id] = np.array(embeddings[video_id]) + embeddings = accumulator.finalize(grouped=not is_image and aggregation is None) + if is_image and aggregation is None: + embeddings = list(embeddings) transformed_modality = TransformedModality( modality, self, self.output_modality_type ) - - if isinstance(embeddings, dict): - transformed_modality.data = list(embeddings.values()) - else: - transformed_modality.data = embeddings - + transformed_modality.data_type = np.float32 + transformed_modality.aggregate_dim = ( + None if aggregation is not None or is_image else (0,) + ) + transformed_modality.data = embeddings return transformed_modality diff --git a/src/main/python/systemds/scuro/representations/swin_video_transformer.py b/src/main/python/systemds/scuro/representations/swin_video_transformer.py index 7bb40e278f3..f57ac18f037 100644 --- a/src/main/python/systemds/scuro/representations/swin_video_transformer.py +++ b/src/main/python/systemds/scuro/representations/swin_video_transformer.py @@ -31,10 +31,17 @@ from systemds.scuro.modality.type import ModalityType from systemds.scuro.drsearch.operator_registry import register_representation from systemds.scuro.dataloader.video_loader import VideoStats +from systemds.scuro.representations.utils import ( + LengthBucketBatchSampler, + OwnerAccumulator, + get_sequence_lengths, + move_batch_to_device, + pin_memory_for, + transformer_inference_context, +) from systemds.scuro.utils.torch_dataset import CustomDataset from systemds.scuro.utils.static_variables import ( - compute_batch_size, get_device, get_device_for_model, ) @@ -43,8 +50,9 @@ @register_representation([ModalityType.VIDEO]) class SwinVideoTransformer(UnimodalRepresentation): _EMBED_DIM = 768 + cache_in_worker = True - def __init__(self, layer_name="avgpool", params=None): + def __init__(self, layer_name="avgpool", batch_size=8, params=None): parameters = { "layer_name": [ "features", @@ -56,19 +64,33 @@ def __init__(self, layer_name="avgpool", params=None): "features.6", "avgpool", ], + "batch_size": [1, 2, 4, 8, 16, 32], } self.data_type = torch.float32 super().__init__("SwinVideoTransformer", ModalityType.EMBEDDING, parameters) if params is not None: layer_name = params.get("layer_name", layer_name) + batch_size = int(params.get("batch_size", batch_size)) self.layer_name = layer_name + self.batch_size = batch_size self.model = swin3d_t(weights=models.video.Swin3D_T_Weights.KINETICS400_V1) self.device = get_device_for_model(self.model, memory_factor=1.5) + self._gpu_id = self.device.index + self._activation_hook = None self.model = self.model.to(self.device) self.model.eval() for param in self.model.parameters(): param.requires_grad = False + @property + def gpu_id(self): + return self._gpu_id + + @gpu_id.setter + def gpu_id(self, gpu_id): + self._gpu_id = gpu_id + self.device = get_device(gpu_id) + def get_output_stats(self, input_stats) -> RepresentationStats: num_instances = getattr(input_stats, "num_instances", 0) return RepresentationStats(num_instances, (self._EMBED_DIM,)) @@ -116,61 +138,57 @@ def estimate_peak_memory_bytes(self, input_stats: VideoStats) -> dict: return {"cpu_peak_bytes": int(cpu_peak), "gpu_peak_bytes": int(gpu_peak)} def transform(self, modality, aggregation=None): - embeddings = {} - swin_output = None + self.model = self.model.to(self.device) + self.model.eval() + self.swin_output = None def get_features(name_): def hook( _module: torch.nn.Module, input_: Tuple[torch.Tensor], output: Any ): - nonlocal swin_output - swin_output = output + self.swin_output = output return hook - sample = modality.data[0] if modality.data else "" - self.batch_size = compute_batch_size( - model=self.model, - device=self.device, - sample_data=sample, - tokenizer=None, - max_seq_length=None, - max_batch_size=128, - ) - - if self.layer_name: + if self.layer_name and self._activation_hook is None: for name, layer in self.model.named_modules(): if name == self.layer_name: - layer.register_forward_hook(get_features(name)) + self._activation_hook = layer.register_forward_hook( + get_features(name) + ) break - dataset = CustomDataset(modality.data, self.data_type, self.device) - for instance in torch.utils.data.DataLoader(dataset): - video_id = instance["id"][0] - frames = instance["data"][0] - embeddings[video_id] = [] - - frames = frames.unsqueeze(0).permute(0, 2, 1, 3, 4) - - _ = self.model(frames) - values = swin_output - pooled = torch.nn.functional.adaptive_avg_pool2d(values, (1, 1)) - - embeddings[video_id].extend( - torch.flatten(pooled, 1) - .detach() - .cpu() - .numpy() - .flatten() - .astype(modality.data_type) - ) - - embeddings[video_id] = np.array(embeddings[video_id]) + dataset = CustomDataset(modality.data, self.data_type, "cpu") + lengths = get_sequence_lengths(modality.data, modality.metadata) + dataloader = torch.utils.data.DataLoader( + dataset, + batch_sampler=LengthBucketBatchSampler( + lengths, self.batch_size, exact=True + ), + pin_memory=pin_memory_for(self.device), + ) + accumulator = OwnerAccumulator(len(dataset), len(dataset), aggregation) + + with transformer_inference_context(self.device): + for batch in dataloader: + batch = move_batch_to_device(batch, self.device) + video_ids = batch["id"].long() + frames = batch["data"].permute(0, 2, 1, 3, 4) + _ = self.model(frames) + values = self.swin_output + if isinstance(values, tuple): + values = values[0] + if values.ndim == 2: + pooled = values + elif self.layer_name.startswith("features"): + pooled = values.mean(dim=tuple(range(1, values.ndim - 1))) + else: + pooled = values.mean(dim=tuple(range(2, values.ndim))) + accumulator.update(pooled, video_ids, video_ids) transformed_modality = TransformedModality( modality, self, self.output_modality_type ) - - transformed_modality.data = list(embeddings.values()) - + transformed_modality.data = accumulator.finalize() + transformed_modality.data_type = np.float32 return transformed_modality diff --git a/src/main/python/systemds/scuro/representations/tfidf.py b/src/main/python/systemds/scuro/representations/tfidf.py index bea18a56024..a68fe50e407 100644 --- a/src/main/python/systemds/scuro/representations/tfidf.py +++ b/src/main/python/systemds/scuro/representations/tfidf.py @@ -41,6 +41,7 @@ def __init__(self, min_df=2, output_file=None, params=None): self.min_df = int(min_df) self.output_file = output_file self.data_type = np.float32 + self.requires_dimensionality_reduction = True def get_output_stats(self, input_stats: TextStats) -> RepresentationStats: vocab_estimate = min( diff --git a/src/main/python/systemds/scuro/representations/timeseries_representations.py b/src/main/python/systemds/scuro/representations/timeseries_representations.py index 2c7fbbe7404..e6aa999e8f2 100644 --- a/src/main/python/systemds/scuro/representations/timeseries_representations.py +++ b/src/main/python/systemds/scuro/representations/timeseries_representations.py @@ -249,10 +249,8 @@ def get_output_stats(self, input_stats): @register_context_representation_operator(ModalityType.PHYSIOLOGICAL) @register_context_representation_operator(ModalityType.AUDIO) class Kurtosis(TimeSeriesRepresentation): - min_input_length = 4 # the fourth moment is undefined below four samples - def __init__(self, params=None): - super().__init__("Kurtosis") + super().__init__("Kurtosis", min_input_length=4) def compute_feature(self, signal, axis=-1): return np.array(stats.kurtosis(signal, fisher=True, bias=True, axis=axis)) diff --git a/src/main/python/systemds/scuro/representations/utils.py b/src/main/python/systemds/scuro/representations/utils.py index 5041e18770c..04e6829d7ab 100644 --- a/src/main/python/systemds/scuro/representations/utils.py +++ b/src/main/python/systemds/scuro/representations/utils.py @@ -20,8 +20,297 @@ # ------------------------------------------------------------- import os import pickle +from bisect import bisect_right +from collections.abc import Sequence +from contextlib import contextmanager import numpy as np +import torch + + +def pool_transformer_output(hidden_state, attention_mask, use_cls=False): + """Pool transformer tokens without including padding tokens.""" + if hidden_state.ndim == 2: + return hidden_state + if hidden_state.ndim != 3: + raise ValueError( + f"Unexpected transformer output shape: {tuple(hidden_state.shape)}" + ) + if use_cls: + return hidden_state[:, 0, :] + + mask = attention_mask.unsqueeze(-1).to( + device=hidden_state.device, dtype=hidden_state.dtype + ) + token_count = mask.sum(dim=1).clamp_min(1) + return (hidden_state * mask).sum(dim=1) / token_count + + +def aggregate_chunk_embeddings(embeddings, aggregation): + """Aggregate all chunks of one instance while they are still on the GPU.""" + name = aggregation.aggregation_function + if name == "mean": + return embeddings.mean(dim=0) + if name == "max": + return embeddings.max(dim=0).values + if name == "min": + return embeddings.min(dim=0).values + if name == "sum": + return embeddings.sum(dim=0) + if name == "median": + return torch.quantile(embeddings, 0.5, dim=0) + if name == "mode": + return embeddings.mode(dim=0).values + raise ValueError(f"Unsupported aggregation function: {name}") + + +class LengthBucketBatchSampler(torch.utils.data.Sampler): + """Build deterministic batches from samples with similar sequence lengths.""" + + def __init__(self, lengths, batch_size, exact=False): + self.lengths = [int(length) for length in lengths] + self.batch_size = max(1, int(batch_size)) + self.exact = exact + self._batches = self._build_batches() + + def _build_batches(self): + indices = sorted(range(len(self.lengths)), key=lambda i: (self.lengths[i], i)) + if not self.exact: + return [ + indices[start : start + self.batch_size] + for start in range(0, len(indices), self.batch_size) + ] + + batches = [] + start = 0 + while start < len(indices): + length = self.lengths[indices[start]] + end = start + while end < len(indices) and self.lengths[indices[end]] == length: + end += 1 + batches.extend( + indices[offset : min(offset + self.batch_size, end)] + for offset in range(start, end, self.batch_size) + ) + start = end + return batches + + def __iter__(self): + return iter(self._batches) + + def __len__(self): + return len(self._batches) + + +class OwnedSequenceDataset(torch.utils.data.Dataset): + """Sequence samples with stable chunk and owner identifiers.""" + + def __init__(self, samples, owner_ids=None): + self.samples = list(samples) + if owner_ids is None: + owner_ids = range(len(self.samples)) + self.owner_ids = [int(owner_id) for owner_id in owner_ids] + if len(self.samples) != len(self.owner_ids): + raise ValueError("Each sequence must have exactly one owner_id") + + def __getitem__(self, chunk_id): + return self.samples[chunk_id], self.owner_ids[chunk_id], chunk_id + + def __len__(self): + return len(self.samples) + + +class FlattenedSequence(Sequence): + """Lazy flattened view over per-owner sequences.""" + + def __init__(self, sequences, lengths): + self.sequences = sequences + self.lengths = tuple(int(length) for length in lengths) + if len(self.sequences) != len(self.lengths): + raise ValueError("Each sequence must have exactly one length") + + self.offsets = [0] + for length in self.lengths: + self.offsets.append(self.offsets[-1] + length) + self._cached_owner = None + self._cached_sequence = None + + @property + def owner_ids(self): + return [ + owner_id + for owner_id, length in enumerate(self.lengths) + for _ in range(length) + ] + + def __getitem__(self, index): + if isinstance(index, slice): + return [self[i] for i in range(*index.indices(len(self)))] + if index < 0: + index += len(self) + if index < 0 or index >= len(self): + raise IndexError(index) + + owner_id = bisect_right(self.offsets, index) - 1 + if owner_id != self._cached_owner: + self._cached_sequence = self.sequences[owner_id] + self._cached_owner = owner_id + return self._cached_sequence[index - self.offsets[owner_id]] + + def __len__(self): + return self.offsets[-1] + + +def get_sequence_lengths(sequences, metadata): + if len(metadata) == len(sequences) and all("length" in md for md in metadata): + return [int(md["length"]) for md in metadata] + return [len(sequence) for sequence in sequences] + + +def flatten_owned_sequences(sequences, lengths=None): + """Flatten per-owner sequences while retaining their owner identifiers.""" + if lengths is not None: + samples = FlattenedSequence(sequences, lengths) + return samples, samples.owner_ids + + samples = [] + owner_ids = [] + for owner_id, owner_samples in enumerate(sequences): + samples.extend(owner_samples) + owner_ids.extend([owner_id] * len(owner_samples)) + return samples, owner_ids + + +def pin_memory_for(device): + device = torch.device(device) + return device.type == "cuda" and torch.cuda.is_available() + + +def move_batch_to_device(batch, device): + non_blocking = pin_memory_for(device) + return { + key: ( + value.to(device, non_blocking=non_blocking) + if isinstance(value, torch.Tensor) + else value + ) + for key, value in batch.items() + } + + +@contextmanager +def inference_context(device): + """Enable inference-only execution and mixed precision on CUDA.""" + device = torch.device(device) + with torch.inference_mode(): + if device.type != "cuda": + yield + return + dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16 + with torch.autocast(device_type="cuda", dtype=dtype): + yield + + +transformer_inference_context = inference_context + + +class OwnerAccumulator: + """Collect chunk vectors on-device and restore or aggregate their owners.""" + + def __init__(self, num_owners, num_chunks, aggregation=None): + self.num_owners = int(num_owners) + self.num_chunks = int(num_chunks) + self.aggregation = aggregation + self._values = None + self._counts = None + self._owner_ids = None + + @property + def aggregation_name(self): + if self.aggregation is None: + return None + return self.aggregation.aggregation_function + + def _initialize(self, vectors): + hidden_dim = vectors.shape[1] + name = self.aggregation_name + if name is None or name in ("median", "mode"): + self._values = torch.empty( + (self.num_chunks, hidden_dim), + device=vectors.device, + dtype=torch.float32, + ) + self._owner_ids = torch.empty( + self.num_chunks, device=vectors.device, dtype=torch.long + ) + else: + fill = 0.0 + if name == "max": + fill = -torch.inf + elif name == "min": + fill = torch.inf + self._values = torch.full( + (self.num_owners, hidden_dim), + fill, + device=vectors.device, + dtype=torch.float32, + ) + self._counts = torch.zeros( + self.num_owners, device=vectors.device, dtype=torch.long + ) + + def update(self, vectors, owner_ids, chunk_ids): + vectors = vectors.detach().float() + owner_ids = owner_ids.to(device=vectors.device, dtype=torch.long) + chunk_ids = chunk_ids.to(device=vectors.device, dtype=torch.long) + if self._values is None: + self._initialize(vectors) + + name = self.aggregation_name + if name is None or name in ("median", "mode"): + self._values.index_copy_(0, chunk_ids, vectors) + self._owner_ids.index_copy_(0, chunk_ids, owner_ids) + return + + self._counts.index_add_( + 0, owner_ids, torch.ones_like(owner_ids, dtype=torch.long) + ) + if name in ("mean", "sum"): + self._values.index_add_(0, owner_ids, vectors) + elif name in ("max", "min"): + indices = owner_ids[:, None].expand_as(vectors) + self._values.scatter_reduce_( + 0, indices, vectors, reduce=f"a{name}", include_self=True + ) + else: + raise ValueError(f"Unsupported aggregation function: {name}") + + def finalize(self, grouped=False): + if self._values is None: + return np.empty((self.num_owners, 0), dtype=np.float32) + + name = self.aggregation_name + if name is not None: + if name == "mean": + values = self._values / self._counts.clamp_min(1).unsqueeze(1) + elif name in ("sum", "max", "min"): + values = self._values + else: + values = torch.stack( + [ + aggregate_chunk_embeddings( + self._values[self._owner_ids == owner], self.aggregation + ) + for owner in range(self.num_owners) + ] + ) + return values.cpu().numpy().astype(np.float32, copy=False) + + values = self._values.cpu().numpy().astype(np.float32, copy=False) + if not grouped: + return values + owner_ids = self._owner_ids.cpu().numpy() + return [values[owner_ids == owner] for owner in range(self.num_owners)] def dense_instance_batch(data): diff --git a/src/main/python/systemds/scuro/representations/vgg.py b/src/main/python/systemds/scuro/representations/vgg.py index c2b56e8d6bd..d73753b2288 100644 --- a/src/main/python/systemds/scuro/representations/vgg.py +++ b/src/main/python/systemds/scuro/representations/vgg.py @@ -27,7 +27,6 @@ from systemds.scuro.drsearch.operator_registry import register_representation import torch.utils.data import torch -import re import torchvision.models as models import numpy as np from systemds.scuro.modality.type import ModalityType @@ -36,6 +35,14 @@ ) from systemds.scuro.dataloader.image_loader import ImageStats from systemds.scuro.representations.representation import RepresentationStats +from systemds.scuro.representations.utils import ( + OwnerAccumulator, + flatten_owned_sequences, + get_sequence_lengths, + inference_context, + move_batch_to_device, + pin_memory_for, +) class Identity(torch.nn.Module): @@ -45,18 +52,25 @@ def forward(self, input_: torch.Tensor) -> torch.Tensor: @register_representation([ModalityType.IMAGE, ModalityType.VIDEO]) class VGG19(UnimodalRepresentation): + supports_aggregation_pushdown = True + cache_in_worker = True + def __init__( self, layer="classifier.0", output_file=None, params=None, batch_size=32 ): self.data_type = torch.bfloat16 self.model = None + self._activation_hook = None + self.activation = None self.gpu_id = None self.device = get_device() self.model = models.vgg19(weights=models.VGG19_Weights.DEFAULT) self.model = self.model.to(self.device) parameters = self._get_parameters() super().__init__("VGG19", ModalityType.EMBEDDING, parameters) + self.params = params if params is not None: + batch_size = int(params.get("batch_size", batch_size)) layer = params.get("layer_name", layer) self.output_file = output_file self.layer_name = layer @@ -81,27 +95,28 @@ def gpu_id(self, gpu_id): def _get_parameters(self): parameters = { + "batch_size": [1, 2, 4, 8, 16, 32, 64, 128], "layer_name": [ "features.35", "classifier.0", "classifier.3", "classifier.6", - ] + ], } - return parameters def estimate_output_memory_bytes(self, input_stats: ImageStats) -> int: - if isinstance(input_stats, VideoStats): - return ( - input_stats.num_instances - * input_stats.max_length - * 4096 - * np.dtype(np.float32).itemsize - ) - return input_stats.num_instances * 4096 * np.dtype(np.float32).itemsize + shape = self.get_output_stats(input_stats).output_shape + return int( + input_stats.num_instances * np.prod(shape) * np.dtype(np.float32).itemsize + ) def get_output_stats(self, input_stats) -> RepresentationStats: + if self.params and "_pushdown_aggregation" in self.params: + return RepresentationStats( + input_stats.num_instances, (4096,), aggregate_dim=None + ) + if isinstance(input_stats, VideoStats): return RepresentationStats( input_stats.num_instances, @@ -121,10 +136,12 @@ def estimate_peak_memory_bytes(self, input_stats: ImageStats) -> dict: * input_stats.max_channels * self.data_type.itemsize ) - model = models.vgg19(weights=models.VGG19_Weights.DEFAULT) - param_bytes = sum(p.numel() for p in model.parameters()) - buffer_bytes = sum(b.numel() for b in model.buffers()) - model_size_bytes = param_bytes * 4 + buffer_bytes * 4 + model_size_bytes = sum( + p.nelement() * p.element_size() for p in self.model.parameters() + ) + model_size_bytes += sum( + b.nelement() * b.element_size() for b in self.model.buffers() + ) return { "cpu_peak_bytes": ( @@ -150,84 +167,68 @@ def transform(self, modality, aggregation=None): self.data_type = torch.float32 if next(self.model.parameters()).dtype != self.data_type: self.model = self.model.to(self.data_type) - - self.activations = {} + self.model = self.model.to(self.device) + self.model.eval() + self.activation = None def get_activation(name_): def hook( _module: torch.nn.Module, input_: Tuple[torch.Tensor], output: Any ): - self.activations[name_] = output + self.activation = output return hook - digit = re.findall(r"\d+", self.layer_name)[0] - if "feature" in self.layer_name: - self.model.features[int(digit)].register_forward_hook( - get_activation(self.layer_name) - ) + if self._activation_hook is None: + for name, layer in self.model.named_modules(): + if name == self.layer_name: + self._activation_hook = layer.register_forward_hook( + get_activation(name) + ) + break + + is_image = modality.modality_type == ModalityType.IMAGE + if is_image: + samples = modality.data + owner_ids = list(range(len(samples))) else: - self.model.classifier[int(digit)].register_forward_hook( - get_activation(self.layer_name) - ) - is_image = len(modality.data[0].shape) == 3 - embeddings = ( - self._transform_image_modality(modality) - if is_image - else self._transform_video_modality(modality) + lengths = get_sequence_lengths(modality.data, modality.metadata) + samples, owner_ids = flatten_owned_sequences(modality.data, lengths) + + dataset = CustomDataset(samples, self.data_type, "cpu") + dataloader = torch.utils.data.DataLoader( + dataset, + batch_size=self.batch_size, + shuffle=False, + pin_memory=pin_memory_for(self.device), ) + owner_by_chunk = torch.tensor(owner_ids, dtype=torch.long) + accumulator = OwnerAccumulator(len(modality.data), len(dataset), aggregation) + + with inference_context(self.device): + for batch in dataloader: + chunk_ids = batch["id"].long() + batch = move_batch_to_device(batch, self.device) + _ = self.model(batch["data"]) + output = self.activation + if output.ndim > 2: + output = torch.nn.functional.adaptive_avg_pool2d(output, (1, 1)) + accumulator.update( + torch.flatten(output, 1), + owner_by_chunk.index_select(0, chunk_ids), + chunk_ids, + ) + + embeddings = accumulator.finalize(grouped=not is_image and aggregation is None) + if is_image and aggregation is None: + embeddings = np.asarray(embeddings) transformed_modality = TransformedModality( modality, self, self.output_modality_type ) - + transformed_modality.data_type = np.float32 + transformed_modality.aggregate_dim = ( + None if aggregation is not None or is_image else (0,) + ) transformed_modality.data = embeddings - return transformed_modality - - def _transform_image_modality(self, modality): - dataset = CustomDataset(modality.data, self.data_type, self.device) - embeddings = [] - for instance in torch.utils.data.DataLoader(dataset, self.batch_size): - frames = instance["data"] - - _ = self.model(frames) - output = self.activations[self.layer_name] - - if len(output.shape) == 4: - output = torch.nn.functional.adaptive_avg_pool2d(output, (1, 1)) - - embeddings.extend(output.detach().cpu().float().numpy().astype(np.float32)) - - return np.array(embeddings) - - def _transform_video_modality(self, modality): - dataset = CustomDataset(modality.data, self.data_type, self.device) - embeddings = {} - for instance in torch.utils.data.DataLoader(dataset): - video_id = instance["id"][0] - frames = instance["data"][0] - embeddings[video_id] = [] - - for start_index in range(0, frames.shape[0], self.batch_size): - end_index = min(start_index + self.batch_size, frames.shape[0]) - frame_batch = frames[start_index:end_index] - - _ = self.model(frame_batch) - output = self.activations[self.layer_name] - - if len(output.shape) == 4: - output = torch.nn.functional.adaptive_avg_pool2d(output, (1, 1)) - - embeddings[video_id].extend( - torch.flatten(output, 1) - .detach() - .cpu() - .float() - .numpy() - .astype(np.float32) - ) - - embeddings[video_id] = np.array(embeddings[video_id]) - - return list(embeddings.values()) diff --git a/src/main/python/systemds/scuro/representations/wav2vec.py b/src/main/python/systemds/scuro/representations/wav2vec.py index c9f4025579a..d50465ca5c0 100644 --- a/src/main/python/systemds/scuro/representations/wav2vec.py +++ b/src/main/python/systemds/scuro/representations/wav2vec.py @@ -29,6 +29,16 @@ from systemds.scuro.representations.unimodal import UnimodalRepresentation from systemds.scuro.drsearch.operator_registry import register_representation from systemds.scuro.utils.memory_utility import get_device +from systemds.scuro.representations.utils import ( + LengthBucketBatchSampler, + OwnerAccumulator, + OwnedSequenceDataset, + move_batch_to_device, + pin_memory_for, + pool_transformer_output, + transformer_inference_context, +) +from torch.utils.data import DataLoader from transformers.utils import logging as transformers_logging @@ -38,14 +48,18 @@ @register_representation(ModalityType.AUDIO) class Wav2Vec(UnimodalRepresentation): cache_in_worker = True - instance_parallel = True + instance_parallel = False MODEL_NAME = "facebook/wav2vec2-base-960h" - def __init__(self, params=None): - super().__init__("Wav2Vec", ModalityType.TIMESERIES, {}) + def __init__(self, batch_size=8, params=None): + parameters = {"batch_size": [1, 2, 4, 8, 16, 32, 64]} + super().__init__("Wav2Vec", ModalityType.TIMESERIES, parameters) + self.batch_size = int((params or {}).get("batch_size", batch_size)) self._processor = None self._model = None + self.gpu_id = None + self.device = get_device() @staticmethod def _from_pretrained(loader_cls, name): @@ -66,29 +80,80 @@ def model(self): self._model = self._from_pretrained(Wav2Vec2Model, self.MODEL_NAME).float() return self._model + @property + def gpu_id(self): + return self._gpu_id + + @gpu_id.setter + def gpu_id(self, gpu_id): + self._gpu_id = gpu_id + self.device = get_device(gpu_id) + def transform(self, modality, aggregation=None): transformed_modality = TransformedModality( modality, self, self.output_modality_type ) - - result = [] - for i, sample in enumerate(modality.data): - sr = modality.metadata[i]["frequency"] - audio_resampled = librosa.resample( - np.array(sample), orig_sr=sr, target_sr=16000 + samples = [ + librosa.resample( + np.asarray(sample), + orig_sr=modality.metadata[owner_id]["frequency"], + target_sr=16000, ) - input = self.processor( - audio_resampled, sampling_rate=16000, return_tensors="pt", padding=True + for owner_id, sample in enumerate(modality.data) + ] + dataset = OwnedSequenceDataset(samples) + lengths = [len(sample) for sample in samples] + + def collate(batch): + audio, owner_ids, chunk_ids = zip(*batch) + inputs = self.processor( + list(audio), + sampling_rate=16000, + return_tensors="pt", + padding=True, + return_attention_mask=True, + ) + inputs = dict(inputs) + inputs["input_values"] = inputs["input_values"].float() + return ( + inputs, + torch.tensor(owner_ids, dtype=torch.long), + torch.tensor(chunk_ids, dtype=torch.long), ) - input.input_values = input.input_values.float() - input.data["input_values"] = input.data["input_values"].float() - with torch.no_grad(): - outputs = self.model(**input) - features = outputs.extract_features - # TODO: check how to get intermediate representations - result.append(torch.flatten(features.mean(dim=1)).detach().cpu().numpy()) - transformed_modality.data = np.array(result) + dataloader = DataLoader( + dataset, + batch_sampler=LengthBucketBatchSampler(lengths, self.batch_size), + collate_fn=collate, + pin_memory=pin_memory_for(self.device), + ) + model = self.model.to(self.device) + model.eval() + accumulator = OwnerAccumulator(len(dataset), len(dataset), aggregation) + + with transformer_inference_context(self.device): + for inputs, owner_ids, chunk_ids in dataloader: + inputs = move_batch_to_device(inputs, self.device) + outputs = model(**inputs) + features = outputs.extract_features + attention_mask = inputs.get("attention_mask") + if attention_mask is not None and hasattr( + model, "_get_feature_vector_attention_mask" + ): + attention_mask = model._get_feature_vector_attention_mask( + features.shape[1], attention_mask + ) + elif attention_mask is None: + attention_mask = torch.ones( + features.shape[:2], + dtype=torch.long, + device=features.device, + ) + pooled = pool_transformer_output(features, attention_mask) + accumulator.update(pooled, owner_ids, chunk_ids) + + transformed_modality.data = accumulator.finalize() + transformed_modality.data_type = np.float32 return transformed_modality def get_output_stats(self, input_stats) -> RepresentationStats: diff --git a/src/main/python/systemds/scuro/representations/window_aggregation.py b/src/main/python/systemds/scuro/representations/window_aggregation.py index 713a398312a..e8189e94312 100644 --- a/src/main/python/systemds/scuro/representations/window_aggregation.py +++ b/src/main/python/systemds/scuro/representations/window_aggregation.py @@ -20,8 +20,11 @@ # ------------------------------------------------------------- import inspect -import numpy as np import math +import os +from collections import defaultdict + +import numpy as np from systemds.scuro.modality.type import DataLayout, ModalityType @@ -50,6 +53,38 @@ def _accepts_axis(compute_feature): return accepts +_WINDOW_FEATURE_BATCH_SIZE = max( + 1, int(os.environ.get("SCURO_WINDOW_FEATURE_BATCH_SIZE", "64")) +) + + +def _compute_window_features(aggregation_function, windows): + """Compute independent windows in bounded vectorized batches when supported.""" + windows = [np.asarray(window) for window in windows] + compute_batched = getattr(aggregation_function, "compute_features_batched", None) + if not callable(compute_batched): + return [aggregation_function.compute_feature(window) for window in windows] + + results = [None] * len(windows) + grouped = defaultdict(list) + for index, window in enumerate(windows): + grouped[(window.shape, window.dtype.str)].append((index, window)) + + for group in grouped.values(): + for start in range(0, len(group), _WINDOW_FEATURE_BATCH_SIZE): + chunk = group[start : start + _WINDOW_FEATURE_BATCH_SIZE] + batch = np.stack([window for _, window in chunk]) + batch_result = np.asarray(compute_batched(batch)) + if batch_result.ndim == 0 or batch_result.shape[0] != len(chunk): + raise ValueError( + f"{aggregation_function.name}.compute_features_batched() must " + "preserve the leading batch dimension" + ) + for row, (index, _) in enumerate(chunk): + results[index] = batch_result[row] + return results + + def nested_aggregation_param_names(agg_cls): if not inspect.isclass(agg_cls): return set() @@ -497,15 +532,13 @@ def window_aggregate_single_level(self, instance, new_length): tail_result = self.aggregation_function.compute_feature(tail) full_result = _append_tail_row(full_result, tail_result) else: - full_result = np.stack( - [ - self.aggregation_function.compute_feature(full_batches[i]) - for i in range(full_batches.shape[0]) - ] - ) + windows = [full_batches[i] for i in range(full_batches.shape[0])] if tail.size: - tail_result = self.aggregation_function.compute_feature(tail) - full_result = _append_tail_row(full_result, tail_result) + windows.append(tail) + features = _compute_window_features(self.aggregation_function, windows) + full_result = np.stack(features[: full_batches.shape[0]]) + if tail.size: + full_result = _append_tail_row(full_result, features[-1]) return full_result @@ -585,8 +618,9 @@ def estimate_peak_memory_bytes(self, input_stats: RepresentationStats) -> dict: return {"cpu_peak_bytes": cpu_peak, "gpu_peak_bytes": 0} def execute(self, modality): - windowed_data = [] + instance_windows = [] for instance in modality.data: + instance = np.asarray(instance) window_size = int(np.ceil(len(instance) / self.num_windows)) padding_size = int(window_size * self.num_windows - len(instance)) pad_width = [(0, 0)] * instance.ndim @@ -594,23 +628,28 @@ def execute(self, modality): instance = np.pad( instance, pad_width=pad_width, mode="constant", constant_values=0 ) - full_batches = instance.reshape( - self.num_windows, window_size, *instance.shape[1:] + instance_windows.append( + instance.reshape(self.num_windows, window_size, *instance.shape[1:]) ) - if _accepts_axis(self.aggregation_function.compute_feature): - f = self.aggregation_function.compute_feature(full_batches, axis=1) - else: - f = np.stack( - [ - self.aggregation_function.compute_feature(full_batches[i]) - for i in range(full_batches.shape[0]) - ] - ) + if _accepts_axis(self.aggregation_function.compute_feature): + windowed_data = [ + self.aggregation_function.compute_feature(windows, axis=1) + for windows in instance_windows + ] + else: + flat_windows = [ + window for windows in instance_windows for window in windows + ] + flat_features = _compute_window_features( + self.aggregation_function, flat_windows + ) + windowed_data = [ + np.stack(flat_features[start : start + self.num_windows]) + for start in range(0, len(flat_features), self.num_windows) + ] - windowed_data.append(f) - windowed_data = _pad_stack(windowed_data) - return windowed_data + return _pad_stack(windowed_data) @register_context_operator( @@ -694,19 +733,29 @@ def estimate_peak_memory_bytes(self, input_stats: RepresentationStats) -> dict: return {"cpu_peak_bytes": cpu_peak, "gpu_peak_bytes": 0} def execute(self, modality): - windowed_data = [] + all_windows = [] + windows_per_instance = [] for instance in modality.data: + instance = np.asarray(instance) indices = np.cumsum(self._window_sizes(len(instance))) - output = [] start = 0 + count = 0 for end in indices: window = instance[start:end] window.setflags(write=False) - output.append(self.aggregation_function.compute_feature(window)) + all_windows.append(window) start = end + count += 1 + windows_per_instance.append(count) + + all_features = _compute_window_features(self.aggregation_function, all_windows) + windowed_data = [] + start = 0 + for count in windows_per_instance: + windowed_data.append(_pad_stack(all_features[start : start + count])) + start += count - windowed_data.append(_pad_stack(output)) windowed_data = _pad_stack(windowed_data) self.assert_output_stats(windowed_data) return windowed_data diff --git a/src/main/python/systemds/scuro/representations/word2vec.py b/src/main/python/systemds/scuro/representations/word2vec.py index fd1e148e117..0954a17e290 100644 --- a/src/main/python/systemds/scuro/representations/word2vec.py +++ b/src/main/python/systemds/scuro/representations/word2vec.py @@ -43,9 +43,9 @@ def get_embedding(sentence, model): @register_representation(ModalityType.TEXT) class W2V(UnimodalRepresentation): - def __init__(self, vector_size=150, min_count=1, output_file=None, params=None): + def __init__(self, vector_size=128, min_count=1, output_file=None, params=None): parameters = { - "vector_size": [50, 100, 150, 200], + "vector_size": [64, 128, 256, 512, 1024], "min_count": [1, 2, 4, 8], } super().__init__("Word2Vec", ModalityType.EMBEDDING, parameters) diff --git a/src/main/python/systemds/scuro/representations/x3d.py b/src/main/python/systemds/scuro/representations/x3d.py index bba22434fc4..fae60d1f1cd 100644 --- a/src/main/python/systemds/scuro/representations/x3d.py +++ b/src/main/python/systemds/scuro/representations/x3d.py @@ -18,25 +18,34 @@ # under the License. # # ------------------------------------------------------------- +import math +from typing import Any, Tuple + +import numpy as np +import torch +import torch.utils.data +import torchvision.models as models +from torchvision.models.video import r3d_18, s3d + +from systemds.scuro.dataloader.video_loader import VideoStats +from systemds.scuro.drsearch.operator_registry import register_representation +from systemds.scuro.modality.transformed import TransformedModality +from systemds.scuro.modality.type import ModalityType +from systemds.scuro.representations.representation import RepresentationStats +from systemds.scuro.representations.unimodal import UnimodalRepresentation +from systemds.scuro.representations.utils import ( + LengthBucketBatchSampler, + get_sequence_lengths, + inference_context, + move_batch_to_device, + pin_memory_for, + save_embeddings, +) from systemds.scuro.utils.static_variables import ( - compute_batch_size, get_device, get_device_for_model, ) from systemds.scuro.utils.torch_dataset import CustomDataset -from systemds.scuro.modality.transformed import TransformedModality -from systemds.scuro.representations.unimodal import UnimodalRepresentation -from systemds.scuro.representations.representation import RepresentationStats -from typing import Tuple, Any, Union -import torch.utils.data -import torch -from torchvision.models.video import r3d_18, s3d -import torchvision.models as models -import numpy as np -from systemds.scuro.modality.type import ModalityType -from systemds.scuro.drsearch.operator_registry import register_representation -from systemds.scuro.dataloader.video_loader import VideoStats -import math class Identity(torch.nn.Module): @@ -46,28 +55,53 @@ def forward(self, input_: torch.Tensor) -> torch.Tensor: @register_representation([ModalityType.VIDEO]) class X3D(UnimodalRepresentation): + cache_in_worker = True + def __init__( - self, layer="classifier.1", model_name="s3d", output_file=None, params=None + self, + layer="classifier.1", + model_name="s3d", + output_file=None, + batch_size=8, + params=None, ): self.data_type = torch.float32 if params is not None: model_name = params.get("model_name", model_name) layer = params.get("layer_name", layer) + batch_size = int(params.get("batch_size", batch_size)) self.model_name = model_name parameters = self._get_parameters() super().__init__("X3D", ModalityType.EMBEDDING, parameters) self.output_file = output_file self.layer_name = layer + self.batch_size = batch_size + self._gpu_id = self.device.index + self._activation_hook = None + self.activation = None self.model.eval() for param in self.model.parameters(): param.requires_grad = False self.model.fc = Identity() + @property + def gpu_id(self): + return self._gpu_id + + @gpu_id.setter + def gpu_id(self, gpu_id): + self._gpu_id = gpu_id + self.device = get_device(gpu_id) + if self.model is not None: + self.model = self.model.to(self.device) + def get_output_stats(self, input_stats) -> RepresentationStats: embedding_dim = 400 * math.floor((max(input_stats.max_length, 14) - 5) / 8) - return RepresentationStats(input_stats.num_instances, (embedding_dim,)) + return RepresentationStats( + input_stats.num_instances, (embedding_dim,), dtype=self.data_type + ) def estimate_output_memory_bytes(self, input_stats: VideoStats) -> int: embedding_dim = 400 * math.floor((max(input_stats.max_length, 14) - 5) / 8) @@ -76,7 +110,8 @@ def estimate_output_memory_bytes(self, input_stats: VideoStats) -> int: def estimate_peak_memory_bytes(self, input_stats: VideoStats) -> dict: temporal = max(input_stats.max_length, 14) input_bytes = ( - self.data_type.itemsize + self.batch_size + * self.data_type.itemsize * input_stats.max_channels * temporal * input_stats.max_height @@ -84,9 +119,11 @@ def estimate_peak_memory_bytes(self, input_stats: VideoStats) -> dict: ) output_bytes = self.estimate_output_memory_bytes(input_stats) n = max(input_stats.num_instances, 1) - output_bytes_batch = output_bytes / n + output_bytes_batch = output_bytes / n * self.batch_size - batch_peak_bytes = (input_bytes + 512 * self.data_type.itemsize) * 2 + batch_peak_bytes = ( + input_bytes + self.batch_size * 512 * self.data_type.itemsize + ) * 2 safety_margin_bytes = 100 * 1024 * 1024 @@ -129,7 +166,11 @@ def model_name(self, model_name): raise NotImplementedError def _get_parameters(self, high_level=True): - parameters = {"model_name": [], "layer_name": []} + parameters = { + "batch_size": [1, 2, 4, 8, 16, 32], + "model_name": [], + "layer_name": [], + } for m in ["r3d", "s3d"]: parameters["model_name"].append(m) @@ -160,65 +201,99 @@ def _get_parameters(self, high_level=True): parameters["layer_name"].append(name) return parameters - def transform(self, modality, aggregation=None): - sample = modality.data[0] if modality.data else "" - self.batch_size = compute_batch_size( - model=self.model, - device=self.device, - sample_data=sample, - tokenizer=None, - max_seq_length=None, - max_batch_size=128, - ) - dataset = CustomDataset(modality.data, self.data_type, self.device) + @staticmethod + def _collate_videos(samples): + video_ids = torch.tensor([sample["id"] for sample in samples]) + target_length = max(14, max(sample["data"].shape[0] for sample in samples)) + videos = [] + for sample in samples: + frames = sample["data"] + if frames.shape[0] < target_length: + pad = torch.zeros( + (target_length - frames.shape[0], *frames.shape[1:]), + dtype=frames.dtype, + ) + frames = torch.cat((frames, pad), dim=0) + videos.append(frames) + return {"id": video_ids, "data": torch.stack(videos)} - embeddings = {} - - activation = None + def transform(self, modality, aggregation=None): + self.model = self.model.to(self.device) + self.model.eval() + self.activation = None def get_features(name_): def hook( _module: torch.nn.Module, input_: Tuple[torch.Tensor], output: Any ): - nonlocal activation - activation = output + self.activation = output return hook - if self.layer_name: + if self.layer_name and self._activation_hook is None: for name, layer in self.model.named_modules(): if name == self.layer_name: - layer.register_forward_hook(get_features(name)) + self._activation_hook = layer.register_forward_hook( + get_features(name) + ) break - for instance in dataset: - video_id = instance["id"] - frames = instance["data"].to(self.device) - embeddings[video_id] = [] - - frames = frames.unsqueeze(0).permute(0, 2, 1, 3, 4) - if frames.shape[2] < 14: - pad_width = (0, 0, 0, 0, 0, 14 - frames.shape[2], 0, 0, 0, 0) - frames = torch.nn.functional.pad(frames, pad_width, mode="constant") - _ = self.model(frames) - values = activation - pooled = torch.nn.functional.adaptive_avg_pool2d(values, (1, 1)) - - embeddings[video_id] = ( - torch.flatten(pooled, 1).detach().cpu().numpy().flatten() - ) + dataset = CustomDataset(modality.data, self.data_type, "cpu") + lengths = [ + max(length, 14) + for length in get_sequence_lengths(modality.data, modality.metadata) + ] + dataloader = torch.utils.data.DataLoader( + dataset, + batch_sampler=LengthBucketBatchSampler( + lengths, self.batch_size, exact=True + ), + collate_fn=self._collate_videos, + pin_memory=pin_memory_for(self.device), + ) + embeddings = [None] * len(dataset) + + with inference_context(self.device): + for batch in dataloader: + batch = move_batch_to_device(batch, self.device) + video_ids = batch["id"].long() + frames = batch["data"].permute(0, 2, 1, 3, 4) + _ = self.model(frames) + values = self.activation + if isinstance(values, tuple): + values = values[0] + if values.ndim > 2: + values = torch.nn.functional.adaptive_avg_pool2d(values, (1, 1)) + vectors = torch.flatten(values, 1).detach().float().cpu().numpy() + for video_id, vector in zip(video_ids.cpu().tolist(), vectors): + embeddings[video_id] = vector + + if self.output_file is not None: + save_embeddings(embeddings, self.output_file) transformed_modality = TransformedModality( modality, self, self.output_modality_type ) - - transformed_modality.data = list(embeddings.values()) - + transformed_modality.data = embeddings + transformed_modality.data_type = np.float32 return transformed_modality class I3D(UnimodalRepresentation): - def __init__(self, layer="blocks.6", model_name="i3d", output_file=None): + _EMBEDDING_DIM = 400 + cache_in_worker = True + + def __init__( + self, + layer="blocks.6", + model_name="i3d", + output_file=None, + batch_size=8, + params=None, + ): + if params is not None: + layer = params.get("layer_name", layer) + batch_size = int(params.get("batch_size", batch_size)) self.model_name = model_name self.model = torch.hub.load( "facebookresearch/pytorchvideo", "i3d_r50", pretrained=True @@ -230,12 +305,42 @@ def __init__(self, layer="blocks.6", model_name="i3d", output_file=None): self.output_file = output_file self.layer_name = layer + self.batch_size = batch_size + self.data_type = torch.float32 + self._gpu_id = self.device.index + self._activation_hook = None + self.features = None self.model.eval() for param in self.model.parameters(): param.requires_grad = False + @property + def gpu_id(self): + return self._gpu_id + + @gpu_id.setter + def gpu_id(self, gpu_id): + self._gpu_id = gpu_id + self.device = get_device(gpu_id) + if self.model is not None: + self.model = self.model.to(self.device) + + def get_output_stats(self, input_stats) -> RepresentationStats: + return RepresentationStats( + input_stats.num_instances, + (self._EMBEDDING_DIM,), + output_shape_is_known=self.layer_name == "blocks.6", + dtype=self.data_type, + ) + + def estimate_output_memory_bytes(self, input_stats: VideoStats) -> int: + return input_stats.num_instances * self._EMBEDDING_DIM * self.data_type.itemsize + def _get_parameters(self, high_level=True): - parameters = {"layer_name": []} + parameters = { + "batch_size": [1, 2, 4, 8, 16, 32], + "layer_name": [], + } if high_level: parameters["layer_name"] = [ @@ -252,52 +357,58 @@ def _get_parameters(self, high_level=True): parameters["layer_name"].append(name) return parameters - def transform(self, modality): - sample = modality.data[0] if modality.data else "" - self.batch_size = compute_batch_size( - model=self.model, - device=self.device, - sample_data=sample, - tokenizer=None, - max_seq_length=None, - max_batch_size=128, - ) - dataset = CustomDataset(modality.data, torch.float32, self.device) - embeddings = {} - - features = None + def transform(self, modality, aggregation=None): + self.model = self.model.to(self.device) + self.model.eval() + self.features = None def get_features(name_): def hook( _module: torch.nn.Module, input_: Tuple[torch.Tensor], output: Any ): - # pooled = torch.nn.functional.adaptive_avg_pool3d(output, 1).squeeze() - nonlocal features - features = output.detach().cpu().numpy() + self.features = output return hook - if self.layer_name: + if self.layer_name and self._activation_hook is None: for name, layer in self.model.named_modules(): if name == self.layer_name: - layer.register_forward_hook(get_features(name)) + self._activation_hook = layer.register_forward_hook( + get_features(name) + ) break - for instance in dataset: - video_id = instance["id"] - frames = instance["data"].to(self.device) - embeddings[video_id] = [] - - batch = torch.transpose(frames, 1, 0) - batch = batch.unsqueeze(0) - _ = self.model(batch) - - embeddings[video_id] = features.flatten() + dataset = CustomDataset(modality.data, self.data_type, "cpu") + dataloader = torch.utils.data.DataLoader( + dataset, + batch_sampler=LengthBucketBatchSampler( + get_sequence_lengths(modality.data, modality.metadata), + self.batch_size, + exact=True, + ), + pin_memory=pin_memory_for(self.device), + ) + embeddings = [None] * len(dataset) + + with inference_context(self.device): + for batch in dataloader: + batch = move_batch_to_device(batch, self.device) + video_ids = batch["id"].long() + frames = batch["data"].permute(0, 2, 1, 3, 4) + _ = self.model(frames) + values = self.features + if isinstance(values, tuple): + values = values[0] + vectors = torch.flatten(values, 1).detach().float().cpu().numpy() + for video_id, vector in zip(video_ids.cpu().tolist(), vectors): + embeddings[video_id] = vector + + if self.output_file is not None: + save_embeddings(embeddings, self.output_file) transformed_modality = TransformedModality( modality, self, self.output_modality_type ) - - transformed_modality.data = list(embeddings.values()) - + transformed_modality.data = embeddings + transformed_modality.data_type = np.float32 return transformed_modality diff --git a/src/main/python/systemds/scuro/utils/checkpointing.py b/src/main/python/systemds/scuro/utils/checkpointing.py index e821798534c..2928dd0963a 100644 --- a/src/main/python/systemds/scuro/utils/checkpointing.py +++ b/src/main/python/systemds/scuro/utils/checkpointing.py @@ -123,7 +123,7 @@ def save(self, results: Any, meta: Dict[str, Any]) -> str: def save_checkpoint(self, results: Any, extra_meta: Dict[str, Any] = {}): meta = {"eval_count": self.eval_count} - # meta.update(extra_meta or {}) + meta.update(extra_meta or {}) self.save(results, meta) def checkpoint_if_due(self, results: Any, extra_meta: Dict[str, Any] = None): diff --git a/src/main/python/systemds/scuro/utils/memory_utility.py b/src/main/python/systemds/scuro/utils/memory_utility.py index 88698fa53cc..ece738d050e 100644 --- a/src/main/python/systemds/scuro/utils/memory_utility.py +++ b/src/main/python/systemds/scuro/utils/memory_utility.py @@ -167,8 +167,28 @@ def get_gpu_memory_mb(device): def gpu_memory_info(): - infos = [] num_gpus = torch.cuda.device_count() + if num_gpus == 0: + return [] + try: + import pynvml + + pynvml.nvmlInit() + try: + infos = [] + for i in range(num_gpus): + handle = pynvml.nvmlDeviceGetHandleByIndex(i) + mem = pynvml.nvmlDeviceGetMemoryInfo(handle) + infos.append( + dict(index=i, free_b=int(mem.free), total_b=int(mem.total)) + ) + return infos + finally: + pynvml.nvmlShutdown() + except Exception: + pass + + infos = [] for i in range(num_gpus): torch.cuda.set_device(i) free_b, total_b = torch.cuda.mem_get_info() diff --git a/src/main/python/tests/scuro/data_generator.py b/src/main/python/tests/scuro/data_generator.py index b30946fb7df..d5965d219b1 100644 --- a/src/main/python/tests/scuro/data_generator.py +++ b/src/main/python/tests/scuro/data_generator.py @@ -76,13 +76,13 @@ def __init__(self, indices, chunk_size, modality_type, data, data_type, metadata 30, max(d.shape[0] for d in data), sum(d.shape[0] for d in data) / len(data), - max(d.shape[1] for d in data), max(d.shape[2] for d in data), + max(d.shape[1] for d in data), max(d.shape[3] for d in data), chunk_size if chunk_size is not None else len(data), len(data), ) - elif modality_type == ModalityType.TIMESERIES: + elif modality_type in (ModalityType.TIMESERIES, ModalityType.PHYSIOLOGICAL): self.stats = TimeseriesStats( max(len(d) for d in data), len(data), @@ -94,8 +94,8 @@ def __init__(self, indices, chunk_size, modality_type, data, data_type, metadata ) elif modality_type == ModalityType.IMAGE: self.stats = ImageStats( - max(d.shape[0] for d in data), max(d.shape[1] for d in data), + max(d.shape[0] for d in data), max(d.shape[2] for d in data), len(data), ( @@ -103,8 +103,8 @@ def __init__(self, indices, chunk_size, modality_type, data, data_type, metadata max(d.shape[1] for d in data), max(d.shape[2] for d in data), ), - average_width=sum(d.shape[0] for d in data) / len(data), - average_height=sum(d.shape[1] for d in data) / len(data), + average_width=sum(d.shape[1] for d in data) / len(data), + average_height=sum(d.shape[0] for d in data) / len(data), average_channels=sum(d.shape[2] for d in data) / len(data), ) @@ -255,6 +255,43 @@ def create_timeseries_data(self, num_instances, sequence_length, num_features=1) ] return data, self.metadata + def create_physiological_data( + self, num_instances, sequence_length, kind="ecg", fs=500.0 + ): + self.modality_type = ModalityType.PHYSIOLOGICAL + rng = np.random.default_rng(7) + data = [] + + for _ in range(num_instances): + t = np.arange(sequence_length) / fs + samples = np.arange(sequence_length) + if kind == "ecg": + signal = rng.normal(0.0, 0.01, sequence_length) + width = max(1.0, 0.02 * fs) + position = int(0.2 * fs) + while position < sequence_length: + signal += np.exp(-(((samples - position) / width) ** 2)) + position += int(rng.uniform(0.7, 0.9) * fs) + elif kind == "eda": + signal = 0.5 + 0.01 * t + rng.normal(0.0, 0.005, sequence_length) + width = max(1.0, 1.5 * fs) + for peak_time in np.arange(5.0, max(t[-1], 5.0), 10.0): + centre = peak_time * fs + signal += np.exp(-(((samples - centre) / width) ** 2)) + elif kind == "resp": + signal = np.sin(2 * np.pi * 0.25 * t) + rng.normal( + 0.0, 0.02, sequence_length + ) + else: + raise ValueError(f"Unsupported physiological signal kind: {kind}") + data.append(signal.astype(self.data_type)) + + self.metadata = [ + self.modality_type.create_metadata(["signal"], data[i]) + for i in range(num_instances) + ] + return data, self.metadata + def create_text_data(self, num_instances, num_sentences_per_instance=1): self.modality_type = ModalityType.TEXT subjects = [ diff --git a/src/main/python/tests/scuro/test_chunked_leaf_execution.py b/src/main/python/tests/scuro/test_chunked_leaf_execution.py new file mode 100644 index 00000000000..91bf0a2f2bf --- /dev/null +++ b/src/main/python/tests/scuro/test_chunked_leaf_execution.py @@ -0,0 +1,352 @@ +# ------------------------------------------------------------- +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# +# ------------------------------------------------------------- + +import os +from types import SimpleNamespace +import unittest + +import numpy as np + + +def _skip_if_session_uses_fork(test_case): + if test_case._session_uses_fork: + test_case.skipTest( + "session is running under SCURO_MP_CONTEXT=fork; creating a " + "worker pool here deadlocks the CUDA-using tests later in the " + "session" + ) + + +from systemds.scuro.drsearch.modality_shared_memory import unlink_shm +from systemds.scuro.drsearch.node_executor import ( + NodeExecutor, + _execute_leaf_batch_worker, +) +from systemds.scuro.drsearch.representation_dag import ( + CSEAwareDAGBuilder, + RepresentationDag, + RepresentationNode, +) +from systemds.scuro.drsearch.task import PerformanceMeasure +from systemds.scuro.modality.transformed import TransformedModality +from systemds.scuro.modality.type import ModalityType +from systemds.scuro.modality.unimodal_modality import UnimodalModality +from systemds.scuro.representations.representation import RepresentationStats +from systemds.scuro.representations.unimodal import UnimodalRepresentation +from tests.scuro.data_generator import TestDataLoader + +NUM_INSTANCES = 12 +CHUNK_SIZE = 4 + + +def _make_modality(chunk_size): + rng = np.random.default_rng(0) + data = [rng.random(160, dtype=np.float32) for _ in range(NUM_INSTANCES)] + metadata = [ + ModalityType.AUDIO.create_metadata(16000, data[i]) for i in range(NUM_INSTANCES) + ] + loader = TestDataLoader( + indices=np.arange(NUM_INSTANCES), + chunk_size=chunk_size, + modality_type=ModalityType.AUDIO, + data=data, + data_type=np.float32, + metadata=metadata, + ) + return UnimodalModality(data_loader=loader) + + +class CountingOperation(UnimodalRepresentation): + """Emits one 4-vector per instance it is given.""" + + def __init__(self, params=None): + super().__init__("CountingOperation", ModalityType.EMBEDDING) + + def get_output_stats(self, input_stats): + return RepresentationStats(input_stats.num_instances, (4,)) + + def estimate_memory_bytes(self, input_stats): + return 1024 + + def estimate_peak_memory_bytes(self, input_stats): + return {"cpu_peak_bytes": 1024, "gpu_peak_bytes": 0} + + def transform(self, modality, aggregation=None): + transformed = TransformedModality( + modality, self, self.output_modality_type, set_data=False + ) + n = len(modality.data) + transformed._data = [np.full(4, float(n), dtype=np.float32) for _ in range(n)] + return transformed + + +class FrameOperation(UnimodalRepresentation): + """A same-named frame encoder that honors pushed-down aggregation.""" + + def __init__(self, params=None): + super().__init__("FrameOperation", ModalityType.EMBEDDING) + + def transform(self, modality, aggregation=None): + transformed = TransformedModality( + modality, self, self.output_modality_type, set_data=False + ) + frame_embedding = np.arange(12, dtype=np.float32).reshape(3, 4) + transformed._data = [frame_embedding.copy() for _ in modality.data] + if aggregation is not None: + return aggregation.transform(transformed) + return transformed + + def get_output_stats(self, input_stats): + return RepresentationStats(input_stats.num_instances, (3, 4)) + + def estimate_memory_bytes(self, input_stats): + return 1024 + + def estimate_peak_memory_bytes(self, input_stats): + return {"cpu_peak_bytes": 1024, "gpu_peak_bytes": 0} + + +class RaggedFrameOperation(FrameOperation): + """Emits variable-length frame sequences without aggregation.""" + + def __init__(self, params=None): + super().__init__(params=params) + self.name = "RaggedFrameOperation" + + def transform(self, modality, aggregation=None): + transformed = TransformedModality( + modality, self, self.output_modality_type, set_data=False + ) + transformed._data = [ + np.full((index % 3 + 1, 4), index, dtype=np.float32) + for index in range(len(modality.data)) + ] + if aggregation is not None: + return aggregation.transform(transformed) + return transformed + + +class InstanceCountingTask: + """Reports how many instances actually reached the task.""" + + def estimate_peak_memory_bytes(self, input_stats): + return {"cpu_peak_bytes": 1024, "gpu_peak_bytes": 0} + + def get_output_stats(self, input_stats): + return RepresentationStats(input_stats.num_instances, (1,)) + + def run(self, data): + count = float(len(data)) + scores = [] + for split in ("train", "val", "test"): + measure = PerformanceMeasure(split, "accuracy") + measure.scores["accuracy"] = [count] + scores.append(measure.compute_averages()) + return scores + + +def _build_dag(modality): + builder = CSEAwareDAGBuilder() + leaf_id = builder.create_leaf_node(modality_id=modality.modality_id) + op_id = builder.create_operation_node(CountingOperation, [leaf_id], {}) + dag = builder.build(op_id) + + task_root_id = f"task_{dag.root_node_id}_0" + task_node = RepresentationNode( + node_id=task_root_id, + operation=None, + inputs=[dag.root_node_id], + parameters={ + "_node_kind": "task", + "_task_idx": 0, + "_dag_root_id": dag.root_node_id, + }, + ) + return [RepresentationDag(nodes=[*dag.nodes, task_node], root_node_id=task_root_id)] + + +class TestChunkedLeafExecution(unittest.TestCase): + def setUp(self): + self._session_uses_fork = os.environ.get("SCURO_MP_CONTEXT") == "fork" + previous = os.environ.get("SCURO_MP_CONTEXT") + os.environ["SCURO_MP_CONTEXT"] = "spawn" + if previous is None: + self.addCleanup(os.environ.pop, "SCURO_MP_CONTEXT", None) + else: + self.addCleanup(os.environ.__setitem__, "SCURO_MP_CONTEXT", previous) + + def _executor(self, modality): + """A NodeExecutor whose pool is torn down even if the test fails.""" + _skip_if_session_uses_fork(self) + executor = NodeExecutor( + dags=_build_dag(modality), + modalities=[modality], + tasks=[InstanceCountingTask()], + max_num_workers=2, + enable_checkpointing=False, + ) + self.addCleanup(executor._pool.shutdown) + return executor + + def test_chunked_leaf_is_not_preloaded(self): + """A streaming modality must not be materialised before scheduling. + + `has_data()` staying False is the observable consequence: the executor + left the leaf alone, and the chunk loop inside `apply_representations` + is what reads the data. + """ + modality = _make_modality(chunk_size=CHUNK_SIZE) + executor = self._executor(modality) + self.assertTrue(executor._loads_in_chunks(modality)) + + executor._load_leaf_modalities() + self.assertFalse( + modality.has_data(), + "chunked leaf was preloaded; BaseLoader.load() would have " + "returned only the first chunk", + ) + + def test_chunked_subset_remains_lazy_and_uses_full_dataset_indices(self): + """Test-only subsets must not index into whichever chunk is resident.""" + modality = _make_modality(chunk_size=CHUNK_SIZE) + modality.extract_raw_data() + self.assertEqual(len(modality.data), CHUNK_SIZE) + + subset_indices = [1, 5, 10] + subset = modality.subset(subset_indices) + + self.assertFalse(subset.has_data()) + self.assertEqual( + subset.data_loader.indices, + [modality.data_loader.indices[i] for i in subset_indices], + ) + transformed = subset.apply_representations([CountingOperation()]) + self.assertEqual(len(transformed["CountingOperation"].data), 3) + + def test_unchunked_leaf_is_still_preloaded(self): + """The skip must be narrow: a non-streaming leaf still loads up front.""" + modality = _make_modality(chunk_size=None) + executor = self._executor(modality) + self.assertFalse(executor._loads_in_chunks(modality)) + + executor._load_leaf_modalities() + self.assertTrue(modality.has_data()) + self.assertEqual(len(modality.data), NUM_INSTANCES) + executor._cleanup_leaf_shared_memory() + + def test_unchunked_ragged_representation_is_padded(self): + modality = _make_modality(chunk_size=None) + transformed = modality.apply_representation(RaggedFrameOperation()) + + self.assertEqual(len(transformed.data), NUM_INSTANCES) + self.assertTrue(all(value.shape == (3, 4) for value in transformed.data)) + np.testing.assert_array_equal(transformed.data[0][1:], np.zeros((2, 4))) + masks = [metadata["attention_masks"] for metadata in transformed.metadata] + np.testing.assert_array_equal(masks[0], np.array([1.0, 0.0, 0.0])) + np.testing.assert_array_equal(masks[1], np.array([1.0, 1.0, 0.0])) + np.testing.assert_array_equal(masks[2], np.array([1.0, 1.0, 1.0])) + + def test_chunked_run_sees_every_instance(self): + """End to end, the search must score the whole dataset. + + This one holds with or without the preload -- `iter_raw_data_chunks` + resets the loader and re-reads everything, so the preload wasted time + and memory rather than truncating results. It is here to pin that + skipping the preload did not cost coverage of the whole dataset. + """ + modality = _make_modality(chunk_size=CHUNK_SIZE) + executor = self._executor(modality) + result = executor.run() + + self.assertEqual(len(result["task_results"]), 1) + entry = result["task_results"][0] + self.assertIsNotNone(entry.val_score, "candidate produced no score at all") + self.assertEqual( + entry.val_score["accuracy"], + float(NUM_INSTANCES), + "the task saw a partial dataset", + ) + + def test_chunked_run_produces_one_metadata_entry_per_instance(self): + """Metadata must not be double-counted. + + TransformedModality seeds its metadata from the source modality's and + the chunk loop appends one entry per instance on top, so a leaf that + arrived carrying preloaded metadata produced len(chunk) extra entries. + """ + modality = _make_modality(chunk_size=CHUNK_SIZE) + # Exactly the state the old preload left the leaf in: carrying the + # first chunk's data and metadata. Without this the modality starts + # empty and the doubling cannot occur, so the test would pass either + # way and prove nothing. + modality.extract_raw_data() + self.assertEqual(len(modality.metadata), CHUNK_SIZE) + + transformed = modality.apply_representations([CountingOperation()]) + out = transformed["CountingOperation"] + + self.assertEqual(len(out.data), NUM_INSTANCES) + self.assertEqual( + len(out.metadata), + NUM_INSTANCES, + "metadata was seeded from the leaf and then appended to per " + "instance, so the preloaded chunk got counted twice", + ) + + def test_leaf_batch_keeps_same_named_nodes_and_pushes_down_aggregation(self): + """Batched frame encoders must yield one 2-D result per DAG node.""" + modality = _make_modality(chunk_size=CHUNK_SIZE) + aggregation = { + "aggregation": "mean", + "target_dimensions": 1, + "aggregate_leading": True, + } + nodes = [ + SimpleNamespace( + node_id=f"frame_node_{index}", + operation=FrameOperation, + parameters={"_pushdown_aggregation": aggregation}, + ) + for index in range(2) + ] + + value = _execute_leaf_batch_worker(nodes, modality, gpu_id=None) + self.addCleanup( + lambda: [ + unlink_shm(info["shm_name"]) + for info in value["shm_info"].values() + if info.get("shm_name") is not None + ] + ) + + self.assertEqual(set(value["results"]), {node.node_id for node in nodes}) + self.assertEqual(value["failed_nodes"], {}) + for transformed in value["results"].values(): + self.assertEqual(len(transformed.data), NUM_INSTANCES) + self.assertEqual( + np.asarray(transformed.data).shape, + (NUM_INSTANCES, 4), + "pushed-down frame aggregation was lost, leaving a 3-D result", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/main/python/tests/scuro/test_lazy_visual_loading.py b/src/main/python/tests/scuro/test_lazy_visual_loading.py new file mode 100644 index 00000000000..e2da64394fb --- /dev/null +++ b/src/main/python/tests/scuro/test_lazy_visual_loading.py @@ -0,0 +1,177 @@ +# ------------------------------------------------------------- +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# +# ------------------------------------------------------------- + +import shutil +import unittest +from unittest.mock import patch + +import numpy as np +import torch + +from systemds.scuro.dataloader.base_loader import LazyFileSequence +from systemds.scuro.dataloader.image_loader import ImageLoader +from systemds.scuro.dataloader.video_loader import VideoLoader +from systemds.scuro.modality.type import ModalityType +from systemds.scuro.modality.unimodal_modality import UnimodalModality +from systemds.scuro.representations.color_histogram import ColorHistogram +from systemds.scuro.representations.optical_flow import OpticalFlow +from systemds.scuro.representations.utils import flatten_owned_sequences +from systemds.scuro.utils.torch_dataset import CustomDataset +from tests.scuro.data_generator import setup_data + + +class TestLazyVisualLoading(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.test_file_path = "test_lazy_visual_data" + cls.data_generator = setup_data( + [ModalityType.IMAGE, ModalityType.VIDEO], + 2, + cls.test_file_path, + ) + + @classmethod + def tearDownClass(cls): + shutil.rmtree(cls.test_file_path, ignore_errors=True) + + def _loader(self, loader_type, indices=None, **kwargs): + modality_type = ( + ModalityType.IMAGE if loader_type is ImageLoader else ModalityType.VIDEO + ) + return loader_type( + self.data_generator.get_modality_path(modality_type), + indices or self.data_generator.indices, + **kwargs, + ) + + def test_unchunked_load_keeps_only_file_references(self): + for loader_type in (ImageLoader, VideoLoader): + with self.subTest(loader=loader_type.__name__): + loader = self._loader(loader_type) + with patch.object( + loader, "_decode_file", wraps=loader._decode_file + ) as decode: + data, metadata = loader.load() + + self.assertIsInstance(data, LazyFileSequence) + self.assertEqual(len(data), 2) + self.assertEqual(len(metadata), 2) + decode.assert_not_called() + + self.assertIsInstance(data[0], np.ndarray) + decode.assert_called_once() + + def test_custom_dataset_decodes_only_the_requested_batch(self): + loader = self._loader(ImageLoader) + with patch.object(loader, "_decode_file", wraps=loader._decode_file) as decode: + data, _ = loader.load() + dataset = CustomDataset(data, torch.float32, "cpu") + dataloader = torch.utils.data.DataLoader(dataset, batch_size=1) + + first_batch = next(iter(dataloader)) + + self.assertEqual(first_batch["data"].shape[0], 1) + decode.assert_called_once() + + def test_modality_subset_preserves_lazy_file_references(self): + loader = self._loader(ImageLoader) + modality = UnimodalModality(loader) + with patch.object(loader, "_decode_file", wraps=loader._decode_file) as decode: + modality.extract_raw_data() + subset = modality.subset([1]) + + self.assertIsInstance(subset.data, LazyFileSequence) + self.assertEqual(len(subset.data), 1) + decode.assert_not_called() + self.assertIsInstance(subset.data[0], np.ndarray) + decode.assert_called_once() + + def test_unchunked_peak_memory_is_not_the_full_corpus(self): + for loader_type in (ImageLoader, VideoLoader): + loader = self._loader(loader_type) + modality = UnimodalModality(loader) + peak = modality.estimate_peak_memory_bytes()["cpu_peak_bytes"] + total = modality.estimate_memory_bytes() + self.assertLess(peak, total) + + def test_chunked_loading_still_returns_decoded_chunks(self): + loader = self._loader(ImageLoader, chunk_size=1) + + first_data, first_metadata = loader.load() + second_data, second_metadata = loader.load() + + self.assertIsInstance(first_data, list) + self.assertNotIsInstance(first_data, LazyFileSequence) + self.assertEqual(len(first_data), 1) + self.assertEqual(len(first_metadata), 1) + self.assertEqual(len(second_data), 1) + self.assertEqual(len(second_metadata), 1) + + def test_histogram_streams_lazy_images_and_videos(self): + for loader_type in (ImageLoader, VideoLoader): + with self.subTest(loader=loader_type.__name__): + loader = self._loader(loader_type) + modality = UnimodalModality(loader) + with patch.object( + loader, "_decode_file", wraps=loader._decode_file + ) as decode: + transformed = modality.apply_representation( + ColorHistogram(bins=4, normalize=True) + ) + + self.assertEqual(len(transformed.data), 2) + self.assertEqual(decode.call_count, 2) + + def test_optical_flow_streams_one_lazy_video_at_a_time(self): + loader = self._loader(VideoLoader, indices=self.data_generator.indices[:1]) + modality = UnimodalModality(loader) + with patch.object(loader, "_decode_file", wraps=loader._decode_file) as decode: + transformed = modality.apply_representation(OpticalFlow()) + + self.assertEqual(len(transformed.data), 1) + self.assertEqual(len(transformed.data[0]), loader.stats.max_length - 1) + decode.assert_called_once() + + def test_flattened_frame_view_caches_only_the_current_owner(self): + class CountingSequences: + def __init__(self): + self.values = [ + np.arange(2).reshape(2, 1), + np.arange(3).reshape(3, 1), + ] + self.reads = [] + + def __getitem__(self, index): + self.reads.append(index) + return self.values[index] + + def __len__(self): + return len(self.values) + + sequences = CountingSequences() + frames, owner_ids = flatten_owned_sequences(sequences, [2, 3]) + + np.testing.assert_array_equal(frames[0], [0]) + np.testing.assert_array_equal(frames[1], [1]) + self.assertEqual(sequences.reads, [0]) + np.testing.assert_array_equal(frames[2], [0]) + self.assertEqual(sequences.reads, [0, 1]) + self.assertEqual(owner_ids, [0, 0, 1, 1, 1]) diff --git a/src/main/python/tests/scuro/test_modality_pad.py b/src/main/python/tests/scuro/test_modality_pad.py new file mode 100644 index 00000000000..b6c0001250c --- /dev/null +++ b/src/main/python/tests/scuro/test_modality_pad.py @@ -0,0 +1,105 @@ +# ------------------------------------------------------------- +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# +# ------------------------------------------------------------- + +import unittest + +import numpy as np + +from systemds.scuro.modality.modality import Modality +from systemds.scuro.modality.type import ModalityType +from systemds.scuro.representations.sum import Sum + + +def _embedding_metadata(embedding_dim): + return [ + { + "data_layout": { + "shape": (embedding_dim,), + "type": np.float32, + "representation": "embedding", + } + } + ] + + +class TestModalityPad(unittest.TestCase): + def test_pad_single_instance_1d_embedding(self): + modality = Modality( + ModalityType.EMBEDDING, + modality_id=1, + metadata=_embedding_metadata(4), + data_type=np.float32, + ) + modality._data = np.array([1.0, 2.0, 3.0, 4.0], dtype=np.float32) + + modality.pad(max_len=6) + + self.assertEqual(modality.data.shape, (1, 6)) + np.testing.assert_array_equal( + modality.data[0, :4], np.array([1.0, 2.0, 3.0, 4.0], dtype=np.float32) + ) + np.testing.assert_array_equal(modality.data[0, 4:], np.zeros(2)) + + def test_pad_2d_embedding_columns(self): + modality = Modality( + ModalityType.EMBEDDING, + modality_id=1, + metadata=_embedding_metadata(3) * 2, + data_type=np.float32, + ) + modality._data = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], dtype=np.float32) + + modality.pad(max_len=5) + + self.assertEqual(modality.data.shape, (2, 5)) + np.testing.assert_array_equal(modality.data[0, :3], np.array([1.0, 2.0, 3.0])) + np.testing.assert_array_equal(modality.data[1, :3], np.array([4.0, 5.0, 6.0])) + + def test_fusion_sum_aligns_mismatched_embedding_sizes(self): + metadata_a = _embedding_metadata(4) + metadata_b = _embedding_metadata(6) + + modality_a = Modality( + ModalityType.EMBEDDING, + modality_id=1, + metadata=metadata_a, + data_type=np.float32, + ) + modality_a._data = np.array([1.0, 2.0, 3.0, 4.0], dtype=np.float32) + + modality_b = Modality( + ModalityType.EMBEDDING, + modality_id=2, + metadata=metadata_b, + data_type=np.float32, + ) + modality_b._data = np.array([[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]], dtype=np.float32) + + fused = Sum().transform([modality_a, modality_b]) + + self.assertEqual(fused.shape, (1, 6)) + np.testing.assert_array_equal( + fused[0], np.array([2.0, 4.0, 6.0, 8.0, 5.0, 6.0], dtype=np.float32) + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/main/python/tests/scuro/test_neural_encoder_batching.py b/src/main/python/tests/scuro/test_neural_encoder_batching.py new file mode 100644 index 00000000000..3094e15f1a2 --- /dev/null +++ b/src/main/python/tests/scuro/test_neural_encoder_batching.py @@ -0,0 +1,156 @@ +# ------------------------------------------------------------- +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# +# ------------------------------------------------------------- +import unittest +from types import SimpleNamespace + +import numpy as np +import torch + +from systemds.scuro.modality.type import ModalityType +from systemds.scuro.representations.aggregated_representation import ( + AggregatedRepresentation, +) +from systemds.scuro.representations.clip import CLIPVisual +from systemds.scuro.representations.resnet import ResNet +from systemds.scuro.representations.vgg import VGG19 + + +class _ResNetModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.scale = torch.nn.Parameter(torch.tensor(1.0)) + self.avgpool = torch.nn.Identity() + + def forward(self, images): + values = images.flatten(2).mean(dim=2)[:, :2] * self.scale + return self.avgpool(values[:, :, None, None]) + + +class _VGGModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.scale = torch.nn.Parameter(torch.tensor(1.0)) + self.classifier = torch.nn.Sequential(torch.nn.Identity()) + + def forward(self, images): + values = images.flatten(2).mean(dim=2)[:, :2] * self.scale + return self.classifier[0](values) + + +def _video_modality(videos): + height, width, channels = videos[0][0].shape + return SimpleNamespace( + modality_type=ModalityType.VIDEO, + modality_id=0, + metadata=[ + ModalityType.VIDEO.create_metadata(30, len(video), width, height, channels) + for video in videos + ], + data_type=np.float32, + transform_time=0, + data=videos, + ) + + +def _representation(representation_class, batch_size): + representation = object.__new__(representation_class) + representation.data_type = torch.float32 + representation.device = torch.device("cpu") + representation.batch_size = batch_size + representation._activation_hook = None + representation.activation = None + representation.output_modality_type = ModalityType.EMBEDDING + if representation_class is ResNet: + representation.model = _ResNetModel() + representation.layer_name = "avgpool" + else: + representation.model = _VGGModel() + representation.layer_name = "classifier.0" + return representation + + +class TestNeuralEncoderBatching(unittest.TestCase): + def setUp(self): + self.videos = [ + [ + np.full((8, 8, 3), 32, dtype=np.uint8), + np.full((8, 8, 3), 64, dtype=np.uint8), + ], + [ + np.full((8, 8, 3), 96, dtype=np.uint8), + np.full((8, 8, 3), 128, dtype=np.uint8), + np.full((8, 8, 3), 160, dtype=np.uint8), + ], + ] + self.aggregation = AggregatedRepresentation("mean") + + def test_global_frame_batching_is_invariant_to_batch_size(self): + for representation_class in (ResNet, VGG19): + with self.subTest(representation=representation_class.__name__): + results = [ + _representation(representation_class, batch_size) + .transform(_video_modality(self.videos), self.aggregation) + .data + for batch_size in (1, 4) + ] + self.assertEqual(results[0].shape, (len(self.videos), 2)) + np.testing.assert_allclose(results[1], results[0]) + + def test_global_frame_batching_matches_patient_by_patient(self): + for representation_class in (ResNet, VGG19): + with self.subTest(representation=representation_class.__name__): + global_result = ( + _representation(representation_class, 3) + .transform(_video_modality(self.videos), self.aggregation) + .data + ) + per_patient = np.concatenate( + [ + _representation(representation_class, 3) + .transform(_video_modality([video]), self.aggregation) + .data + for video in self.videos + ], + axis=0, + ) + np.testing.assert_allclose(global_result, per_patient) + + def test_frame_encoders_support_aggregation_pushdown(self): + input_stats = SimpleNamespace(num_instances=3) + for representation_class, hidden_dim in ( + (ResNet, 512), + (VGG19, 4096), + (CLIPVisual, 512), + ): + with self.subTest(representation=representation_class.__name__): + representation = object.__new__(representation_class) + representation.params = {"_pushdown_aggregation": {}} + representation.data_type = torch.float32 + + self.assertTrue(representation.supports_aggregation_pushdown) + stats = representation.get_output_stats(input_stats) + self.assertEqual(stats.num_instances, input_stats.num_instances) + self.assertEqual(stats.output_shape, (hidden_dim,)) + self.assertIsNone(stats.aggregate_dim) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/main/python/tests/scuro/test_operator_registry.py b/src/main/python/tests/scuro/test_operator_registry.py index 93afba342b0..99bfeb2f2a4 100644 --- a/src/main/python/tests/scuro/test_operator_registry.py +++ b/src/main/python/tests/scuro/test_operator_registry.py @@ -75,7 +75,10 @@ from systemds.scuro.representations.mel_spectrogram import MelSpectrogram from systemds.scuro.representations.spectrogram import Spectrogram from systemds.scuro.representations.hadamard import Hadamard +from systemds.scuro.representations.image_bind import ImageBind from systemds.scuro.representations.resnet import ResNet +from systemds.scuro.representations.openface import OpenFace +from systemds.scuro.representations.color_histogram import ColorHistogram from systemds.scuro.representations.multimodal_attention_fusion import AttentionFusion from systemds.scuro.representations.physiological_window import ( AdaptiveWindow, @@ -89,6 +92,7 @@ def test_audio_representations_in_registry(self): assert registry.get_representations(ModalityType.AUDIO) == [ MelSpectrogram, MFCC, + ImageBind, Spectrogram, Wav2Vec, Spectral, @@ -101,8 +105,11 @@ def test_video_representations_in_registry(self): registry = Registry() assert registry.get_representations(ModalityType.VIDEO) == [ ResNet, + OpenFace, + ImageBind, SwinVideoTransformer, X3D, + ColorHistogram, VGG19, CLIPVisual, ] diff --git a/src/main/python/tests/scuro/test_transformer_text_aggregation.py b/src/main/python/tests/scuro/test_transformer_text_aggregation.py new file mode 100644 index 00000000000..ded26859917 --- /dev/null +++ b/src/main/python/tests/scuro/test_transformer_text_aggregation.py @@ -0,0 +1,353 @@ +# ------------------------------------------------------------- +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# +# ------------------------------------------------------------- +import copy +import unittest +from types import SimpleNamespace + +import numpy as np +import torch + +from systemds.scuro.drsearch.representation_dag import ( + CSEAwareDAGBuilder, + pushdown_aggregation, +) +from systemds.scuro.representations.aggregated_representation import ( + AggregatedRepresentation, +) +from systemds.scuro.representations.bert import Bert +from systemds.scuro.representations.clip import CLIPText +from systemds.scuro.representations.representation import RepresentationStats +from systemds.scuro.representations.utils import pool_transformer_output + + +class _BatchEncoding(dict): + @property + def data(self): + return self + + def to(self, device): + for key, value in self.items(): + self[key] = value.to(device) + return self + + +class _Tokenizer: + def __call__(self, batch, **kwargs): + ids = torch.tensor([int(text) for text in batch]) + input_ids = ids.unsqueeze(1).repeat(1, 3) + attention_mask = torch.tensor([[1, 1, 0]]).repeat(len(batch), 1) + return _BatchEncoding( + input_ids=input_ids, + attention_mask=attention_mask, + offset_mapping=torch.zeros((len(batch), 3, 2), dtype=torch.long), + ) + + +class _BertModel: + def __call__(self, input_ids, attention_mask): + ids = input_ids[:, 0].float() + hidden = torch.stack( + ( + torch.stack((ids, ids + 10), dim=1), + torch.stack((ids + 100, ids + 200), dim=1), + torch.full((len(ids), 2), 1000.0, device=ids.device), + ), + dim=1, + ) + return SimpleNamespace(last_hidden_state=hidden) + + +class _DynamicTokenizer: + def __call__(self, batch, **kwargs): + tokens = [[int(token) for token in text.split()] for text in batch] + max_length = kwargs.get("max_length") + if max_length is not None: + tokens = [values[:max_length] for values in tokens] + + if kwargs.get("return_tensors") != "pt": + return _BatchEncoding( + input_ids=tokens, + attention_mask=[[1] * len(values) for values in tokens], + ) + + padded_length = max(len(values) for values in tokens) + input_ids = torch.zeros((len(tokens), padded_length), dtype=torch.long) + attention_mask = torch.zeros_like(input_ids) + for row, values in enumerate(tokens): + input_ids[row, : len(values)] = torch.tensor(values) + attention_mask[row, : len(values)] = 1 + return _BatchEncoding( + input_ids=input_ids, + attention_mask=attention_mask, + offset_mapping=torch.zeros( + (len(tokens), padded_length, 2), dtype=torch.long + ), + ) + + +class _IntermediateBertModel: + def __init__(self, representation): + self.representation = representation + + def __call__(self, input_ids, attention_mask): + values = input_ids.float() + hidden = torch.stack((values, values + 10), dim=-1) + hidden = torch.where( + attention_mask.unsqueeze(-1).bool(), + hidden, + torch.full_like(hidden, 1000), + ) + self.representation.bert_output = hidden + return SimpleNamespace(last_hidden_state=hidden) + + +class _CLIPProcessor: + def __call__(self, text, **kwargs): + ids = torch.tensor([int(value) for value in text]) + return _BatchEncoding( + input_ids=ids.unsqueeze(1), + attention_mask=torch.tensor([[1, 1, 0]]).repeat(len(text), 1), + ) + + +class _CLIPTextModel: + def __init__(self, representation): + self.representation = representation + + def __call__(self, input_ids, attention_mask): + ids = input_ids[:, 0].float() + self.representation.clip_output = torch.stack( + ( + torch.stack((ids, ids + 2), dim=1), + torch.stack((ids + 2, ids + 4), dim=1), + torch.full((len(ids), 2), 1000.0, device=ids.device), + ), + dim=1, + ) + + +class _CLIPModel: + def __init__(self, representation): + self.text_model = _CLIPTextModel(representation) + + +class TestTransformerTextAggregation(unittest.TestCase): + def test_token_pooling_selects_cls_or_masked_mean(self): + hidden = torch.tensor( + [ + [[1.0, 2.0], [3.0, 4.0], [1000.0, 1000.0]], + [[5.0, 6.0], [1000.0, 1000.0], [1000.0, 1000.0]], + ] + ) + mask = torch.tensor([[1, 1, 0], [1, 0, 0]]) + + np.testing.assert_allclose( + pool_transformer_output(hidden, mask, use_cls=True).numpy(), + [[1.0, 2.0], [5.0, 6.0]], + ) + np.testing.assert_allclose( + pool_transformer_output(hidden, mask).numpy(), + [[2.0, 3.0], [5.0, 6.0]], + ) + + def test_output_stats_describe_pooled_chunk_vectors(self): + raw_stats = SimpleNamespace(num_instances=3) + context_stats = RepresentationStats(3, (5, 77)) + + bert_plain = Bert().get_output_stats(raw_stats) + bert_context = Bert().get_output_stats(context_stats) + clip_plain = CLIPText().get_output_stats(raw_stats) + clip_context = CLIPText().get_output_stats(context_stats) + + self.assertEqual(bert_plain.output_shape, (768,)) + self.assertEqual(bert_context.output_shape, (5, 768)) + self.assertEqual(bert_context.aggregate_dim, (0,)) + self.assertEqual(clip_plain.output_shape, (512,)) + self.assertEqual(clip_context.output_shape, (5, 512)) + self.assertEqual(clip_context.aggregate_dim, (0,)) + + def test_bert_cls_aggregation_is_independent_of_batch_boundaries(self): + representation = Bert(batch_size=2, max_seq_length=3) + result = representation.create_embeddings( + ["0", "1", "2", "3", "4"], + _BertModel(), + _Tokenizer(), + AggregatedRepresentation("mean"), + ) + + np.testing.assert_allclose(result, [2.0, 12.0]) + self.assertEqual(result.shape, (2,)) + + def test_global_batching_returns_one_vector_per_patient(self): + chunks = ["1", "3 5", "7 9 11", "2", "4 6"] + owner_ids = [0, 0, 1, 2, 2] + representation = Bert(layer="intermediate", batch_size=3, max_seq_length=8) + result = representation.create_embeddings( + chunks, + _IntermediateBertModel(representation), + _DynamicTokenizer(), + AggregatedRepresentation("mean"), + owner_ids=owner_ids, + num_owners=3, + ) + + self.assertEqual(result.shape, (3, 2)) + + def test_global_batching_is_invariant_to_batch_size(self): + chunks = ["1", "3 5", "7 9 11", "2", "4 6"] + owner_ids = [0, 0, 1, 2, 2] + results = [] + for batch_size in (1, 2, 4): + representation = Bert( + layer="intermediate", + batch_size=batch_size, + max_seq_length=8, + ) + results.append( + representation.create_embeddings( + chunks, + _IntermediateBertModel(representation), + _DynamicTokenizer(), + AggregatedRepresentation("mean"), + owner_ids=owner_ids, + num_owners=3, + ) + ) + + for result in results[1:]: + np.testing.assert_allclose(result, results[0]) + + def test_dynamic_padding_does_not_change_embeddings(self): + short = "2 4" + representation = Bert(layer="intermediate", batch_size=2, max_seq_length=8) + model = _IntermediateBertModel(representation) + tokenizer = _DynamicTokenizer() + + alone = representation.create_embeddings([short], model, tokenizer)[0] + mixed = representation.create_embeddings( + [short, "10 20 30 40 50"], model, tokenizer + )[0] + + np.testing.assert_allclose(mixed, alone) + + def test_global_batching_matches_patient_by_patient_batching(self): + patient_chunks = [["1", "3 5"], ["7 9 11"], ["2", "4 6"]] + chunks = [chunk for patient in patient_chunks for chunk in patient] + owner_ids = [ + owner_id for owner_id, patient in enumerate(patient_chunks) for _ in patient + ] + aggregation = AggregatedRepresentation("mean") + representation = Bert(layer="intermediate", batch_size=3, max_seq_length=8) + model = _IntermediateBertModel(representation) + tokenizer = _DynamicTokenizer() + + global_result = representation.create_embeddings( + chunks, + model, + tokenizer, + aggregation, + owner_ids=owner_ids, + num_owners=len(patient_chunks), + ) + per_patient_result = np.stack( + [ + representation.create_embeddings(patient, model, tokenizer, aggregation) + for patient in patient_chunks + ] + ) + + np.testing.assert_allclose(global_result, per_patient_result) + + def test_clip_intermediate_layer_uses_masked_mean_before_chunk_aggregation(self): + representation = CLIPText(batch_size=2, layer_name="intermediate") + representation.processor = _CLIPProcessor() + result = representation.create_text_embeddings( + ["0", "2", "4"], + _CLIPModel(representation), + AggregatedRepresentation("mean"), + ) + + np.testing.assert_allclose(result, [3.0, 5.0]) + self.assertEqual(result.shape, (2,)) + + def test_pushdown_keeps_shared_plain_and_aggregated_paths_distinct(self): + builder = CSEAwareDAGBuilder() + leaf_id = builder.create_leaf_node("transformer_pushdown") + bert = Bert() + bert_id = builder.create_operation_node( + Bert, [leaf_id], bert.get_current_parameters() + ) + aggregation = AggregatedRepresentation( + "mean", target_dimensions=1, aggregate_leading=True + ) + aggregation_id = builder.create_operation_node( + AggregatedRepresentation, + [bert_id], + aggregation.get_current_parameters(), + ) + + plain_dag = builder.build(bert_id) + aggregated_dag = builder.build(aggregation_id) + aggregation_params = copy.deepcopy( + aggregated_dag.get_node_by_id(aggregation_id).parameters + ) + + pushdown_aggregation([plain_dag, aggregated_dag]) + + plain_node = plain_dag.get_node_by_id(bert_id) + self.assertIs(plain_node.operation, Bert) + self.assertNotIn("_pushdown_aggregation", plain_node.parameters) + + pushed_node = aggregated_dag.get_node_by_id(aggregation_id) + self.assertIs(pushed_node.operation, Bert) + self.assertEqual(pushed_node.inputs, [leaf_id]) + self.assertEqual( + pushed_node.parameters["_pushdown_aggregation"], aggregation_params + ) + self.assertIsNone(aggregated_dag.get_node_by_id(bert_id)) + + def test_clip_text_supports_aggregation_pushdown(self): + builder = CSEAwareDAGBuilder() + leaf_id = builder.create_leaf_node("clip_pushdown") + clip_id = builder.create_operation_node( + CLIPText, [leaf_id], CLIPText().get_current_parameters() + ) + aggregation = AggregatedRepresentation("max", target_dimensions=1) + aggregation_id = builder.create_operation_node( + AggregatedRepresentation, + [clip_id], + aggregation.get_current_parameters(), + ) + dag = builder.build(aggregation_id) + + pushdown_aggregation([dag]) + + pushed_node = dag.get_node_by_id(aggregation_id) + self.assertIs(pushed_node.operation, CLIPText) + pushed_aggregation = AggregatedRepresentation( + params=pushed_node.parameters["_pushdown_aggregation"] + ) + self.assertEqual(pushed_aggregation.aggregation_function, "max") + self.assertIsNone(dag.get_node_by_id(clip_id)) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/main/python/tests/scuro/test_unimodal_optimizer.py b/src/main/python/tests/scuro/test_unimodal_optimizer.py index 11c3aa29ea6..41717b4e3a8 100644 --- a/src/main/python/tests/scuro/test_unimodal_optimizer.py +++ b/src/main/python/tests/scuro/test_unimodal_optimizer.py @@ -28,11 +28,24 @@ from systemds.scuro.drsearch.unimodal_optimizer import UnimodalOptimizer from systemds.scuro.representations.covarep_audio_features import ZeroCrossing +from systemds.scuro.representations.covarep_audio_features import ( + Spectral, + RMSE, + Pitch, +) from systemds.scuro.representations.resnet import ResNet from systemds.scuro.representations.mel_spectrogram import MelSpectrogram +from systemds.scuro.representations.mfcc import MFCC +from systemds.scuro.representations.mlp_averaging import MLPAveraging +from systemds.scuro.representations.spectrogram import Spectrogram from systemds.scuro.representations.tfidf import TfIdf from systemds.scuro.representations.bow import BoW from systemds.scuro.representations.bert import Bert +from systemds.scuro.representations.word2vec import W2V +from tests.scuro.test_unimodal_representations import ( + PHYSIOLOGICAL_REPRESENTATIONS, + TIMESERIES_REPRESENTATIONS, +) from systemds.scuro.modality.unimodal_modality import UnimodalModality from tests.scuro.data_generator import ( ModalityRandomDataGenerator, @@ -62,6 +75,37 @@ ModalityType.EMBEDDING: [], } +#: Every registered representation that runs without downloading a pretrained +#: model. The transformer- and CNN-based ones (Bert, RoBERTa, CLIP, GloVe, X3D, +#: VGG19, Swin, Wav2Vec) are deliberately absent: they pull hundreds of MB over +#: the network, which is the same reason the video representation test in +#: test_unimodal_representations.py is commented out. +FULL_TEXT_REPRESENTATIONS = [BoW, TfIdf, W2V] +FULL_AUDIO_REPRESENTATIONS = [ + MFCC, + MelSpectrogram, + Spectrogram, + Spectral, + RMSE, + Pitch, + ZeroCrossing, +] +FULL_IMAGE_REPRESENTATIONS = [ColorHistogram] +FULL_TIMESERIES_REPRESENTATIONS = TIMESERIES_REPRESENTATIONS +FULL_PHYSIOLOGICAL_REPRESENTATIONS = PHYSIOLOGICAL_REPRESENTATIONS + + +def registry_for(modality_type, representations): + """A registry holding `representations` for one modality and nothing else. + + Every modality type has to be present: the optimizer looks its modality up + directly, and a partial dict would raise a KeyError rather than search an + empty space. + """ + registry = {m_type: [] for m_type in ModalityType} + registry[modality_type] = representations + return registry + class TestUnimodalRepresentationOptimizer(unittest.TestCase): data_generator = None @@ -89,6 +133,38 @@ def test_unimodal_optimizer_for_text_modality(self): ) self.optimize_unimodal_representation_for_modality([text]) + def test_bow_and_tfidf_require_dimensionality_reduction_before_task(self): + text_data, text_md = ModalityRandomDataGenerator().create_text_data( + self.num_instances, 10 + ) + text = UnimodalModality( + TestDataLoader( + self.indices, None, ModalityType.TEXT, text_data, str, text_md + ) + ) + + dimensionality_reduction_operators = {ModalityType.EMBEDDING: [MLPAveraging]} + for representation in (BoW, TfIdf): + with self.subTest(representation=representation.__name__), patch.object( + Registry, + "_representations", + registry_for(ModalityType.TEXT, [representation]), + ), patch.object( + Registry, + "_dimensionality_reduction_operators", + dimensionality_reduction_operators, + ): + optimizer = UnimodalOptimizer( + [text], self.tasks, False, enable_checkpointing=False + ) + _, _, task_dags = optimizer._build_execution_dags_for_modality(text) + + self.assertGreater(len(task_dags), 0) + for dag in task_dags: + task_node = dag.get_node_by_id(dag.root_node_id) + task_input = dag.get_node_by_id(task_node.inputs[0]) + self.assertIs(task_input.operation, MLPAveraging) + def test_unimodal_optimizer_for_image_modality(self): image_data, image_md = ModalityRandomDataGenerator().create_visual_modality( self.num_instances, 1, 10, 10 @@ -133,7 +209,7 @@ def test_unimodal_optimizer_for_audio_modality(self): def test_unimodal_optimizer_for_video_modality(self): video_data, video_md = ModalityRandomDataGenerator().create_visual_modality( - self.num_instances, 10, 10 + self.num_instances, 10, 10, 10 ) video = UnimodalModality( TestDataLoader( @@ -142,6 +218,122 @@ def test_unimodal_optimizer_for_video_modality(self): ) self.optimize_unimodal_representation_for_modality([video]) + # ------------------------------------------------------------------ + # Every registered representation, run through the optimizer + # ------------------------------------------------------------------ + # + # The tests above check that the optimizer runs at all. These check that no + # individual representation breaks it: a rep whose get_output_stats, + # preconditions or transform disagree with what the executor expects takes + # the whole search down, and with a two-representation registry that would + # never surface. + + def _optimize_with_registry(self, modality, registry): + with patch.object(Registry, "_representations", registry): + Registry() + unimodal_optimizer = UnimodalOptimizer( + [modality], + self.tasks, + False, + k=1, + max_num_workers=1, + enable_checkpointing=False, + ) + unimodal_optimizer.optimize() + + self.assertIn( + modality.modality_id, + unimodal_optimizer.operator_performance.modality_ids, + ) + result, _ = unimodal_optimizer.operator_performance.get_k_best_results( + modality, self.tasks[0], "accuracy" + ) + self.assertEqual(len(result), 1) + return unimodal_optimizer + + def test_unimodal_optimizer_with_all_text_representations(self): + text_data, text_md = ModalityRandomDataGenerator().create_text_data( + self.num_instances, 10 + ) + text = UnimodalModality( + TestDataLoader( + self.indices, None, ModalityType.TEXT, text_data, str, text_md + ) + ) + self._optimize_with_registry( + text, registry_for(ModalityType.TEXT, FULL_TEXT_REPRESENTATIONS) + ) + + def test_unimodal_optimizer_with_all_audio_representations(self): + audio_data, audio_md = ModalityRandomDataGenerator().create_audio_data( + self.num_instances, 4000 + ) + audio = UnimodalModality( + TestDataLoader( + self.indices, None, ModalityType.AUDIO, audio_data, np.float32, audio_md + ) + ) + self._optimize_with_registry( + audio, registry_for(ModalityType.AUDIO, FULL_AUDIO_REPRESENTATIONS) + ) + + def test_unimodal_optimizer_with_all_image_representations(self): + image_data, image_md = ModalityRandomDataGenerator().create_visual_modality( + self.num_instances, 1, 10, 10 + ) + image = UnimodalModality( + TestDataLoader( + self.indices, None, ModalityType.IMAGE, image_data, np.float32, image_md + ) + ) + self._optimize_with_registry( + image, registry_for(ModalityType.IMAGE, FULL_IMAGE_REPRESENTATIONS) + ) + + def test_unimodal_optimizer_with_all_timeseries_representations(self): + ts_data, ts_md = ModalityRandomDataGenerator().create_timeseries_data( + self.num_instances, 256 + ) + timeseries = UnimodalModality( + TestDataLoader( + self.indices, + None, + ModalityType.TIMESERIES, + ts_data, + np.float32, + ts_md, + ) + ) + optimizer = self._optimize_with_registry( + timeseries, + registry_for(ModalityType.TIMESERIES, FULL_TIMESERIES_REPRESENTATIONS), + ) + # A search over windowed timeseries always proposes some configurations + # the input cannot express (a lag longer than the window, a moment on a + # two-sample window). Those must be pruned up front, not executed. + self.assertGreater(len(optimizer.pruned), 0) + + def test_unimodal_optimizer_with_all_physiological_representations(self): + data, md = ModalityRandomDataGenerator().create_physiological_data( + self.num_instances, 2000, kind="ecg", fs=500.0 + ) + physiological = UnimodalModality( + TestDataLoader( + self.indices, + None, + ModalityType.PHYSIOLOGICAL, + data, + np.float32, + md, + ) + ) + self._optimize_with_registry( + physiological, + registry_for( + ModalityType.PHYSIOLOGICAL, FULL_PHYSIOLOGICAL_REPRESENTATIONS + ), + ) + def test_aggregation_pushdown_preserves_dag_id_and_bert_node_parameters(self): builder = CSEAwareDAGBuilder() modality_id = "test_modality_agg_pushdown" @@ -184,12 +376,13 @@ def test_aggregation_pushdown_preserves_dag_id_and_bert_node_parameters(self): pushdown_aggregation([dag]) self.assertEqual(dag.dag_id, expected_dag_id) - self.assertEqual(dag.root_node_id, bert_id) + self.assertEqual(dag.root_node_id, agg_id) self.assertEqual(len(dag.nodes), 2) - self.assertIsNone(dag.get_node_by_id(agg_id)) + self.assertIsNone(dag.get_node_by_id(bert_id)) - bert_after = dag.get_node_by_id(bert_id) + bert_after = dag.get_node_by_id(agg_id) self.assertIsNotNone(bert_after) + self.assertIs(bert_after.operation, Bert) self.assertEqual(bert_after.inputs, [leaf_id]) self.assertIn("_pushdown_aggregation", bert_after.parameters) self.assertEqual( diff --git a/src/main/python/tests/scuro/test_unimodal_representations.py b/src/main/python/tests/scuro/test_unimodal_representations.py index 59bef40ef64..27e09d48711 100644 --- a/src/main/python/tests/scuro/test_unimodal_representations.py +++ b/src/main/python/tests/scuro/test_unimodal_representations.py @@ -42,10 +42,15 @@ ModalityRandomDataGenerator, ) from systemds.scuro.modality.type import ModalityType +from systemds.scuro.modality.transformed import TransformedModality +from systemds.scuro.representations.representation import RepresentationStats +from systemds.scuro.representations.tabular_features import TabularFeatures +from systemds.scuro.representations.word2vec import W2V from systemds.scuro.representations.timeseries_representations import ( Mean, Max, Min, + Sum, Kurtosis, Skew, Std, @@ -56,6 +61,80 @@ Quantile, ZeroCrossingRate, BandpowerFFT, + LastValue, + TransitionCount, + ObservationDensity, +) +from systemds.scuro.representations.physiological_representations import ( + SDNN, + RMSSD, + pNN, + RRPerMinute, + HRVBandPower, + HRVVLF, + HRVLF, + HRVHF, + HRVLFHF, + PoincareSD1, + PoincareSD2, + SCLSlope, + SCLDynamicRange, + SCRPeaksPerMinute, + SCRAverageAmplitude, + SCRAverageDuration, + BreathingRate, + BreathIntervalRMSSD, + BreathAmplitude, +) + +TIMESERIES_REPRESENTATIONS = [ + Mean, + Min, + Max, + Sum, + Std, + Skew, + Quantile, + Kurtosis, + RMS, + ZeroCrossingRate, + LastValue, + TransitionCount, + ObservationDensity, + ACF, + FrequencyMagnitude, + SpectralCentroid, + BandpowerFFT, +] + + +ECG_REPRESENTATIONS = [ + SDNN, + RMSSD, + pNN, + RRPerMinute, + HRVBandPower, + HRVVLF, + HRVLF, + HRVHF, + HRVLFHF, + PoincareSD1, + PoincareSD2, +] +EDA_REPRESENTATIONS = [ + SCLSlope, + SCLDynamicRange, + SCRPeaksPerMinute, + SCRAverageAmplitude, + SCRAverageDuration, +] +RESPIRATION_REPRESENTATIONS = [ + BreathingRate, + BreathIntervalRMSSD, + BreathAmplitude, +] +PHYSIOLOGICAL_REPRESENTATIONS = ( + ECG_REPRESENTATIONS + EDA_REPRESENTATIONS + RESPIRATION_REPRESENTATIONS ) @@ -147,21 +226,7 @@ def test_audio_representations(self): assert (audio.data[i] == original_data[i]).all() def test_timeseries_representations(self): - ts_representations = [ - Mean(), - Max(), - Min(), - Kurtosis(), - Skew(), - Std(), - RMS(), - ACF(), - FrequencyMagnitude(), - SpectralCentroid(), - Quantile(), - ZeroCrossingRate(), - BandpowerFFT(), - ] + ts_representations = [cls() for cls in TIMESERIES_REPRESENTATIONS] ts_data, ts_md = ModalityRandomDataGenerator().create_timeseries_data( self.num_instances, 100 ) @@ -182,6 +247,417 @@ def test_timeseries_representations(self): for i in range(self.num_instances): assert (ts.data[i] == original_data[i]).all() + def _create_timeseries_modality(self, data, metadata): + modality = UnimodalModality( + TestDataLoader( + np.array(range(len(data))), + None, + ModalityType.TIMESERIES, + data, + np.float32, + metadata, + ) + ) + modality.extract_raw_data() + return modality + + def test_timeseries_output_stats_match_transformed_data(self): + sequence_length = 64 + ts_data, ts_md = ModalityRandomDataGenerator().create_timeseries_data( + self.num_instances, sequence_length + ) + ts = self._create_timeseries_modality(ts_data, ts_md) + input_stats = RepresentationStats(self.num_instances, (sequence_length,)) + + for representation_class in TIMESERIES_REPRESENTATIONS: + with self.subTest(representation=representation_class.__name__): + representation = representation_class() + transformed = ts.apply_representation(representation) + stats = representation.get_output_stats(input_stats) + + self.assertEqual(stats.num_instances, self.num_instances) + self.assertEqual( + np.array(transformed.data).shape[0], self.num_instances + ) + self.assertEqual( + np.array(transformed.data).shape[1:], tuple(stats.output_shape) + ) + + def test_timeseries_batched_path_matches_per_instance_path(self): + sequence_length = 48 + ts_data, _ = ModalityRandomDataGenerator().create_timeseries_data( + self.num_instances, sequence_length + ) + batch = np.stack(ts_data) + + for representation_class in TIMESERIES_REPRESENTATIONS: + with self.subTest(representation=representation_class.__name__): + representation = representation_class() + batched = np.asarray(representation.compute_features_batched(batch)) + if batched.ndim == 1: + batched = batched[:, None] + per_instance = np.stack( + [ + np.atleast_1d(representation.compute_feature(instance)) + for instance in ts_data + ] + ) + np.testing.assert_allclose(batched, per_instance, rtol=1e-5, atol=1e-5) + + def test_timeseries_representations_on_variable_length_instances(self): + lengths = [40 + 7 * i for i in range(self.num_instances)] + ts_data = [np.random.rand(length).astype(np.float32) for length in lengths] + ts_md = [ + ModalityType.TIMESERIES.create_metadata(["signal"], instance) + for instance in ts_data + ] + ts = self._create_timeseries_modality(ts_data, ts_md) + + for representation_class in TIMESERIES_REPRESENTATIONS: + with self.subTest(representation=representation_class.__name__): + transformed = ts.apply_representation(representation_class()) + self.assertEqual( + np.array(transformed.data).shape[0], self.num_instances + ) + self.assertTrue(np.isfinite(np.array(transformed.data)).all()) + + spectrum = ts.apply_representation(FrequencyMagnitude()) + self.assertEqual(np.array(spectrum.data).shape[1], max(lengths) // 2 + 1) + + def test_timeseries_representations_on_degenerate_signals(self): + for name, instance in [ + ("constant", np.full(32, 3.5, dtype=np.float32)), + ("zeros", np.zeros(32, dtype=np.float32)), + ]: + ts_data = [instance.copy() for _ in range(self.num_instances)] + ts_md = [ + ModalityType.TIMESERIES.create_metadata(["signal"], d) for d in ts_data + ] + ts = self._create_timeseries_modality(ts_data, ts_md) + for representation_class in TIMESERIES_REPRESENTATIONS: + if representation_class in (Skew, Kurtosis): + continue + with self.subTest( + signal=name, representation=representation_class.__name__ + ): + transformed = ts.apply_representation(representation_class()) + self.assertTrue(np.isfinite(transformed.data).all()) + + def test_timeseries_minimum_input_length_preconditions(self): + for representation_class in TIMESERIES_REPRESENTATIONS: + representation = representation_class() + minimum = representation.min_input_length + with self.subTest(representation=representation_class.__name__): + self.assertIsNone( + representation.check_preconditions( + RepresentationStats(self.num_instances, (minimum,)) + ) + ) + if minimum > 1: + rejection = representation.check_preconditions( + RepresentationStats(self.num_instances, (minimum - 1,)) + ) + self.assertIsNotNone(rejection) + self.assertIn(str(minimum), rejection) + + def test_acf_rejects_and_narrows_lags_the_input_cannot_express(self): + acf = ACF(k=10) + self.assertIsNotNone( + acf.check_preconditions(RepresentationStats(self.num_instances, (5,))) + ) + self.assertIsNone( + acf.check_preconditions(RepresentationStats(self.num_instances, (50,))) + ) + + candidates = [1, 2, 5, 10, 20] + self.assertEqual( + acf.filter_parameter_domain( + "k", candidates, RepresentationStats(self.num_instances, (5,)) + ), + [1, 2], + ) + # Never empty: a node with no candidates left would stop being tunable. + self.assertEqual( + acf.filter_parameter_domain( + "k", [5, 10], RepresentationStats(self.num_instances, (2,)) + ), + [1], + ) + # Unrelated parameters pass through untouched. + self.assertEqual( + acf.filter_parameter_domain( + "unrelated", candidates, RepresentationStats(self.num_instances, (5,)) + ), + candidates, + ) + + def test_spectral_operators_bind_sampling_rate_to_the_input(self): + for representation in [SpectralCentroid(), BandpowerFFT()]: + with self.subTest(representation=representation.name): + representation.configure_for_input( + RepresentationStats(self.num_instances, (10,), sampling_rate=250.0) + ) + self.assertEqual(representation.fs, 250.0) + + # An input that does not know its rate must not reset the bound one. + representation.configure_for_input( + RepresentationStats(self.num_instances, (10,), sampling_rate=None) + ) + self.assertEqual(representation.fs, 250.0) + self.assertEqual(representation.get_current_parameters()["fs"], 250.0) + + def test_bandpower_band_is_clamped_to_nyquist(self): + self.assertEqual(BandpowerFFT(band_low=0.5, band_width=1.0).band_high, 1.0) + self.assertEqual(BandpowerFFT(band_low=0.0, band_width=0.25).band_high, 0.25) + + def test_quantile_returns_one_column_per_requested_quantile(self): + quantiles = [0.25, 0.5, 0.75] + sequence_length = 64 + ts_data, ts_md = ModalityRandomDataGenerator().create_timeseries_data( + self.num_instances, sequence_length + ) + ts = self._create_timeseries_modality(ts_data, ts_md) + + quantile = Quantile(quantile=quantiles) + transformed = ts.apply_representation(quantile) + stats = quantile.get_output_stats( + RepresentationStats(self.num_instances, (sequence_length,)) + ) + self.assertEqual(tuple(stats.output_shape), (len(quantiles),)) + self.assertEqual( + np.array(transformed.data).shape, (self.num_instances, len(quantiles)) + ) + # np.quantile prepends its own axis; the columns must still come back in + # the requested order rather than transposed. + np.testing.assert_allclose( + transformed.data[0], + np.quantile(ts_data[0], quantiles), + rtol=1e-5, + atol=1e-5, + ) + + def _create_physiological_modality(self, data, metadata): + modality = UnimodalModality( + TestDataLoader( + np.array(range(len(data))), + None, + ModalityType.PHYSIOLOGICAL, + data, + np.float32, + metadata, + ) + ) + modality.extract_raw_data() + return modality + + def test_physiological_representations_output_shapes(self): + data, md = ModalityRandomDataGenerator().create_physiological_data( + self.num_instances, 2000, kind="ecg", fs=500.0 + ) + physiological = self._create_physiological_modality(data, md) + + for representation_class in PHYSIOLOGICAL_REPRESENTATIONS: + with self.subTest(representation=representation_class.__name__): + transformed = physiological.apply_representation(representation_class()) + transformed_data = np.asarray(transformed.data) + self.assertEqual(transformed_data.shape, (self.num_instances, 1)) + self.assertTrue(np.isfinite(transformed_data).all()) + + def test_ecg_features_recover_the_generated_heart_rate(self): + """The generator lays down R peaks at 0.7-0.9 s intervals, i.e. 66-86 + bpm. Anything outside that means the detector is locking onto the noise + floor rather than the beats.""" + fs = 500.0 + data, md = ModalityRandomDataGenerator().create_physiological_data( + self.num_instances, int(fs * 20), kind="ecg", fs=fs + ) + physiological = self._create_physiological_modality(data, md) + transformed_data = np.asarray(physiological.data) + heart_rate = physiological.apply_representation(RRPerMinute(fs=fs)) + self.assertTrue( + ( + (np.array(heart_rate.data) > 60.0) & (np.array(heart_rate.data) < 95.0) + ).all() + ) + + # Jittered intervals mean real, non-zero variability. + for representation_class in [SDNN, RMSSD, PoincareSD1, PoincareSD2]: + with self.subTest(representation=representation_class.__name__): + transformed = physiological.apply_representation( + representation_class(fs=fs) + ) + self.assertTrue((np.array(transformed.data) > 0.0).all()) + + def test_eda_features_recover_the_generated_scr_peaks(self): + """One SCR bump every 10 s is 6 per minute.""" + fs = 4.0 + data, md = ModalityRandomDataGenerator().create_physiological_data( + self.num_instances, 400, kind="eda", fs=fs + ) + physiological = self._create_physiological_modality(data, md) + + peaks_per_minute = physiological.apply_representation(SCRPeaksPerMinute(fs=fs)) + np.testing.assert_allclose(peaks_per_minute.data, 6.0) + + for representation_class in [SCRAverageAmplitude, SCRAverageDuration]: + with self.subTest(representation=representation_class.__name__): + transformed = physiological.apply_representation( + representation_class(fs=fs) + ) + self.assertTrue((np.array(transformed.data) > 0.0).all()) + + # A rising tonic level has a positive slope and a non-degenerate range. + self.assertTrue( + (np.array(physiological.apply_representation(SCLSlope()).data) > 0.0).all() + ) + self.assertTrue( + ( + np.array(physiological.apply_representation(SCLDynamicRange()).data) + > 0.0 + ).all() + ) + + def test_respiration_features_recover_the_generated_breathing_rate(self): + """0.25 Hz is 15 breaths per minute.""" + fs = 500.0 + data, md = ModalityRandomDataGenerator().create_physiological_data( + self.num_instances, int(fs * 20), kind="resp", fs=fs + ) + physiological = self._create_physiological_modality(data, md) + + breathing_rate = physiological.apply_representation(BreathingRate(fs=fs)) + np.testing.assert_allclose(breathing_rate.data, 15.0, rtol=0.1) + + amplitude = physiological.apply_representation(BreathAmplitude(fs=fs)) + np.testing.assert_allclose(amplitude.data, 2.0, rtol=0.2) + + def test_scl_slope_recovers_a_known_linear_trend(self): + ramp = (3.0 * np.arange(100) + 2.0).astype(np.float32) + np.testing.assert_allclose(SCLSlope().compute_feature(ramp), 3.0, rtol=1e-4) + np.testing.assert_allclose( + SCLDynamicRange().compute_feature(np.array([1.0, 5.0, -2.0])), 7.0 + ) + + def test_physiological_representations_on_degenerate_signals(self): + """A flat signal has no beats, no SCRs and no breaths. Every detector + has to fall back to 0.0 instead of dividing by an empty interval list.""" + for name, instance in [ + ("constant", np.ones(400, dtype=np.float32)), + ("zeros", np.zeros(400, dtype=np.float32)), + ("single_sample", np.array([0.5], dtype=np.float32)), + ]: + data = [instance.copy() for _ in range(self.num_instances)] + md = [ + ModalityType.PHYSIOLOGICAL.create_metadata(["signal"], d) for d in data + ] + physiological = self._create_physiological_modality(data, md) + for representation_class in PHYSIOLOGICAL_REPRESENTATIONS: + with self.subTest( + signal=name, representation=representation_class.__name__ + ): + transformed = physiological.apply_representation( + representation_class() + ) + self.assertTrue(np.isfinite(transformed.data).all()) + np.testing.assert_allclose(transformed.data, 0.0, atol=1e-6) + + def test_tabular_features(self): + data_generator = ModalityRandomDataGenerator() + data_generator.modality_type = ModalityType.EMBEDDING + rows = [[1.0, 2.0, 3.0] for _ in range(self.num_instances)] + + modality = TransformedModality(data_generator, "test_transformation") + modality.data = rows + modality.metadata = [ + ModalityType.EMBEDDING.create_metadata(np.asarray(row)) for row in rows + ] + + tabular_features = TabularFeatures() + transformed = tabular_features.transform(modality) + self.assertEqual(transformed.data.shape, (self.num_instances, 3)) + self.assertEqual(transformed.data.dtype, np.float32) + np.testing.assert_allclose(transformed.data, np.asarray(rows)) + + def test_word2vec_representation(self): + vector_size = 20 + text_data, text_md = ModalityRandomDataGenerator().create_text_data( + self.num_instances, 3 + ) + text = UnimodalModality( + TestDataLoader( + self.indices, None, ModalityType.TEXT, text_data, str, text_md + ) + ) + transformed = text.apply_representation(W2V(vector_size=vector_size)) + transformed_data = np.asarray(transformed.data) + self.assertEqual(transformed_data.shape, (self.num_instances, vector_size)) + self.assertTrue(np.isfinite(transformed_data).all()) + + def test_audio_representations_on_a_signal_shorter_than_one_frame(self): + """librosa's default frame is 2048 samples. A shorter instance must + still come back as a single frame rather than an empty array a + downstream aggregation would then reduce over nothing.""" + audio = self._create_audio_modality(signal_length=8) + + for representation in [ + MFCC(), + MelSpectrogram(), + Spectrogram(), + Spectral(), + ZeroCrossing(), + RMSE(), + Pitch(), + ]: + with self.subTest(representation=representation.name): + transformed = representation.transform(audio) + self.assertEqual(len(transformed.data), self.num_instances) + for instance in transformed.data: + self.assertEqual(instance.shape[0], 1) + self.assertTrue(np.isfinite(instance).all()) + + def test_color_histogram_color_spaces_and_normalization(self): + image_data, image_md = ModalityRandomDataGenerator().create_visual_modality( + self.num_instances, 1, height=8, width=8 + ) + image = UnimodalModality( + TestDataLoader( + self.indices, None, ModalityType.IMAGE, image_data, np.float32, image_md + ) + ) + image.extract_raw_data() + + for color_space in ["RGB", "HSV", "GRAY"]: + with self.subTest(color_space=color_space): + representation = ColorHistogram( + color_space=color_space, bins=4, normalize=True + ) + transformed = np.asarray(representation.transform(image).data) + self.assertEqual( + transformed.shape, + (self.num_instances, representation.calculate_hist_dim()), + ) + np.testing.assert_allclose(transformed.sum(axis=1), 1.0, rtol=1e-5) + + # A single-colour image puts every pixel in one bin -- the degenerate + # case a downstream model can learn nothing from. + uniform = [ + np.full((8, 8, 3), 7, dtype=np.uint8) for _ in range(self.num_instances) + ] + uniform_md = [ + ModalityType.IMAGE.create_metadata(8, 8, 3) + for _ in range(self.num_instances) + ] + uniform_image = UnimodalModality( + TestDataLoader( + self.indices, None, ModalityType.IMAGE, uniform, np.float32, uniform_md + ) + ) + uniform_image.extract_raw_data() + histogram = np.asarray( + ColorHistogram(bins=4, normalize=True).transform(uniform_image).data + ) + np.testing.assert_array_equal((histogram > 0).sum(axis=1), 1) + def test_image_representations(self): image_data, image_md = ModalityRandomDataGenerator().create_visual_modality( self.num_instances, 1, height=8, width=8 @@ -248,6 +724,10 @@ def test_chunked_video_representations(self): assert len(r.metadata) == self.num_instances -# TODO: add unit tests for the other representations +# TODO: the representations still untested here are the ones that download a +# pretrained model at construction time -- Bert, RoBERTa, CLIPText/CLIPVisual, +# GloVe, W2V's larger variants, Wav2Vec, VGG19, X3D and SwinVideoTransformer. +# They need either a cached-model fixture or a network-marked test suite, which +# is also why test_video_representations above is commented out. if __name__ == "__main__": unittest.main() diff --git a/src/main/python/tests/scuro/test_window_operations.py b/src/main/python/tests/scuro/test_window_operations.py index a8c86374801..c6a258fb465 100644 --- a/src/main/python/tests/scuro/test_window_operations.py +++ b/src/main/python/tests/scuro/test_window_operations.py @@ -29,13 +29,45 @@ from tests.scuro.data_generator import ModalityRandomDataGenerator, TestDataLoader from systemds.scuro.modality.type import ModalityType from systemds.scuro.modality.unimodal_modality import UnimodalModality +from systemds.scuro.representations.aggregate import Aggregation +from systemds.scuro.representations.timeseries_representations import ( + FrequencyMagnitude, + Mean, + Quantile, + Std, +) +from systemds.scuro.representations.physiological_window import ( + AdaptiveWindow, + PhysiologicalEventWindow, +) from systemds.scuro.representations.window_aggregation import ( StaticWindow, DynamicWindow, WindowAggregation, + resolve_aggregation_function, ) +class _FakeModality: + """The smallest surface a context operator's execute() actually touches. + + Lets a test hand a window operator hand-built instances -- an empty one, a + flat one -- without going through a loader that would reject them first. + """ + + def __init__(self, data, metadata=None): + self.data = data + self.metadata = metadata or [ + ModalityType.TIMESERIES.create_metadata(["signal"], np.asarray(instance)) + for instance in data + ] + + def get_data_layout(self): + from systemds.scuro.modality.type import DataLayout + + return DataLayout.SINGLE_LEVEL + + class TestWindowOperations(unittest.TestCase): @classmethod def setUpClass(cls): @@ -141,6 +173,372 @@ def test_window_aggregation_on_2d_modality(self): windowed_modality = embedding_modality.context(window_operator) + def _timeseries_modality(self, signal_length=100): + return self.data_generator.create1DModality( + self.num_instances, signal_length, ModalityType.TIMESERIES + ) + + # ------------------------------------------------------------------ + # WindowAggregation: window size against signal length + # ------------------------------------------------------------------ + + def test_window_size_of_one_leaves_the_signal_unchanged(self): + signal_length = 60 + modality = self._timeseries_modality(signal_length) + windowed = np.asarray( + WindowAggregation("mean", window_size=1).execute(modality) + ) + self.assertEqual(windowed.shape, (self.num_instances, signal_length)) + np.testing.assert_allclose(windowed, modality.data, rtol=1e-5, atol=1e-5) + + def test_window_larger_than_the_signal_collapses_to_one_padded_window(self): + """The single short window is zero-padded up to window_size, so a mean + over it divides by the nominal size and not by the sample count. That + dilution is the behaviour a caller has to be able to rely on.""" + signal_length = 100 + window_size = 250 + modality = self._timeseries_modality(signal_length) + + windowed = np.asarray( + WindowAggregation("mean", window_size=window_size).execute(modality) + ) + self.assertEqual(windowed.shape, (self.num_instances, 1)) + np.testing.assert_allclose( + windowed[:, 0], + modality.data.sum(axis=1) / window_size, + rtol=1e-5, + atol=1e-5, + ) + + def test_window_size_that_does_not_divide_the_signal_keeps_a_tail_window(self): + signal_length = 100 + window_size = 7 + modality = self._timeseries_modality(signal_length) + + windowed = np.asarray( + WindowAggregation("mean", window_size=window_size).execute(modality) + ) + self.assertEqual( + windowed.shape, + (self.num_instances, math.ceil(signal_length / window_size)), + ) + # The tail window covers only the samples that are actually there. + tail_start = (windowed.shape[1] - 1) * window_size + np.testing.assert_allclose( + windowed[:, -1], + modality.data[:, tail_start:].mean(axis=1), + rtol=1e-5, + atol=1e-5, + ) + + def test_batched_and_per_instance_paths_agree(self): + """Equal-length numeric instances take a vectorized path; anything else + falls back to the per-instance loop. A window's value must not depend on + which of the two ran.""" + modality = self._timeseries_modality(300) + for window_size in [1, 7, 10, 300, 400]: + with self.subTest(window_size=window_size): + batched = np.asarray( + WindowAggregation("mean", window_size=window_size).execute(modality) + ) + per_instance_operator = WindowAggregation( + "mean", window_size=window_size + ) + per_instance = np.stack( + [ + per_instance_operator.window_aggregate_single_level( + np.asarray(instance), + math.ceil(len(instance) / window_size), + ) + for instance in modality.data + ] + ) + np.testing.assert_allclose(batched, per_instance, rtol=1e-5, atol=1e-5) + + def test_window_aggregation_without_padding_returns_one_array_per_instance(self): + modality = self._timeseries_modality(100) + windowed = WindowAggregation("mean", window_size=10, pad=False).execute( + modality + ) + self.assertIsInstance(windowed, list) + self.assertEqual(len(windowed), self.num_instances) + for instance in windowed: + self.assertEqual(np.asarray(instance).shape, (10,)) + + # ------------------------------------------------------------------ + # Preconditions + # ------------------------------------------------------------------ + + def test_empty_instance_is_rejected(self): + """An empty instance has no window to reduce. Failing loudly beats + returning an empty feature the model would only choke on later.""" + empty = _FakeModality( + [np.array([], dtype=np.float32)], + [ + ModalityType.TIMESERIES.create_metadata( + ["signal"], np.zeros(1, dtype=np.float32) + ) + ], + ) + with self.assertRaises(ValueError): + WindowAggregation("mean", window_size=10).execute(empty) + + def test_invalid_aggregation_function_is_rejected(self): + for operator in [WindowAggregation, StaticWindow, DynamicWindow]: + with self.subTest(operator=operator.__name__): + with self.assertRaises(ValueError): + operator(aggregation_function=object()) + # Aggregation itself only knows a fixed set of names. + with self.assertRaises(ValueError): + Aggregation("not_an_aggregation") + + # ------------------------------------------------------------------ + # StaticWindow / DynamicWindow: window count against signal length + # ------------------------------------------------------------------ + + def test_single_window_reduces_the_whole_signal(self): + modality = self._timeseries_modality(100) + for operator_class in [StaticWindow, DynamicWindow]: + with self.subTest(operator=operator_class.__name__): + windowed = np.asarray( + operator_class("mean", num_windows=1).execute(modality) + ) + self.assertEqual(windowed.shape, (self.num_instances, 1)) + np.testing.assert_allclose( + windowed[:, 0], modality.data.mean(axis=1), rtol=1e-5, atol=1e-5 + ) + + def test_static_window_pads_beyond_the_signal_length(self): + """StaticWindow honours num_windows literally: asking for more windows + than there are samples zero-pads rather than clamping.""" + signal_length = 100 + num_windows = 250 + modality = self._timeseries_modality(signal_length) + + operator = StaticWindow("mean", num_windows=num_windows) + windowed = np.asarray(operator.execute(modality)) + self.assertEqual(windowed.shape, (self.num_instances, num_windows)) + self.assertEqual( + tuple( + operator.get_output_stats( + RepresentationStats(self.num_instances, (signal_length,)) + ).output_shape + ), + (num_windows,), + ) + + def test_dynamic_window_clamps_num_windows_to_the_signal_length(self): + """DynamicWindow splits the signal into geometrically growing windows, + which cannot be shorter than one sample -- so the count is capped.""" + signal_length = 100 + modality = self._timeseries_modality(signal_length) + + operator = DynamicWindow("mean", num_windows=250) + windowed = np.asarray(operator.execute(modality)) + self.assertEqual(windowed.shape, (self.num_instances, signal_length)) + self.assertEqual( + tuple( + operator.get_output_stats( + RepresentationStats(self.num_instances, (signal_length,)) + ).output_shape + ), + (signal_length,), + ) + + def test_window_operators_on_variable_length_instances(self): + """A window *count* is length-independent, so ragged instances still + stack into one rectangular block.""" + num_windows = 5 + instances = [ + np.random.rand(100 + 37 * i).astype(np.float32) + for i in range(self.num_instances) + ] + modality = _FakeModality(instances) + + for operator_class in [StaticWindow, DynamicWindow]: + with self.subTest(operator=operator_class.__name__): + windowed = np.asarray( + operator_class("mean", num_windows=num_windows).execute(modality) + ) + self.assertEqual(windowed.shape, (self.num_instances, num_windows)) + self.assertTrue(np.isfinite(windowed).all()) + + # ------------------------------------------------------------------ + # Aggregation functions + # ------------------------------------------------------------------ + + def test_window_operators_accept_a_representation_as_aggregation(self): + """A window may reduce with a full representation, not just a named + aggregation -- and then the per-window feature can be multi-valued, so + get_output_stats has to carry that extra shape.""" + signal_length = 100 + modality = self._timeseries_modality(signal_length) + input_stats = RepresentationStats(self.num_instances, (signal_length,)) + + window_size = 10 + for aggregation, expected_feature_shape in [ + (Mean(), ()), + (Std(), ()), + (Quantile(), ()), + (FrequencyMagnitude(), (window_size // 2 + 1,)), + ]: + with self.subTest(aggregation=aggregation.name): + operator = WindowAggregation(aggregation, window_size=window_size) + windowed = np.asarray(operator.execute(modality)) + stats = operator.get_output_stats(input_stats) + + expected = ( + signal_length // window_size, + *expected_feature_shape, + ) + self.assertEqual(windowed.shape, (self.num_instances, *expected)) + self.assertEqual(tuple(stats.output_shape), expected) + + for operator_class in [StaticWindow, DynamicWindow]: + with self.subTest(operator=operator_class.__name__): + operator = operator_class(Std(), num_windows=5) + windowed = np.asarray(operator.execute(modality)) + self.assertEqual(windowed.shape, (self.num_instances, 5)) + + def test_window_operator_current_parameters_expose_the_nested_aggregation(self): + """The tuner reads its search space from get_current_parameters, so a + representation used as an aggregation has to surface its own parameters + under a prefixed name rather than disappearing behind the class.""" + window = WindowAggregation("mean", window_size=16) + parameters = window.get_current_parameters() + self.assertEqual(parameters["window_size"], 16) + self.assertIs(parameters["aggregation_function"], Aggregation) + self.assertEqual( + parameters["aggregation_function_aggregation_function"], "mean" + ) + + nested = WindowAggregation(Quantile(quantile=0.5), window_size=16) + nested_parameters = nested.get_current_parameters() + self.assertIs(nested_parameters["aggregation_function"], Quantile) + self.assertEqual(nested_parameters["aggregation_function_quantile"], 0.5) + + static = StaticWindow("max", num_windows=7) + self.assertEqual(static.get_current_parameters()["num_windows"], 7) + + def test_resolve_aggregation_function(self): + self.assertEqual(resolve_aggregation_function("mean", None), "mean") + self.assertEqual( + resolve_aggregation_function("mean", {"aggregation_function": "max"}), "max" + ) + # A class is instantiated ... + self.assertIsInstance( + resolve_aggregation_function("mean", {"aggregation_function": Mean}), Mean + ) + # ... and its prefixed parameters are threaded into the instance. + resolved = resolve_aggregation_function( + "mean", + { + "aggregation_function": Quantile, + "aggregation_function_quantile": 0.25, + }, + ) + self.assertIsInstance(resolved, Quantile) + self.assertEqual(resolved.quantile, 0.25) + + # ------------------------------------------------------------------ + # Data-dependent windows + # ------------------------------------------------------------------ + + def test_adaptive_window_output_shape_is_only_an_estimate(self): + """The window count depends on the signal's local variance, so it is + not knowable from statistics. The operator must say so rather than + report a shape the executor would then assert against.""" + signal_length = 300 + modality = self._timeseries_modality(signal_length) + operator = AdaptiveWindow( + "mean", base_window_size=64, overlap=0.5, min_window_size=16 + ) + + stats = operator.get_output_stats( + RepresentationStats(self.num_instances, (signal_length,)) + ) + self.assertFalse(stats.output_shape_is_known) + self.assertEqual(stats.num_instances, self.num_instances) + + windowed = np.asarray(operator.execute(modality)) + self.assertEqual(windowed.shape[0], self.num_instances) + self.assertGreater(windowed.shape[1], 0) + self.assertTrue(np.isfinite(windowed).all()) + + def test_adaptive_window_clamps_a_floor_above_the_nominal_size(self): + """A minimum larger than the nominal window is contradictory; it + collapses onto the nominal size instead of silently inverting.""" + operator = AdaptiveWindow( + "mean", base_window_size=8, overlap=0.5, min_window_size=64 + ) + self.assertEqual(operator.base_window_size, 8) + self.assertEqual(operator.min_window_size, 8) + + def test_adaptive_window_always_advances(self): + """With a small window and a high overlap the stride truncates to zero + samples, which would never move the cursor. The floor of one sample is + what keeps execute() from spinning forever.""" + modality = _FakeModality( + [np.ones(200, dtype=np.float32) for _ in range(self.num_instances)] + ) + for base_window_size, overlap in [(8, 1.0), (1, 0.9), (4, 0.99)]: + with self.subTest(base_window_size=base_window_size, overlap=overlap): + windowed = np.asarray( + AdaptiveWindow( + "mean", + base_window_size=base_window_size, + overlap=overlap, + min_window_size=1, + ).execute(modality) + ) + self.assertEqual(windowed.shape[0], self.num_instances) + self.assertGreater(windowed.shape[1], 0) + + def test_physiological_event_window_splits_on_detected_events(self): + signal_length = 300 + modality = self._timeseries_modality(signal_length) + operator = PhysiologicalEventWindow( + "mean", event_threshold=0.5, min_distance=32 + ) + + stats = operator.get_output_stats( + RepresentationStats(self.num_instances, (signal_length,)) + ) + self.assertFalse(stats.output_shape_is_known) + + windowed = np.asarray(operator.execute(modality)) + self.assertEqual(windowed.shape[0], self.num_instances) + self.assertGreater(windowed.shape[1], 0) + self.assertTrue(np.isfinite(windowed).all()) + + def test_physiological_event_window_falls_back_when_no_event_is_found(self): + """A flat signal has no peaks to split on, so the operator falls back to + an even split instead of producing zero windows.""" + min_distance = 32 + for name, instance in [ + ("constant", np.ones(200, dtype=np.float32)), + ("zeros", np.zeros(200, dtype=np.float32)), + ]: + with self.subTest(signal=name): + modality = _FakeModality( + [instance.copy() for _ in range(self.num_instances)] + ) + windowed = np.asarray( + PhysiologicalEventWindow( + "mean", event_threshold=0.5, min_distance=min_distance + ).execute(modality) + ) + self.assertEqual( + windowed.shape, + (self.num_instances, len(instance) // min_distance), + ) + self.assertTrue(np.isfinite(windowed).all()) + + def test_physiological_event_window_floors_min_distance(self): + self.assertEqual( + PhysiologicalEventWindow("mean", min_distance=0).min_distance, 1 + ) + def verify_window_operation( self, aggregation, modality, windowed_modality, window_size ): diff --git a/src/main/python/tests/scuro/test_window_representation_batching.py b/src/main/python/tests/scuro/test_window_representation_batching.py new file mode 100644 index 00000000000..e561830d081 --- /dev/null +++ b/src/main/python/tests/scuro/test_window_representation_batching.py @@ -0,0 +1,141 @@ +# ------------------------------------------------------------- +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# +# ------------------------------------------------------------- +import math +import unittest + +import numpy as np + +from systemds.scuro.modality.type import DataLayout, ModalityType +from systemds.scuro.representations.mel_spectrogram import MelSpectrogram +from systemds.scuro.representations.mfcc import MFCC +from systemds.scuro.representations.representation import RepresentationStats +from systemds.scuro.representations.unimodal import UnimodalRepresentation +from systemds.scuro.representations.window_aggregation import ( + DynamicWindow, + StaticWindow, + WindowAggregation, +) + + +class _FakeModality: + def __init__(self, data): + self.data = np.asarray(data, dtype=np.float32) + self.metadata = [ + ModalityType.TIMESERIES.create_metadata(["signal"], instance) + for instance in self.data + ] + + def get_data_layout(self): + return DataLayout.SINGLE_LEVEL + + +class _BatchedMean(UnimodalRepresentation): + def __init__(self): + super().__init__("BatchedMean", ModalityType.EMBEDDING) + self.scalar_calls = 0 + self.batch_calls = 0 + + def compute_feature(self, signal): + self.scalar_calls += 1 + return np.asarray(signal).mean() + + def compute_features_batched(self, data): + self.batch_calls += 1 + data = np.asarray(data) + return data.mean(axis=tuple(range(1, data.ndim))) + + def transform(self, modality, aggregation=None): + raise NotImplementedError + + def get_output_stats(self, input_stats): + return RepresentationStats(input_stats.num_instances, (1,)) + + +class TestWindowRepresentationBatching(unittest.TestCase): + def setUp(self): + rng = np.random.default_rng(7) + self.data = rng.normal(size=(4, 100)).astype(np.float32) + self.modality = _FakeModality(self.data) + + def test_static_window_batches_across_instances_and_windows(self): + representation = _BatchedMean() + result = StaticWindow(representation, num_windows=20).execute(self.modality) + expected = self.data.reshape(4, 20, 5).mean(axis=2) + + np.testing.assert_allclose(result, expected, rtol=1e-6, atol=1e-6) + self.assertEqual(representation.scalar_calls, 0) + self.assertGreater(representation.batch_calls, 0) + self.assertLess(representation.batch_calls, self.data.shape[0] * 20) + + def test_dynamic_window_batches_equal_shape_windows(self): + representation = _BatchedMean() + operator = DynamicWindow(representation, num_windows=16) + result = operator.execute(self.modality) + + expected = [] + for instance in self.data: + ends = np.cumsum(operator._window_sizes(len(instance))) + starts = np.concatenate(([0], ends[:-1])) + expected.append( + [instance[start:end].mean() for start, end in zip(starts, ends)] + ) + + np.testing.assert_allclose(result, np.asarray(expected), rtol=1e-6, atol=1e-6) + self.assertEqual(representation.scalar_calls, 0) + self.assertGreater(representation.batch_calls, 0) + self.assertLess(representation.batch_calls, self.data.shape[0] * 16) + + def test_window_aggregation_batches_full_windows_and_preserves_tail(self): + representation = _BatchedMean() + operator = WindowAggregation(representation, window_size=7) + result = operator.execute(self.modality) + expected = np.stack( + [ + [ + instance[start : min(start + 7, len(instance))].mean() + for start in range(0, len(instance), 7) + ] + for instance in self.data + ] + ) + + self.assertEqual(result.shape[1], math.ceil(self.data.shape[1] / 7)) + np.testing.assert_allclose(result, expected, rtol=1e-6, atol=1e-6) + self.assertEqual(representation.scalar_calls, 0) + + def test_audio_representations_preserve_results_when_batched(self): + rng = np.random.default_rng(11) + windows = rng.normal(size=(3, 64)).astype(np.float32) + representations = ( + MFCC(n_mfcc=4, n_mels=8, hop_length=8, n_fft=32), + MelSpectrogram(n_mels=8, hop_length=8, n_fft=32), + ) + for representation in representations: + with self.subTest(representation=representation.name): + batched = representation.compute_features_batched(windows) + per_window = np.stack( + [representation.compute_feature(window) for window in windows] + ) + np.testing.assert_allclose(batched, per_window, rtol=1e-5, atol=1e-5) + + +if __name__ == "__main__": + unittest.main()