From d28e52b1e924007844211a215b30706b380f9699 Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Mon, 12 Oct 2026 15:57:19 +0200 Subject: [PATCH 01/26] decode visual modality more efficiently --- .../systemds/scuro/dataloader/base_loader.py | 27 ++++++ .../systemds/scuro/dataloader/image_loader.py | 6 +- .../systemds/scuro/dataloader/video_loader.py | 24 ++++- .../scuro/drsearch/modality_shared_memory.py | 30 ++++++ .../systemds/scuro/drsearch/node_executor.py | 94 ++++++++++++++++--- .../scuro/representations/pHash_alignment.py | 6 +- 6 files changed, 171 insertions(+), 16 deletions(-) diff --git a/src/main/python/systemds/scuro/dataloader/base_loader.py b/src/main/python/systemds/scuro/dataloader/base_loader.py index 9deaa666826..b58b95b4294 100644 --- a/src/main/python/systemds/scuro/dataloader/base_loader.py +++ b/src/main/python/systemds/scuro/dataloader/base_loader.py @@ -24,6 +24,7 @@ import pickle from abc import ABC, abstractmethod from collections.abc import Sequence +from concurrent.futures import ThreadPoolExecutor from typing import Callable, Iterator, List, Optional, Tuple, Union import math from numbers import Integral @@ -31,6 +32,16 @@ import numpy as np +def decode_threads() -> int: + explicit = os.environ.get("SCURO_DECODE_THREADS") + if explicit: + try: + return max(1, int(explicit)) + except ValueError: + pass + return max(1, min(16, os.cpu_count() or 1)) + + class LazyFileSequence(Sequence): """List-like file references decoded only when a sample is requested.""" @@ -49,6 +60,13 @@ def __len__(self): def subset(self, indices): return type(self)([self.file_names[i] for i in indices], self.decoder) + def materialize(self, num_threads: Optional[int] = None) -> list: + num_threads = decode_threads() if num_threads is None else int(num_threads) + if num_threads <= 1 or len(self) <= 1: + return [self.decoder(file_name) for file_name in self.file_names] + with ThreadPoolExecutor(max_workers=min(num_threads, len(self))) as pool: + return list(pool.map(self.decoder, self.file_names)) + class BaseLoader(ABC): def __init__( @@ -254,6 +272,13 @@ def get_file_names(self, indices=None): def extract(self, file: str, index: Optional[Union[str, List[str]]] = None): pass + def estimate_decoded_bytes(self, metadata=None) -> Optional[int]: + """ + Bytes of all instances once decoded (of the instances `metadata` + describes, if given), or None if not known before decoding. + """ + return None + @staticmethod def file_sanity_check(file): """ @@ -271,6 +296,7 @@ def file_sanity_check(file): def resolve_data_type(data_type): if isinstance(data_type, str): if data_type.lower() in [ + "uint8", "float16", "float32", "float64", @@ -282,6 +308,7 @@ def resolve_data_type(data_type): else: raise ValueError(f"Unsupported data_type string: {data_type}") elif data_type in [ + np.uint8, np.float16, np.float32, np.float64, diff --git a/src/main/python/systemds/scuro/dataloader/image_loader.py b/src/main/python/systemds/scuro/dataloader/image_loader.py index 35394448dfb..b94fa52bd34 100644 --- a/src/main/python/systemds/scuro/dataloader/image_loader.py +++ b/src/main/python/systemds/scuro/dataloader/image_loader.py @@ -46,7 +46,7 @@ def __init__( self, source_path: str, indices: List[str], - data_type: Union[np.dtype, str] = np.float16, + data_type: Union[np.dtype, str] = np.uint8, chunk_size: Optional[int] = None, load=True, ext=".jpg", @@ -139,6 +139,10 @@ def get_stats(self, source_path: str): average_channels, ) + def estimate_decoded_bytes(self, metadata=None) -> int: + metadata = self._all_metadata if metadata is None else metadata + return int(sum(int(md["height"]) * int(md["width"]) * 3 for md in metadata)) + def estimate_peak_memory_bytes(self) -> dict: n = self.chunk_size if self.chunk_size is not None else 1 per_instance = ( diff --git a/src/main/python/systemds/scuro/dataloader/video_loader.py b/src/main/python/systemds/scuro/dataloader/video_loader.py index 1f499a09b3e..028c89a4ba3 100644 --- a/src/main/python/systemds/scuro/dataloader/video_loader.py +++ b/src/main/python/systemds/scuro/dataloader/video_loader.py @@ -69,7 +69,7 @@ def __init__( self, source_path: str, indices: List[str], - data_type: Union[np.dtype, str] = np.float16, + data_type: Union[np.dtype, str] = np.uint8, chunk_size: Optional[int] = None, load=True, fps=None, @@ -158,7 +158,8 @@ def _decode_file(self, file: str): 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) + data_type = np.dtype(self._data_type) + scale_denominator = data_type.type(255.0) if data_type.kind == "f" else None frames = [] idx = 0 while True: @@ -168,7 +169,10 @@ def _decode_file(self, file: str): 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 + if scale_denominator is not None: + frame = frame.astype(data_type) / scale_denominator + elif data_type != frame.dtype: + frame = frame.astype(data_type) frames.append(frame) idx += 1 finally: @@ -259,6 +263,20 @@ def get_stats(self, source_path: str): shape_variance, ) + def estimate_decoded_bytes(self, metadata=None) -> int: + itemsize = np.dtype(self._data_type).itemsize + metadata = self._all_metadata if metadata is None else metadata + return int( + sum( + int(md["length"]) + * int(md["height"]) + * int(md["width"]) + * int(md["num_channels"]) + for md in metadata + ) + * itemsize + ) + def estimate_peak_memory_bytes(self) -> dict: stats = self.stats n = self.chunk_size if self.chunk_size is not None else 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 1f78ff67421..bc2d443a7c0 100644 --- a/src/main/python/systemds/scuro/drsearch/modality_shared_memory.py +++ b/src/main/python/systemds/scuro/drsearch/modality_shared_memory.py @@ -437,6 +437,36 @@ def add_shared_memory_candidate(data: Any, resident_bytes: int = 0) -> bool: return None, None, 0, resident_bytes +def move_arrays_to_shared_memory(arrays: List[np.ndarray]): + if not _is_shared_memory_candidate(arrays): + return None, None, 0 + data_nbytes = _estimate_ndarray_list_nbytes(arrays) + if data_nbytes <= 0 or not _should_offload(data_nbytes): + return None, None, 0 + + dtype = arrays[0].dtype + offsets = [] + total_elems = 0 + for arr in arrays: + offsets.append((total_elems, int(arr.size), tuple(arr.shape))) + total_elems += int(arr.size) + + shm = shared_memory.SharedMemory(create=True, size=data_nbytes) + try: + buffer = np.ndarray((total_elems,), dtype=dtype, buffer=shm.buf) + for i, (start, size, _) in enumerate(offsets): + buffer[start : start + size] = arrays[i].reshape(-1) + arrays[i] = None + del buffer + except BaseException: + shm.close() + shm.unlink() + raise + name = shm.name + shm.close() + return SharedArrayList(name, dtype.str, offsets, total_elems), name, data_nbytes + + _SHARED_MEMORY_WRAPPER_TYPES = ( SharedStringList, SharedGroupedArrayList, diff --git a/src/main/python/systemds/scuro/drsearch/node_executor.py b/src/main/python/systemds/scuro/drsearch/node_executor.py index 4be1b669af0..c797fd0fcc4 100644 --- a/src/main/python/systemds/scuro/drsearch/node_executor.py +++ b/src/main/python/systemds/scuro/drsearch/node_executor.py @@ -26,11 +26,15 @@ import torch +from systemds.scuro.dataloader.base_loader import LazyFileSequence from systemds.scuro.modality.modality import Modality from systemds.scuro.drsearch.modality_result_cache import RefCountResultCache from systemds.scuro.drsearch.modality_shared_memory import ( + SHARED_MEMORY_MIN_BYTES, + _should_offload, add_shared_memory_candidate, collect_shm_names_from_payload, + move_arrays_to_shared_memory, unlink_shm, ) from systemds.scuro.drsearch.node_scheduler import MemoryAwareNodeScheduler @@ -70,6 +74,14 @@ _MAX_NODE_RETRIES = int(os.environ.get("SCURO_MAX_NODE_RETRIES", "3")) +def _materialize_lazy_leaves() -> bool: + return os.environ.get("SCURO_MATERIALIZE_LAZY_LEAVES", "1") == "1" + + +def _lazy_leaf_budget_fraction() -> float: + return float(os.environ.get("SCURO_LAZY_LEAF_BUDGET_FRACTION", "0.5")) + + 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}") @@ -347,11 +359,29 @@ def _run(): } -def _load_leaf_worker(modality: Any) -> Dict[str, Any]: +def _load_leaf_worker(modality: Any, materialize: bool = False) -> Dict[str, Any]: if hasattr(modality, "extract_raw_data") and not modality.has_data(): modality.extract_raw_data() data = modality.data + if materialize and isinstance(data, LazyFileSequence): + arrays = data.materialize() + shared, shm_name, shm_bytes = move_arrays_to_shared_memory(arrays) + if shared is not None: + return { + "data": shared, + "metadata": modality.metadata, + "shm_name": shm_name, + "shm_bytes": shm_bytes, + } + if sum(getattr(a, "nbytes", 0) for a in arrays) < SHARED_MEMORY_MIN_BYTES: + return { + "data": arrays, + "metadata": modality.metadata, + "shm_name": None, + "materialized": True, + } + resident_bytes = 0 try: resident_bytes = modality.estimate_peak_memory_bytes()["cpu_peak_bytes"] @@ -390,8 +420,8 @@ def _dispatch_leaf_batch(payload, gpu_id): def _dispatch_load_leaf(payload, _gpu_id): - (modality,) = payload - return _load_leaf_worker(modality) + modality, materialize = payload + return _load_leaf_worker(modality, materialize) _WORKER_DISPATCH = { @@ -498,6 +528,7 @@ def __init__( self._leaf_shm_names: List[str] = [] self._leaf_shm_owners: Dict[str, Any] = {} self._leaf_original_data: Dict[str, Any] = {} + self._leaf_unshared_originals: List[tuple] = [] self._task_results: Dict[str, ResultEntry] = {} self._owns_pool = worker_pool is None @@ -557,22 +588,50 @@ def _retain_for_submit(self, parent_ids: List[str], payload: Any) -> List[str]: names = list(dict.fromkeys(names)) return self._result_cache.retain_shm_names(names) + def _should_materialize_leaf(self, modality: Any) -> bool: + if not _materialize_lazy_leaves(): + return False + data = modality.data if modality.has_data() else None + if data is not None and not isinstance(data, LazyFileSequence): + return False + estimate = getattr( + getattr(modality, "data_loader", None), "estimate_decoded_bytes", None + ) + if not callable(estimate): + return False + metadata = None + if data is not None and len(modality.metadata or []) == len(data): + metadata = modality.metadata + try: + decoded_bytes = estimate(metadata) + except (KeyError, TypeError, ValueError): + return False + if not decoded_bytes: + return False + budget = self.scheduler.memory_budget["cpu"] * _lazy_leaf_budget_fraction() + if self._leaf_cached_bytes + decoded_bytes > budget: + return False + return decoded_bytes < SHARED_MEMORY_MIN_BYTES or _should_offload(decoded_bytes) + def _load_leaf_modalities(self) -> None: for modality in self._modalities: if self._loads_in_chunks(modality): continue + original = None if getattr(modality, "has_data", None) and modality.has_data(): original = modality.data - wrapped, shm_name, _, _ = add_shared_memory_candidate(original) - if shm_name is not None: - modality._data = wrapped - self._leaf_shm_names.append(shm_name) - self._leaf_shm_owners[shm_name] = modality - self._leaf_original_data[shm_name] = original - continue + if not isinstance(original, LazyFileSequence): + wrapped, shm_name, _, _ = add_shared_memory_candidate(original) + if shm_name is not None: + modality._data = wrapped + self._leaf_shm_names.append(shm_name) + self._leaf_shm_owners[shm_name] = modality + self._leaf_original_data[shm_name] = original + continue + materialize = self._should_materialize_leaf(modality) attempts = 0 while True: - self._pool.submit("load_leaf", (modality,), gpu_id=None) + self._pool.submit("load_leaf", (modality, materialize), gpu_id=None) jr = self._pool.wait() if jr.ok: break @@ -584,10 +643,20 @@ def _load_leaf_modalities(self) -> None: ) modality._data = jr.value["data"] modality.metadata = jr.value["metadata"] + if jr.value.get("materialized"): + self._leaf_unshared_originals.append((modality, original)) shm_name = jr.value.get("shm_name") if shm_name is not None: self._leaf_shm_names.append(shm_name) self._leaf_shm_owners[shm_name] = modality + if original is not None: + self._leaf_original_data[shm_name] = original + shm_bytes = int(jr.value.get("shm_bytes", 0) or 0) + if shm_bytes: + # The decoded leaf now occupies RAM (/dev/shm) for the + # whole run; the scheduler only knew one lazy instance. + self._leaf_cached_bytes += shm_bytes + self._sync_cache_accounting() @staticmethod def _loads_in_chunks(modality: Any) -> bool: @@ -607,6 +676,9 @@ def _cleanup_leaf_shared_memory(self) -> None: unlink_shm(shm_name) self._leaf_shm_names = [] self._leaf_shm_owners = {} + for modality, original in self._leaf_unshared_originals: + modality._data = original + self._leaf_unshared_originals = [] def _submit_node( self, node_id: str, allow_gpu_worker_for_cpu: bool = False diff --git a/src/main/python/systemds/scuro/representations/pHash_alignment.py b/src/main/python/systemds/scuro/representations/pHash_alignment.py index f70a12a770f..6185ee0b92a 100644 --- a/src/main/python/systemds/scuro/representations/pHash_alignment.py +++ b/src/main/python/systemds/scuro/representations/pHash_alignment.py @@ -34,7 +34,11 @@ def compute_descriptor(self, segment): if segment.ndim == 4: # For videos descriptors = [] for s in segment: - frame = (s * 255).astype(np.uint8, copy=True) + frame = ( + (s * 255).astype(np.uint8, copy=True) + if np.issubdtype(s.dtype, np.floating) + else np.asarray(s, dtype=np.uint8) + ) descriptors.append(self.hasher.compute(frame)) return descriptors raise ("PHashAlignment is only implemented for ndim=3 or ndim=4") From ba65e22812e4ac9fc3b209c2d734e3349a161062 Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Mon, 12 Oct 2026 16:16:01 +0200 Subject: [PATCH 02/26] avoid aggregation pushdown if other consumers exist --- .../scuro/drsearch/representation_dag.py | 21 +++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/src/main/python/systemds/scuro/drsearch/representation_dag.py b/src/main/python/systemds/scuro/drsearch/representation_dag.py index efc4692dcdf..3d3a48d553e 100644 --- a/src/main/python/systemds/scuro/drsearch/representation_dag.py +++ b/src/main/python/systemds/scuro/drsearch/representation_dag.py @@ -649,7 +649,21 @@ def get_consumer_count(dags: List[RepresentationDag]) -> Dict[str, int]: return consumer_count +def _consumers_by_node( + dag_group: List[RepresentationDag], +) -> Dict[str, set]: + consumers: Dict[str, set] = defaultdict(set) + for dag in dag_group: + consumers[dag.root_node_id].add(None) + for node in dag.nodes: + for input_id in node.inputs: + consumers[input_id].add(node.node_id) + return consumers + + def pushdown_aggregation(dag_group: List[RepresentationDag]) -> List[RepresentationDag]: + consumers = _consumers_by_node(dag_group) + changed = False for dag in dag_group: agg_nodes = [ n @@ -674,11 +688,18 @@ def pushdown_aggregation(dag_group: List[RepresentationDag]) -> List[Representat ): continue + if consumers[input_id] != {agg_node.node_id}: + continue + 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 + changed = True + + if changed: + for dag in dag_group: dag.nodes = dag.filter_connected_nodes(dag.nodes) return dag_group From 2b908f8349a39806b690ad4eb5da0c98161a7919 Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Mon, 12 Oct 2026 16:30:50 +0200 Subject: [PATCH 03/26] improve openface represenations --- .../scuro/representations/openface.py | 146 +++++++++++------- 1 file changed, 86 insertions(+), 60 deletions(-) diff --git a/src/main/python/systemds/scuro/representations/openface.py b/src/main/python/systemds/scuro/representations/openface.py index b2d82cc6e91..c96759bd318 100644 --- a/src/main/python/systemds/scuro/representations/openface.py +++ b/src/main/python/systemds/scuro/representations/openface.py @@ -518,12 +518,19 @@ def _chunks(self, count): yield start, min(start + self.FRAME_BATCH_SIZE, count) def _detect_faces(self, images_bgr): + """Return the top-scoring face crop and its 15-value detection per image. + + Only the highest-scoring box is used downstream, and NMS always keeps + the highest-scoring box first, so NMS is skipped: an argmax over the + scores selects the same detection and only that prior is decoded. + """ from openface.Pytorch_Retinaface.layers.functions.prior_box import PriorBox from openface.Pytorch_Retinaface.utils.box_utils import decode, decode_landm - from openface.Pytorch_Retinaface.utils.nms.py_cpu_nms import py_cpu_nms detector = self._face_detector + device = detector.device variance = detector.cfg["variance"] + mean = torch.tensor(self.DETECTOR_MEAN_BGR, dtype=torch.float32, device=device) faces = [None] * len(images_bgr) detections = [None] * len(images_bgr) priors_by_shape = {} @@ -536,44 +543,43 @@ def _detect_faces(self, images_bgr): priors_by_shape[(height, width)] = ( PriorBox(detector.cfg, image_size=(height, width)) .forward() - .to(detector.device) + .to(device) .data ) priors = priors_by_shape[(height, width)] - box_scale = torch.Tensor([width, height, width, height]).to(detector.device) - landmark_scale = torch.Tensor([width, height] * 5).to(detector.device) - - batch = torch.cat( - [ - detector.preprocess_image( - self._pad_to(images_bgr[i], height, width) - )[0] - for i in chunk - ] + box_scale = torch.tensor( + [width, height] * 2, dtype=torch.float32, device=device + ) + landmark_scale = torch.tensor( + [width, height] * 5, dtype=torch.float32, device=device ) + + batch = np.stack([self._pad_to(images_bgr[i], height, width) for i in chunk]) + batch = torch.from_numpy(batch).to(device) + batch = (batch.float() - mean).permute(0, 3, 1, 2).contiguous() with torch.no_grad(): locations, confidences, landmark_offsets = detector.model(batch) - for row, index in enumerate(chunk): - boxes = decode(locations.data[row], priors, variance) - boxes = (boxes * box_scale).cpu().numpy() - scores = confidences[row].data.cpu().numpy()[:, 1] - points = decode_landm(landmark_offsets.data[row], priors, variance) - points = (points * landmark_scale).cpu().numpy() - - keep = np.where(scores > detector.confidence_threshold)[0] - boxes, points, scores = boxes[keep], points[keep], scores[keep] - dets = np.hstack((boxes, scores[:, np.newaxis])).astype( - np.float32, copy=False + rows = torch.arange(len(chunk), device=device) + scores, best = confidences[:, :, 1].max(dim=1) + best_priors = priors[best] + boxes = decode(locations[rows, best], best_priors, variance) + points = decode_landm(landmark_offsets[rows, best], best_priors, variance) + best_dets = torch.cat( + (boxes * box_scale, scores[:, None], points * landmark_scale), dim=1 ) - keep = py_cpu_nms(dets, detector.nms_threshold) - dets = np.concatenate((dets[keep], points[keep]), axis=1) + best_dets = best_dets.cpu().numpy() - if len(dets) == 0 or dets[0][4] < detector.vis_threshold: + for row, index in enumerate(chunk): + det = best_dets[row] + if ( + det[4] <= detector.confidence_threshold + or det[4] < detector.vis_threshold + ): continue - bbox = dets[0][:4].astype(int) + bbox = det[:4].astype(int) faces[index] = images_bgr[index][bbox[1] : bbox[3], bbox[0] : bbox[2]] - detections[index] = dets + detections[index] = det return faces, detections def _predict_multitask(self, faces, needs_backbone): @@ -582,9 +588,12 @@ def _predict_multitask(self, faces, needs_backbone): self._ensure_backbone_hook() outputs = ([], [], [], []) for start, stop in self._chunks(len(faces)): - batch = torch.cat( - [predictor.preprocess(face) for face in faces[start:stop]] - ) + batch = torch.stack( + [ + predictor.transform(cv2.cvtColor(face, cv2.COLOR_BGR2RGB)) + for face in faces[start:stop] + ] + ).to(predictor.device) self._backbone_output = None with torch.no_grad(): emotion, gaze, action_units = predictor.model(batch) @@ -603,46 +612,63 @@ def _landmark_features(self, images_bgr, detections): features = np.zeros( (len(images_bgr), len(self.LANDMARK_COLUMNS)), dtype=np.float32 ) - crops = [] - for index, dets in enumerate(detections): - det = dets[0] + max_value = torch.tensor(255.0, device=alignment.device) + indices, matrices = [], [] + for index, det in enumerate(detections): if det[4] < self.vis_threshold: continue x1, y1, x2, y2 = det[:4].astype(int) scale = min(x2 - x1, y2 - y1) / 200 * 1.05 - crops.append( - (index,) - + alignment.preprocess( - images_bgr[index], - float(scale), - float((x2 + x1) / 2), - float((y2 + y1) / 2), + indices.append(index) + matrices.append( + alignment.getCropMatrix.process( + float(scale), float((x2 + x1) / 2), float((y2 + y1) / 2) ) ) - for start, stop in self._chunks(len(crops)): - chunk = crops[start:stop] + for start, stop in self._chunks(len(indices)): + chunk = indices[start:stop] + crops = np.stack( + [ + alignment.transformPerspective.process(images_bgr[index], matrix) + for index, matrix in zip(chunk, matrices[start:stop]) + ] + ) + crops = torch.from_numpy(crops).to(alignment.device) + # A device-side divisor keeps CUDA from rewriting the division as a + # reciprocal multiply, so the inputs match OpenFace's CPU preprocessing. + crops = crops.float().permute(0, 3, 1, 2) / max_value * 2.0 - 1.0 with torch.no_grad(): - output = alignment.alignment(torch.cat([c[1] for c in chunk])) - points_batch = alignment.denorm_points(output[-1]).data.cpu().numpy() - for (index, _, matrix), points in zip(chunk, points_batch): - points = alignment.postprocess(points, np.linalg.inv(matrix)) - points = np.asarray(points, dtype=np.float32) - if points.shape != (self.NUM_LANDMARKS, 2): - raise RuntimeError( - "OpenFace 3.0 returned unexpected landmark dimensions: " - f"{points.shape}" - ) - height, width = images_bgr[index].shape[:2] - points = points.copy() - points[:, 0] /= width - points[:, 1] /= height - features[index] = points.reshape(-1) + output = alignment.alignment(crops) + points = alignment.denorm_points(output[-1]).cpu().numpy() + if points.shape[1:] != (self.NUM_LANDMARKS, 2): + raise RuntimeError( + "OpenFace 3.0 returned unexpected landmark dimensions: " + f"{points.shape[1:]}" + ) + inverse = np.linalg.inv(np.stack(matrices[start:stop])) + x, y = points[..., 0], points[..., 1] + image_points = np.stack( + ( + inverse[:, 0, 0, None] * x + + inverse[:, 0, 1, None] * y + + inverse[:, 0, 2, None], + inverse[:, 1, 0, None] * x + + inverse[:, 1, 1, None] * y + + inverse[:, 1, 2, None], + ), + axis=-1, + ) + image_sizes = np.array( + [images_bgr[index].shape[1::-1] for index in chunk], dtype=np.float32 + ) + image_points /= image_sizes[:, None, :] + features[chunk] = image_points.reshape(len(chunk), -1) return features @classmethod - def _detection_features(cls, detections, image_shape): - detection = np.asarray(detections[0], dtype=np.float32) + def _detection_features(cls, detection, image_shape): + detection = np.asarray(detection, dtype=np.float32) if detection.size < len(cls.DETECTION_COLUMNS): raise RuntimeError( "OpenFace 3.0 returned unexpected face-detection dimensions: " From 5d0b8a1ce189d3463e966a87f701632ebd64c234 Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Tue, 13 Oct 2026 08:06:46 +0200 Subject: [PATCH 04/26] improve visual preprocessing --- .../systemds/scuro/representations/clip.py | 23 ++-- .../systemds/scuro/representations/resnet.py | 14 +- .../representations/swin_video_transformer.py | 14 +- .../systemds/scuro/representations/vgg.py | 14 +- .../scuro/representations/video_batching.py | 3 + .../systemds/scuro/representations/x3d.py | 30 +++-- .../systemds/scuro/utils/torch_dataset.py | 124 +++++++++++++++++- 7 files changed, 183 insertions(+), 39 deletions(-) diff --git a/src/main/python/systemds/scuro/representations/clip.py b/src/main/python/systemds/scuro/representations/clip.py index 55f841e42e8..689aece4a78 100644 --- a/src/main/python/systemds/scuro/representations/clip.py +++ b/src/main/python/systemds/scuro/representations/clip.py @@ -20,7 +20,6 @@ # ------------------------------------------------------------- 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 @@ -47,8 +46,10 @@ from systemds.scuro.utils.static_variables import get_device from systemds.scuro.utils.torch_dataset import ( - CustomDataset, + FrameDataset, + FramePreprocessor, TextDataset, + collate_frames, TextSpanDataset, ) from systemds.scuro.utils.static_variables import ( @@ -279,14 +280,8 @@ def hook(model, input, output): return transformed_modality def create_visual_embeddings(self, modality, aggregation=None): - clip_transform = transforms.Compose( - [ - transforms.ToPILImage(), - transforms.Resize(256), - transforms.CenterCrop(224), - transforms.ToTensor(), - transforms.ConvertImageDtype(dtype=self.data_type), - ] + preprocess = FramePreprocessor.from_hf_processor( + self.processor, dtype=self.data_type ) is_image = modality.modality_type == ModalityType.IMAGE if is_image: @@ -296,11 +291,12 @@ def create_visual_embeddings(self, modality, aggregation=None): lengths = get_sequence_lengths(modality.data, modality.metadata) samples, owner_ids = flatten_owned_sequences(modality.data, lengths) - dataset = CustomDataset(samples, self.data_type, "cpu", tf=clip_transform) + dataset = FrameDataset(samples) dataloader = DataLoader( dataset, batch_size=self.batch_size, shuffle=False, + collate_fn=collate_frames, pin_memory=pin_memory_for(self.device), ) owner_by_chunk = torch.tensor(owner_ids, dtype=torch.long) @@ -309,10 +305,7 @@ def create_visual_embeddings(self, modality, aggregation=None): with transformer_inference_context(self.device): for batch in dataloader: chunk_ids = batch["id"].long() - inputs = self.processor( - images=batch["data"], return_tensors="pt", do_rescale=False - ) - inputs = move_batch_to_device(dict(inputs), self.device) + inputs = {"pixel_values": preprocess(batch["data"], self.device)} if self.layer_name != "": _ = self.model.vision_model(**inputs) output = self.clip_output diff --git a/src/main/python/systemds/scuro/representations/resnet.py b/src/main/python/systemds/scuro/representations/resnet.py index 00b27a22602..382ea757902 100644 --- a/src/main/python/systemds/scuro/representations/resnet.py +++ b/src/main/python/systemds/scuro/representations/resnet.py @@ -26,10 +26,13 @@ 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.utils.torch_dataset import ( + FrameDataset, + FramePreprocessor, + collate_frames, +) from systemds.scuro.modality.transformed import TransformedModality from systemds.scuro.representations.unimodal import UnimodalRepresentation from typing import Tuple, Any @@ -248,11 +251,13 @@ def hook( lengths = get_sequence_lengths(modality.data, modality.metadata) samples, owner_ids = flatten_owned_sequences(modality.data, lengths) - dataset = CustomDataset(samples, self.data_type, "cpu") + dataset = FrameDataset(samples) + preprocess = FramePreprocessor(dtype=self.data_type) dataloader = torch.utils.data.DataLoader( dataset, batch_size=self.batch_size, shuffle=False, + collate_fn=collate_frames, pin_memory=pin_memory_for(self.device), ) owner_by_chunk = torch.tensor(owner_ids, dtype=torch.long) @@ -261,8 +266,7 @@ def hook( 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"]) + _ = self.model(preprocess(batch["data"], self.device)) output = self.activation if output.ndim > 2: output = torch.nn.functional.adaptive_avg_pool2d(output, (1, 1)) 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 2172e6af4ab..11d87d37a4e 100644 --- a/src/main/python/systemds/scuro/representations/swin_video_transformer.py +++ b/src/main/python/systemds/scuro/representations/swin_video_transformer.py @@ -45,7 +45,11 @@ VideoMemoryProfile, ) -from systemds.scuro.utils.torch_dataset import CustomDataset +from systemds.scuro.utils.torch_dataset import ( + FrameDataset, + FramePreprocessor, + collate_frames, +) from systemds.scuro.utils.static_variables import ( get_device, get_device_for_model, @@ -139,13 +143,15 @@ def hook( ) break - dataset = CustomDataset(modality.data, self.data_type, "cpu") + dataset = FrameDataset(modality.data, video=True) + preprocess = FramePreprocessor(dtype=self.data_type) lengths = get_sequence_lengths(modality.data, modality.metadata) accumulator = OwnerAccumulator(len(dataset), len(dataset), aggregation) def forward(model, batch): video_ids = batch["id"].long() - _ = model(batch["data"].permute(0, 2, 1, 3, 4)) + frames = preprocess(batch["data"], next(model.parameters()).device) + _ = model(frames.permute(0, 2, 1, 3, 4)) values = self.swin_output if isinstance(values, tuple): values = values[0] @@ -160,7 +166,7 @@ def forward(model, batch): accumulator.update(pooled.to(self.device), video_ids, video_ids) with transformer_inference_context(self.device): - self._run_batches(dataset, lengths, forward) + self._run_batches(dataset, lengths, forward, collate_fn=collate_frames) transformed_modality = TransformedModality( modality, self, self.output_modality_type diff --git a/src/main/python/systemds/scuro/representations/vgg.py b/src/main/python/systemds/scuro/representations/vgg.py index 214b56a59b0..baee8ba044f 100644 --- a/src/main/python/systemds/scuro/representations/vgg.py +++ b/src/main/python/systemds/scuro/representations/vgg.py @@ -19,7 +19,11 @@ # # ------------------------------------------------------------- from systemds.scuro.utils.converter import numpy_dtype_to_torch_dtype -from systemds.scuro.utils.torch_dataset import CustomDataset +from systemds.scuro.utils.torch_dataset import ( + FrameDataset, + FramePreprocessor, + collate_frames, +) from systemds.scuro.modality.transformed import TransformedModality from systemds.scuro.dataloader.video_loader import VideoStats from systemds.scuro.representations.unimodal import UnimodalRepresentation @@ -43,7 +47,6 @@ flatten_owned_sequences, get_sequence_lengths, inference_context, - move_batch_to_device, pin_memory_for, ) @@ -200,11 +203,13 @@ def hook( lengths = get_sequence_lengths(modality.data, modality.metadata) samples, owner_ids = flatten_owned_sequences(modality.data, lengths) - dataset = CustomDataset(samples, self.data_type, "cpu") + dataset = FrameDataset(samples) + preprocess = FramePreprocessor(dtype=self.data_type) dataloader = torch.utils.data.DataLoader( dataset, batch_size=self.batch_size, shuffle=False, + collate_fn=collate_frames, pin_memory=pin_memory_for(self.device), ) owner_by_chunk = torch.tensor(owner_ids, dtype=torch.long) @@ -213,8 +218,7 @@ def hook( 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"]) + _ = self.model(preprocess(batch["data"], self.device)) output = self.activation if output.ndim > 2: output = torch.nn.functional.adaptive_avg_pool2d(output, (1, 1)) diff --git a/src/main/python/systemds/scuro/representations/video_batching.py b/src/main/python/systemds/scuro/representations/video_batching.py index 2b9802e956c..d277e2854ac 100644 --- a/src/main/python/systemds/scuro/representations/video_batching.py +++ b/src/main/python/systemds/scuro/representations/video_batching.py @@ -83,6 +83,9 @@ def max_batch_frames(self, value): self._max_batch_frames = None if value is None else max(1, int(value)) def _release_cache_on_new_shape(self, frames) -> None: + if not isinstance(frames, torch.Tensor): + # Clips of differing frame sizes stay a list until preprocessing. + return shape = tuple(frames.shape) if frames.is_cuda and shape != getattr(self, "_last_batch_shape", None): if getattr(self, "_last_batch_shape", None) is not None: diff --git a/src/main/python/systemds/scuro/representations/x3d.py b/src/main/python/systemds/scuro/representations/x3d.py index d3f6b50e16b..e863e805c86 100644 --- a/src/main/python/systemds/scuro/representations/x3d.py +++ b/src/main/python/systemds/scuro/representations/x3d.py @@ -50,7 +50,12 @@ get_device, get_device_for_model, ) -from systemds.scuro.utils.torch_dataset import CustomDataset +from systemds.scuro.utils.torch_dataset import ( + FrameDataset, + FramePreprocessor, + collate_frames, + stack_or_list, +) class Identity(torch.nn.Module): @@ -192,10 +197,10 @@ def _get_parameters(self, high_level=True): @staticmethod def _collate_videos(samples): + """Pad raw clips to a common length; the padding is zeroed after preprocessing.""" video_ids = torch.tensor([sample["id"] for sample in samples]) - target_length = max( - X3D._MIN_FRAMES, max(sample["data"].shape[0] for sample in samples) - ) + lengths = torch.tensor([sample["data"].shape[0] for sample in samples]) + target_length = max(X3D._MIN_FRAMES, int(lengths.max())) videos = [] for sample in samples: frames = sample["data"] @@ -206,7 +211,7 @@ def _collate_videos(samples): ) frames = torch.cat((frames, pad), dim=0) videos.append(frames) - return {"id": video_ids, "data": torch.stack(videos)} + return {"id": video_ids, "data": stack_or_list(videos), "lengths": lengths} def transform(self, modality, aggregation=None): self.model = self.model.to(self.device) @@ -229,13 +234,17 @@ def hook( ) break - dataset = CustomDataset(modality.data, self.data_type, "cpu") + dataset = FrameDataset(modality.data, video=True) + preprocess = FramePreprocessor(dtype=self.data_type) lengths = get_sequence_lengths(modality.data, modality.metadata) embeddings = [None] * len(dataset) def forward(model, batch): video_ids = batch["id"].long() - _ = model(batch["data"].permute(0, 2, 1, 3, 4)) + frames = preprocess(batch["data"], next(model.parameters()).device) + for i, length in enumerate(batch["lengths"].tolist()): + frames[i, length:] = 0 + _ = model(frames.permute(0, 2, 1, 3, 4)) values = self.activation if isinstance(values, tuple): values = values[0] @@ -371,12 +380,14 @@ def hook( ) break - dataset = CustomDataset(modality.data, self.data_type, "cpu") + dataset = FrameDataset(modality.data, video=True) + preprocess = FramePreprocessor(dtype=self.data_type) embeddings = [None] * len(dataset) def forward(model, batch): video_ids = batch["id"].long() - _ = model(batch["data"].permute(0, 2, 1, 3, 4)) + frames = preprocess(batch["data"], next(model.parameters()).device) + _ = model(frames.permute(0, 2, 1, 3, 4)) values = self.features if isinstance(values, tuple): values = values[0] @@ -389,6 +400,7 @@ def forward(model, batch): dataset, get_sequence_lengths(modality.data, modality.metadata), forward, + collate_fn=collate_frames, ) if self.output_file is not None: diff --git a/src/main/python/systemds/scuro/utils/torch_dataset.py b/src/main/python/systemds/scuro/utils/torch_dataset.py index fea59088d7b..7dcdea6f504 100644 --- a/src/main/python/systemds/scuro/utils/torch_dataset.py +++ b/src/main/python/systemds/scuro/utils/torch_dataset.py @@ -18,14 +18,136 @@ # under the License. # # ------------------------------------------------------------- -from typing import Dict +from typing import Dict, Sequence import numpy as np import torch import torchvision.transforms as transforms +from torchvision.transforms import InterpolationMode +from torchvision.transforms.v2 import functional as TF from systemds.scuro.modality.type import ModalityType +IMAGENET_MEAN = (0.485, 0.456, 0.406) +IMAGENET_STD = (0.229, 0.224, 0.225) + +# PIL resample codes, as Hugging Face image processors store them. +_PIL_INTERPOLATION = { + 0: InterpolationMode.NEAREST, + 1: InterpolationMode.LANCZOS, + 2: InterpolationMode.BILINEAR, + 3: InterpolationMode.BICUBIC, +} + + +class FrameDataset(torch.utils.data.Dataset): + """Raw frames as channel-last tensors; preprocessing is left to FramePreprocessor. + + Each item is an image (H, W[, C]) or, with ``video=True``, a clip (T, H, W[, C]). + Grayscale items get a channel axis of size 1. + """ + + def __init__(self, data, video=False): + self.data = data + self.frame_ndim = 4 if video else 3 + + def __getitem__(self, index) -> Dict[str, object]: + frames = np.asarray(self.data[index]) + if frames.ndim == self.frame_ndim - 1: + frames = frames[..., None] + return {"id": index, "data": torch.from_numpy(np.ascontiguousarray(frames))} + + def __len__(self) -> int: + return len(self.data) + + +def stack_or_list(tensors: Sequence[torch.Tensor]): + """Stack equally shaped tensors; keep a list when shapes differ.""" + if all(t.shape == tensors[0].shape for t in tensors): + return torch.stack(list(tensors)) + return list(tensors) + + +def collate_frames(samples): + return { + "id": torch.tensor([sample["id"] for sample in samples]), + "data": stack_or_list([sample["data"] for sample in samples]), + } + + +class FramePreprocessor: + """Resize, center crop and normalize batches of raw frames on their target device. + + Matches transforms.Resize(resize) -> CenterCrop(crop) -> ToTensor -> Normalize on + PIL images: frames are resized as uint8 with antialiasing and rounded back to + uint8 before scaling to [0, 1]. Float frames are taken to be in [0, 1]. + """ + + _MAX_CHUNK_BYTES = 256 * 1024 * 1024 + + def __init__( + self, + resize=256, + crop=224, + mean=IMAGENET_MEAN, + std=IMAGENET_STD, + interpolation=InterpolationMode.BILINEAR, + dtype=torch.float32, + ): + self.resize = int(resize) + self.crop = int(crop) + self.mean = list(mean) + self.std = list(std) + self.interpolation = interpolation + self.dtype = dtype + + @classmethod + def from_hf_processor(cls, processor, dtype=torch.float32): + image_processor = getattr(processor, "image_processor", processor) + crop = image_processor.crop_size + return cls( + resize=image_processor.size["shortest_edge"], + crop=crop["height"] if isinstance(crop, dict) else crop, + mean=image_processor.image_mean, + std=image_processor.image_std, + interpolation=_PIL_INTERPOLATION.get( + int(image_processor.resample), InterpolationMode.BICUBIC + ), + dtype=dtype, + ) + + def __call__(self, frames, device) -> torch.Tensor: + """(..., H, W, C) frames, or a list of them -> (..., 3, crop, crop) on device.""" + if isinstance(frames, (list, tuple)): + return torch.stack([self(f, device) for f in frames]) + + device = torch.device(device) + # Antialiased interpolation is not implemented on MPS. + work_device = torch.device("cpu") if device.type == "mps" else device + frames = frames.to(work_device, non_blocking=frames.is_pinned()) + leading = frames.shape[:-3] + frames = frames.reshape(-1, *frames.shape[-3:]) + frame_bytes = 4 * 3 * frames.shape[1] * frames.shape[2] + chunk = max(1, self._MAX_CHUNK_BYTES // frame_bytes) + output = torch.cat([self._process(part) for part in frames.split(chunk)]) + return output.reshape(*leading, *output.shape[1:]).to(device) + + def _process(self, frames: torch.Tensor) -> torch.Tensor: + x = frames.permute(0, 3, 1, 2) + if x.shape[1] == 1: + x = x.expand(-1, 3, -1, -1) + elif x.shape[1] > 3: + x = x[:, :3] + if x.is_floating_point(): + x = x.float().mul(255.0).round_().clamp_(0, 255).to(torch.uint8) + x = TF.resize( + x, [self.resize], interpolation=self.interpolation, antialias=True + ) + x = TF.center_crop(x, [self.crop, self.crop]) + x = x.to(self.dtype).div_(255.0) + # Frames arrive channel-last; hand the model the NCHW layout it got before. + return TF.normalize(x, self.mean, self.std).contiguous() + class CustomDataset(torch.utils.data.Dataset): def __init__(self, data, data_type, device, size=None, tf=None): From d667b30c84645ad9fd61e9d2c046ff48c59a03bb Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Tue, 13 Oct 2026 08:19:36 +0200 Subject: [PATCH 05/26] adapt number of gpu slots for representation nodes --- .../scuro/drsearch/multimodal_ga_optimizer.py | 5 ++--- .../python/systemds/scuro/drsearch/node_executor.py | 5 ++--- .../systemds/scuro/drsearch/node_scheduler.py | 5 ++--- .../python/systemds/scuro/drsearch/worker_pool.py | 13 +++++++++++++ 4 files changed, 19 insertions(+), 9 deletions(-) diff --git a/src/main/python/systemds/scuro/drsearch/multimodal_ga_optimizer.py b/src/main/python/systemds/scuro/drsearch/multimodal_ga_optimizer.py index 0c182338101..66d86023e58 100644 --- a/src/main/python/systemds/scuro/drsearch/multimodal_ga_optimizer.py +++ b/src/main/python/systemds/scuro/drsearch/multimodal_ga_optimizer.py @@ -70,6 +70,7 @@ from systemds.scuro.drsearch.worker_pool import ( PersistentWorkerPool, create_mp_context, + default_gpu_slots_per_device, default_threads_per_worker, ) from systemds.scuro.representations.aggregated_representation import ( @@ -739,9 +740,7 @@ def _start_parallel_runtime(self, task_name: str) -> None: ctx=create_mp_context(), threads_per_worker=self.threads_per_worker, gpu_devices=[info["index"] for info in gpu_memory_info()], - gpu_slots_per_device=int( - os.environ.get("SCURO_GPU_SLOTS_PER_DEVICE", "1") - ), + gpu_slots_per_device=default_gpu_slots_per_device(), gpu_demand_fraction=1.0, ) except Exception: diff --git a/src/main/python/systemds/scuro/drsearch/node_executor.py b/src/main/python/systemds/scuro/drsearch/node_executor.py index c797fd0fcc4..1e8a073ade0 100644 --- a/src/main/python/systemds/scuro/drsearch/node_executor.py +++ b/src/main/python/systemds/scuro/drsearch/node_executor.py @@ -46,6 +46,7 @@ from systemds.scuro.drsearch.worker_pool import ( PersistentWorkerPool, create_mp_context, + default_gpu_slots_per_device, default_threads_per_worker, ) from systemds.scuro.representations.aggregated_representation import ( @@ -540,9 +541,7 @@ def __init__( 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_slots_per_device=default_gpu_slots_per_device(), gpu_demand_fraction=1.0, ) self._pool = worker_pool diff --git a/src/main/python/systemds/scuro/drsearch/node_scheduler.py b/src/main/python/systemds/scuro/drsearch/node_scheduler.py index b69c50f43cc..45f626cb939 100644 --- a/src/main/python/systemds/scuro/drsearch/node_scheduler.py +++ b/src/main/python/systemds/scuro/drsearch/node_scheduler.py @@ -34,6 +34,7 @@ stats_dtype, stats_itemsize, ) +from systemds.scuro.drsearch.worker_pool import default_gpu_slots_per_device from systemds.scuro.utils.memory_utility import gpu_memory_info from systemds.scuro.utils.static_variables import DEBUG @@ -81,9 +82,7 @@ def __init__( 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_per_device = default_gpu_slots_per_device() 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 diff --git a/src/main/python/systemds/scuro/drsearch/worker_pool.py b/src/main/python/systemds/scuro/drsearch/worker_pool.py index 4ab50abfebb..ff68bef3645 100644 --- a/src/main/python/systemds/scuro/drsearch/worker_pool.py +++ b/src/main/python/systemds/scuro/drsearch/worker_pool.py @@ -56,6 +56,19 @@ def default_threads_per_worker() -> int: return _resolve_thread_count(1) +DEFAULT_GPU_SLOTS_PER_DEVICE = 4 + + +def default_gpu_slots_per_device() -> int: + explicit = os.environ.get("SCURO_GPU_SLOTS_PER_DEVICE") + if explicit: + try: + return max(1, int(explicit)) + except ValueError: + pass + return DEFAULT_GPU_SLOTS_PER_DEVICE + + def set_thread_env_before_spawn(num_threads: int) -> None: num_threads = _resolve_thread_count(num_threads) for var in _THREAD_ENV_VARS: From 4ec4d6301a403135a2f8281120e2127830513ecf Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Tue, 13 Oct 2026 08:34:45 +0200 Subject: [PATCH 06/26] avoid model loading --- .../systemds/scuro/representations/bert.py | 35 +++++++-- .../systemds/scuro/representations/clip.py | 62 +++++++++------- .../scuro/representations/openface.py | 8 ++- .../systemds/scuro/representations/resnet.py | 71 ++++++++----------- .../representations/swin_video_transformer.py | 29 ++++++-- .../systemds/scuro/representations/utils.py | 24 +++++++ .../systemds/scuro/representations/vgg.py | 26 +++---- .../systemds/scuro/representations/x3d.py | 42 ++++++----- 8 files changed, 183 insertions(+), 114 deletions(-) diff --git a/src/main/python/systemds/scuro/representations/bert.py b/src/main/python/systemds/scuro/representations/bert.py index d5b851dd2a1..c1e0c8e0e4f 100644 --- a/src/main/python/systemds/scuro/representations/bert.py +++ b/src/main/python/systemds/scuro/representations/bert.py @@ -54,6 +54,10 @@ class BertFamily(UnimodalRepresentation): supports_aggregation_pushdown = True cache_in_worker = True + MODEL_BYTES = None + HIDDEN_SIZE = 768 + NUM_LAYERS = 12 + INTERMEDIATE_SIZE = 3072 def __init__( self, @@ -143,9 +147,25 @@ def estimate_output_memory_bytes(self, input_stats): input_stats.num_instances * np.prod(output_stats) * self.data_type.itemsize ) - def estimate_peak_memory_bytes(self, input_stats): + def _model_spec(self): + if self.MODEL_BYTES is not None: + return ( + self.MODEL_BYTES, + self.HIDDEN_SIZE, + self.NUM_LAYERS, + self.INTERMEDIATE_SIZE, + ) model = AutoModel.from_pretrained(self.model_name) - params_bytes = model.get_memory_footprint() + cfg = model.config + return ( + model.get_memory_footprint(), + cfg.hidden_size, + cfg.num_hidden_layers, + getattr(cfg, "intermediate_size", 4 * cfg.hidden_size), + ) + + def estimate_peak_memory_bytes(self, input_stats): + params_bytes, hidden_size, num_layers, intermediate_size = self._model_spec() tokenizer_bytes = 60 * 1024 * 1024 # rough upper bound output_bytes = self.estimate_output_memory_bytes(input_stats) @@ -178,11 +198,6 @@ def estimate_peak_memory_bytes(self, input_stats): + 128 * 1024 * 1024 ) - cfg = model.config - hidden_size = cfg.hidden_size - num_layers = cfg.num_hidden_layers - intermediate_size = getattr(cfg, "intermediate_size", 4 * hidden_size) - batch_tokens = self.batch_size * self.max_seq_length activations_per_token_per_layer = ( self.C * (hidden_size + intermediate_size) * self.data_type.itemsize @@ -354,6 +369,8 @@ def collate(samples): @register_representation(ModalityType.TEXT) @register_expensive_representation(ModalityType.TEXT) class Bert(BertFamily): + MODEL_BYTES = 437_937_152 + def __init__( self, layer="cls", @@ -399,6 +416,8 @@ def __init__( @register_representation(ModalityType.TEXT) @register_expensive_representation(ModalityType.TEXT) class RoBERTa(BertFamily): + MODEL_BYTES = 498_590_752 + def __init__( self, layer="cls", @@ -504,6 +523,8 @@ def __init__( # @register_representation(ModalityType.TEXT) class ELECTRA(BertFamily): + MODEL_BYTES = 435_574_784 + def __init__( self, layer="cls", diff --git a/src/main/python/systemds/scuro/representations/clip.py b/src/main/python/systemds/scuro/representations/clip.py index 689aece4a78..0e7c2951be8 100644 --- a/src/main/python/systemds/scuro/representations/clip.py +++ b/src/main/python/systemds/scuro/representations/clip.py @@ -26,6 +26,7 @@ from systemds.scuro.representations.representation import RepresentationStats from systemds.scuro.representations.unimodal import UnimodalRepresentation from systemds.scuro.representations.utils import ( + LazyModelMixin, LengthBucketBatchSampler, OwnerAccumulator, OwnedSequenceDataset, @@ -60,21 +61,28 @@ ) from torch.utils.data import DataLoader +CLIP_MODEL_NAME = "openai/clip-vit-base-patch32" +# Memory footprint (fp32 parameters and buffers) of the full CLIPModel. +CLIP_MODEL_BYTES = 605_110_268 + @register_representation([ModalityType.VIDEO, ModalityType.IMAGE]) @register_expensive_representation([ModalityType.VIDEO, ModalityType.IMAGE]) -class CLIPVisual(UnimodalRepresentation): +class CLIPVisual(LazyModelMixin, UnimodalRepresentation): OUTPUT_MODALITY_TYPE = ModalityType.EMBEDDING supports_aggregation_pushdown = True cache_in_worker = True + MODEL_BYTES = CLIP_MODEL_BYTES + HIDDEN_SIZE = 768 + NUM_LAYERS = 12 + PATCH_SIZE = 32 def __init__(self, output_file=None, batch_size=32, layer_name="", params=None): parameters = self._get_parameters() super().__init__("CLIPVisual", self.OUTPUT_MODALITY_TYPE, 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") + self._processor = None if params is not None: self.batch_size = int((params or {}).get("batch_size", batch_size)) self.layer_name = params.get("layer_name", layer_name) @@ -96,6 +104,15 @@ def gpu_id(self, gpu_id): self._gpu_id = gpu_id self.device = get_device(gpu_id) + def _build_model(self): + return CLIPModel.from_pretrained(CLIP_MODEL_NAME) + + @property + def processor(self): + if self._processor is None: + self._processor = CLIPProcessor.from_pretrained(CLIP_MODEL_NAME) + return self._processor + def _get_parameters(self): parameters = { # "batch_size": [1, 2, 4, 8, 16, 32, 64, 128], @@ -171,8 +188,7 @@ def estimate_peak_memory_bytes(self, input_stats) -> dict: ) max_frames = int(getattr(input_stats, "max_length", 1)) - model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32") - model_bytes = int(model.get_memory_footprint()) + model_bytes = self.MODEL_BYTES per_image_payload = EMB_DIM * out_dtype_size per_image_item = per_image_payload + NP_ARRAY_HEADER_BYTES + PY_LIST_SLOT_BYTES @@ -213,10 +229,9 @@ def estimate_peak_memory_bytes(self, input_stats) -> dict: cpu_peak = model_bytes + outputs_retained + cpu_transient - cfg = model.vision_model.config - hidden_size = int(cfg.hidden_size) - num_layers = int(cfg.num_hidden_layers) - patch = int(cfg.patch_size) + hidden_size = self.HIDDEN_SIZE + num_layers = self.NUM_LAYERS + patch = self.PATCH_SIZE seq_len = (224 // patch) ** 2 + 1 gpu_input = batch_size * 3 * 224 * 224 * out_dtype_size @@ -340,6 +355,11 @@ def _pool_visual_output(self, output: torch.Tensor) -> torch.Tensor: class CLIPText(UnimodalRepresentation): supports_aggregation_pushdown = True cache_in_worker = True + MODEL_BYTES = CLIP_MODEL_BYTES + HIDDEN_SIZE = 512 + NUM_LAYERS = 12 + INTERMEDIATE_SIZE = 2048 + NUM_HEADS = 8 def __init__(self, output_file=None, batch_size=32, layer_name="", params=None): if params is not None: @@ -425,13 +445,10 @@ def get_output_stats(self, input_stats) -> RepresentationStats: def estimate_peak_memory_bytes(self, input_stats) -> dict: output_bytes = self.estimate_output_memory_bytes(input_stats) - model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32") - - cfg = model.text_model.config - hidden_size = cfg.hidden_size - num_layers = cfg.num_hidden_layers - intermediate_size = getattr(cfg, "intermediate_size", 4 * hidden_size) - num_heads = cfg.num_attention_heads + hidden_size = self.HIDDEN_SIZE + num_layers = self.NUM_LAYERS + intermediate_size = self.INTERMEDIATE_SIZE + num_heads = self.NUM_HEADS dtype_size = self.data_type.itemsize batch_tokens = self.batch_size * self.max_seq_length hidden_ffn_bytes = ( @@ -464,7 +481,7 @@ def estimate_peak_memory_bytes(self, input_stats) -> dict: input_bytes_all_instances = self.batch_size * per_instance_input_bytes batch_output_bytes = self.batch_size * 512 * np.dtype(np.float32).itemsize cpu_peak = ( - model.get_memory_footprint() + self.MODEL_BYTES + 100 * 1024 * 1024 + output_bytes + batch_peak_bytes @@ -472,10 +489,7 @@ def estimate_peak_memory_bytes(self, input_stats) -> dict: + input_bytes_all_instances ) gpu_peak = ( - model.get_memory_footprint() - + batch_peak_bytes - + activations_bytes - + batch_output_bytes + self.MODEL_BYTES + batch_peak_bytes + activations_bytes + batch_output_bytes ) return {"cpu_peak_bytes": cpu_peak, "gpu_peak_bytes": gpu_peak} @@ -484,11 +498,9 @@ def transform(self, modality, aggregation=None): modality, self, self.output_modality_type ) if self.processor is None: - self.processor = CLIPProcessor.from_pretrained( - "openai/clip-vit-base-patch32" - ) + self.processor = CLIPProcessor.from_pretrained(CLIP_MODEL_NAME) if self.model is None: - self.model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32") + self.model = CLIPModel.from_pretrained(CLIP_MODEL_NAME) self.model = self.model.to(self.device) self.model.eval() self.clip_output = None diff --git a/src/main/python/systemds/scuro/representations/openface.py b/src/main/python/systemds/scuro/representations/openface.py index c96759bd318..bbb1e6240a6 100644 --- a/src/main/python/systemds/scuro/representations/openface.py +++ b/src/main/python/systemds/scuro/representations/openface.py @@ -554,7 +554,9 @@ def _detect_faces(self, images_bgr): [width, height] * 5, dtype=torch.float32, device=device ) - batch = np.stack([self._pad_to(images_bgr[i], height, width) for i in chunk]) + batch = np.stack( + [self._pad_to(images_bgr[i], height, width) for i in chunk] + ) batch = torch.from_numpy(batch).to(device) batch = (batch.float() - mean).permute(0, 3, 1, 2).contiguous() with torch.no_grad(): @@ -564,7 +566,9 @@ def _detect_faces(self, images_bgr): scores, best = confidences[:, :, 1].max(dim=1) best_priors = priors[best] boxes = decode(locations[rows, best], best_priors, variance) - points = decode_landm(landmark_offsets[rows, best], best_priors, variance) + points = decode_landm( + landmark_offsets[rows, best], best_priors, variance + ) best_dets = torch.cat( (boxes * box_scale, scores[:, None], points * landmark_scale), dim=1 ) diff --git a/src/main/python/systemds/scuro/representations/resnet.py b/src/main/python/systemds/scuro/representations/resnet.py index 382ea757902..8a30c9bd342 100644 --- a/src/main/python/systemds/scuro/representations/resnet.py +++ b/src/main/python/systemds/scuro/representations/resnet.py @@ -22,6 +22,7 @@ from systemds.scuro.dataloader.video_loader import VideoStats from systemds.scuro.representations.representation import RepresentationStats from systemds.scuro.representations.utils import ( + LazyModelMixin, OwnerAccumulator, flatten_owned_sequences, get_sequence_lengths, @@ -55,9 +56,16 @@ def forward(self, input_: torch.Tensor) -> torch.Tensor: @register_representation([ModalityType.IMAGE, ModalityType.VIDEO]) @register_expensive_representation([ModalityType.IMAGE, ModalityType.VIDEO]) -class ResNet(UnimodalRepresentation): +class ResNet(LazyModelMixin, UnimodalRepresentation): supports_aggregation_pushdown = True cache_in_worker = True + MODEL_BYTES = { + "ResNet18": 46_796_608, + "ResNet34": 87_259_072, + "ResNet50": 102_441_032, + "ResNet101": 178_618_848, + "ResNet152": 241_378_168, + } def __init__( self, @@ -68,7 +76,6 @@ def __init__( params=None, ): self.data_type = torch.float32 - self.model = None self._activation_hook = None self.activation = None self.gpu_id = None @@ -86,11 +93,6 @@ def __init__( self.params = params self.output_file = output_file - self.model.eval() - for param in self.model.parameters(): - param.requires_grad = False - - self.model.fc = Identity() @property def gpu_id(self): @@ -100,7 +102,7 @@ def gpu_id(self): def gpu_id(self, gpu_id): self._gpu_id = gpu_id self.device = get_device(gpu_id) - if self.model is not None: + if self._model is not None: self.model = self.model.to(self.device) @property @@ -109,32 +111,25 @@ def model_name(self): @model_name.setter def model_name(self, model_name): - self._model_name = model_name - if model_name == "ResNet18": - model = models.resnet18(weights=models.ResNet18_Weights.DEFAULT) - self.model = model.to(self.device) - self.model = self.model.to(self.data_type) - - elif model_name == "ResNet34": - model = models.resnet34(weights=models.ResNet34_Weights.DEFAULT) - self.model = model.to(self.device) - self.model = self.model.to(self.data_type) - elif model_name == "ResNet50": - model = models.resnet50(weights=models.ResNet50_Weights.DEFAULT) - self.model = model.to(self.device) - self.model = self.model.to(self.data_type) - - elif model_name == "ResNet101": - model = models.resnet101(weights=models.ResNet101_Weights.DEFAULT) - self.model = model.to(self.device) - self.model = self.model.to(self.data_type) - - elif model_name == "ResNet152": - model = models.resnet152(weights=models.ResNet152_Weights.DEFAULT) - self.model = model.to(self.device) - self.model = self.model.to(self.data_type) - else: + if model_name not in self.MODEL_BYTES: raise NotImplementedError + self._model_name = model_name + self._model = None + + def _build_model(self): + factory, weights = { + "ResNet18": (models.resnet18, models.ResNet18_Weights), + "ResNet34": (models.resnet34, models.ResNet34_Weights), + "ResNet50": (models.resnet50, models.ResNet50_Weights), + "ResNet101": (models.resnet101, models.ResNet101_Weights), + "ResNet152": (models.resnet152, models.ResNet152_Weights), + }[self.model_name] + model = factory(weights=weights.DEFAULT) + model.eval() + for param in model.parameters(): + param.requires_grad = False + model.fc = Identity() + return model.to(self.device).to(self.data_type) def estimate_output_memory_bytes(self, input_stats: ImageStats) -> int: shape = self.get_output_stats(input_stats).output_shape @@ -177,15 +172,7 @@ def estimate_peak_memory_bytes(self, input_stats: ImageStats) -> dict: safety_margin_bytes = 100 * 1024 * 1024 - param_size = 0 - for param in self.model.parameters(): - param_size += param.nelement() * param.element_size() - - buffer_size = 0 - for buffer in self.model.buffers(): - buffer_size += buffer.nelement() * buffer.element_size() - - size_all_bytes = param_size + buffer_size + size_all_bytes = self.MODEL_BYTES[self.model_name] cpu_peak = ( size_all_bytes * 2 * self.data_type.itemsize 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 11d87d37a4e..5d5abaac2ea 100644 --- a/src/main/python/systemds/scuro/representations/swin_video_transformer.py +++ b/src/main/python/systemds/scuro/representations/swin_video_transformer.py @@ -35,6 +35,7 @@ ) from systemds.scuro.dataloader.video_loader import VideoStats from systemds.scuro.representations.utils import ( + LazyModelMixin, OwnerAccumulator, get_sequence_lengths, transformer_inference_context, @@ -58,9 +59,12 @@ @register_representation([ModalityType.VIDEO]) @register_expensive_representation([ModalityType.VIDEO]) -class SwinVideoTransformer(FrameBudgetedVideoModel, UnimodalRepresentation): +class SwinVideoTransformer( + LazyModelMixin, FrameBudgetedVideoModel, UnimodalRepresentation +): OUTPUT_MODALITY_TYPE = ModalityType.EMBEDDING _EMBED_DIM = 768 + MODEL_BYTES = 127_384_024 cache_in_worker = True _MEMORY_PROFILE = VideoMemoryProfile( clip_frame_bytes=1 * MIB, @@ -93,14 +97,24 @@ def __init__( self.batch_size = batch_size # 32 clips of 300 frames would need ~175 GiB; see FrameBudgetedVideoModel. self.max_batch_frames = max_batch_frames - 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.device = get_device() + self._gpu_id = None + # Until the scheduler places it, the model picks a device when built. + self._placed = False self._activation_hook = None - self.model = self.model.to(self.device) - self.model.eval() - for param in self.model.parameters(): + + def _build_model(self): + model = swin3d_t(weights=models.video.Swin3D_T_Weights.KINETICS400_V1) + if not self._placed: + self.device = get_device_for_model(model, memory_factor=1.5) + self._gpu_id = self.device.index + model.eval() + for param in model.parameters(): param.requires_grad = False + return model.to(self.device) + + def _model_bytes(self) -> int: + return self.MODEL_BYTES @property def gpu_id(self): @@ -109,6 +123,7 @@ def gpu_id(self): @gpu_id.setter def gpu_id(self, gpu_id): self._gpu_id = gpu_id + self._placed = True self.device = get_device(gpu_id) def get_output_stats(self, input_stats) -> RepresentationStats: diff --git a/src/main/python/systemds/scuro/representations/utils.py b/src/main/python/systemds/scuro/representations/utils.py index 4ad52e926f9..5f3a3d56f73 100644 --- a/src/main/python/systemds/scuro/representations/utils.py +++ b/src/main/python/systemds/scuro/representations/utils.py @@ -175,6 +175,30 @@ def __len__(self): return self.offsets[-1] +class LazyModelMixin: + """Loads the pretrained model on first use of ``self.model``. + + The optimizer and scheduler instantiate operators in the main process only + to read their parameter grids, output shapes and memory estimates, so + those must come from class-level constants rather than the model. + """ + + _model = None + + def _build_model(self): + raise NotImplementedError + + @property + def model(self): + if self._model is None: + self._model = self._build_model() + return self._model + + @model.setter + def model(self, model): + self._model = model + + 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] diff --git a/src/main/python/systemds/scuro/representations/vgg.py b/src/main/python/systemds/scuro/representations/vgg.py index baee8ba044f..a88cfc6be9b 100644 --- a/src/main/python/systemds/scuro/representations/vgg.py +++ b/src/main/python/systemds/scuro/representations/vgg.py @@ -43,6 +43,7 @@ from systemds.scuro.dataloader.image_loader import ImageStats from systemds.scuro.representations.representation import RepresentationStats from systemds.scuro.representations.utils import ( + LazyModelMixin, OwnerAccumulator, flatten_owned_sequences, get_sequence_lengths, @@ -58,22 +59,20 @@ def forward(self, input_: torch.Tensor) -> torch.Tensor: @register_representation([ModalityType.IMAGE, ModalityType.VIDEO]) @register_expensive_representation([ModalityType.IMAGE, ModalityType.VIDEO]) -class VGG19(UnimodalRepresentation): +class VGG19(LazyModelMixin, UnimodalRepresentation): OUTPUT_MODALITY_TYPE = ModalityType.EMBEDDING supports_aggregation_pushdown = True cache_in_worker = True + MODEL_BYTES = 574_668_960 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", self.OUTPUT_MODALITY_TYPE, parameters) self.params = params @@ -82,13 +81,15 @@ def __init__( layer = params.get("layer_name", layer) self.output_file = output_file self.layer_name = layer - self.model.eval() self.batch_size = batch_size - for param in self.model.parameters(): + def _build_model(self): + model = models.vgg19(weights=models.VGG19_Weights.DEFAULT) + model.eval() + for param in model.parameters(): param.requires_grad = False - - self.model.fc = Identity() + model.fc = Identity() + return model.to(self.device) @property def gpu_id(self): @@ -98,7 +99,7 @@ def gpu_id(self): def gpu_id(self, gpu_id): self._gpu_id = gpu_id self.device = get_device(gpu_id) - if self.model is not None: + if self._model is not None: self.model = self.model.to(self.device) def _get_parameters(self): @@ -144,12 +145,7 @@ def estimate_peak_memory_bytes(self, input_stats: ImageStats) -> dict: * input_stats.max_channels * self.data_type.itemsize ) - 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() - ) + model_size_bytes = self.MODEL_BYTES return { "cpu_peak_bytes": ( diff --git a/src/main/python/systemds/scuro/representations/x3d.py b/src/main/python/systemds/scuro/representations/x3d.py index e863e805c86..4d9799109c1 100644 --- a/src/main/python/systemds/scuro/representations/x3d.py +++ b/src/main/python/systemds/scuro/representations/x3d.py @@ -37,6 +37,7 @@ from systemds.scuro.representations.representation import RepresentationStats from systemds.scuro.representations.unimodal import UnimodalRepresentation from systemds.scuro.representations.utils import ( + LazyModelMixin, get_sequence_lengths, inference_context, save_embeddings, @@ -65,8 +66,9 @@ def forward(self, input_: torch.Tensor) -> torch.Tensor: @register_representation([ModalityType.VIDEO]) @register_expensive_representation([ModalityType.VIDEO]) -class X3D(FrameBudgetedVideoModel, UnimodalRepresentation): +class X3D(LazyModelMixin, FrameBudgetedVideoModel, UnimodalRepresentation): cache_in_worker = True + MODEL_BYTES = {"r3d": 133_524_448, "s3d": 33_365_672} _MEMORY_PROFILES = { "r3d": VideoMemoryProfile( clip_frame_bytes=42 * MIB, @@ -97,6 +99,9 @@ def __init__( model_name = params.get("model_name", model_name) layer = params.get("layer_name", layer) batch_size = int(params.get("batch_size", batch_size)) + self.device = get_device() + self._gpu_id = None + self._placed = False self.model_name = model_name parameters = self._get_parameters() super().__init__("X3D", ModalityType.EMBEDDING, parameters) @@ -105,14 +110,25 @@ def __init__( self.layer_name = layer self.batch_size = batch_size self.max_batch_frames = max_batch_frames - self._gpu_id = self.device.index self._activation_hook = None self.activation = None - self.model.eval() - for param in self.model.parameters(): + + def _build_model(self): + if self.model_name == "r3d": + model = r3d_18(pretrained=True) + else: + model = s3d(weights=models.video.S3D_Weights.DEFAULT) + if not self._placed: + self.device = get_device_for_model(model, memory_factor=1.5) + self._gpu_id = self.device.index + model.eval() + for param in model.parameters(): param.requires_grad = False + model.fc = Identity() + return model.to(self.device) - self.model.fc = Identity() + def _model_bytes(self) -> int: + return self.MODEL_BYTES[self.model_name] @property def gpu_id(self): @@ -121,8 +137,9 @@ def gpu_id(self): @gpu_id.setter def gpu_id(self, gpu_id): self._gpu_id = gpu_id + self._placed = True self.device = get_device(gpu_id) - if self.model is not None: + if self._model is not None: self.model = self.model.to(self.device) def get_output_stats(self, input_stats) -> RepresentationStats: @@ -147,17 +164,10 @@ def model_name(self): @model_name.setter def model_name(self, model_name): - self._model_name = model_name - if model_name == "r3d": - self.model = r3d_18(pretrained=True) - self.device = get_device_for_model(self.model, memory_factor=1.5) - self.model = self.model.to(self.device) - elif model_name == "s3d": - self.model = s3d(weights=models.video.S3D_Weights.DEFAULT) - self.device = get_device_for_model(self.model, memory_factor=1.5) - self.model = self.model.to(self.device) - else: + if model_name not in self.MODEL_BYTES: raise NotImplementedError + self._model_name = model_name + self._model = None def _get_parameters(self, high_level=True): parameters = { From 74a1e725679e058d999a2f7d9c748d37712d2d7b Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Tue, 13 Oct 2026 09:03:30 +0200 Subject: [PATCH 07/26] keep models over workers --- .../systemds/scuro/drsearch/node_executor.py | 146 ++++++++++++++++-- .../systemds/scuro/drsearch/node_scheduler.py | 57 ++++++- .../systemds/scuro/drsearch/worker_pool.py | 44 +++++- .../scuro/representations/image_bind.py | 12 +- .../systemds/scuro/representations/resnet.py | 3 +- .../representations/swin_video_transformer.py | 5 +- .../systemds/scuro/representations/utils.py | 75 +++++++++ .../systemds/scuro/representations/vgg.py | 3 +- .../systemds/scuro/representations/x3d.py | 7 +- 9 files changed, 318 insertions(+), 34 deletions(-) diff --git a/src/main/python/systemds/scuro/drsearch/node_executor.py b/src/main/python/systemds/scuro/drsearch/node_executor.py index 1e8a073ade0..753683ede67 100644 --- a/src/main/python/systemds/scuro/drsearch/node_executor.py +++ b/src/main/python/systemds/scuro/drsearch/node_executor.py @@ -83,6 +83,10 @@ def _lazy_leaf_budget_fraction() -> float: return float(os.environ.get("SCURO_LAZY_LEAF_BUDGET_FRACTION", "0.5")) +def _keep_worker_models() -> bool: + return os.environ.get("SCURO_KEEP_WORKER_MODELS", "1") == "1" + + 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}") @@ -126,17 +130,21 @@ def _run_gpu_op(fn, gpu_id: Optional[int], cleanup_after: bool = True): _WORKER_OP_CACHE: Dict[str, Any] = {} +def _operation_cache_key(node) -> Optional[str]: + if not getattr(node.operation, "cache_in_worker", False): + return None + try: + params = sorted( + (kv for kv in node.parameters.items() if not kv[0].startswith("_")), + key=lambda kv: kv[0], + ) + return f"{node.operation.__module__}.{node.operation.__qualname__}|{params!r}" + except Exception: + return None + + def _instantiate_operation(node): - cache_key = None - if getattr(node.operation, "cache_in_worker", False): - try: - params_repr = repr(sorted(node.parameters.items(), key=lambda kv: kv[0])) - cache_key = ( - f"{node.operation.__module__}.{node.operation.__qualname__}" - f"|{params_repr}" - ) - except Exception: - cache_key = None + cache_key = _operation_cache_key(node) if cache_key is not None and cache_key in _WORKER_OP_CACHE: return _WORKER_OP_CACHE[cache_key] operation = node.operation(params=node.parameters) @@ -398,16 +406,42 @@ def _load_leaf_worker(modality: Any, materialize: bool = False) -> Dict[str, Any return {"data": data, "metadata": modality.metadata, "shm_name": shm_name} +def _retained_gpu_bytes() -> int: + if not torch.cuda.is_available() or not torch.cuda.is_initialized(): + return 0 + return int(torch.cuda.memory_allocated(0)) + + +def _release_worker_models() -> None: + _WORKER_OP_CACHE.clear() + if torch.cuda.is_available() and torch.cuda.is_initialized(): + cleanup_gpu(0) + + def _dispatch_node(payload, gpu_id): node, input_mods = payload + cache_key = _operation_cache_key(node) + if cache_key is not None and cache_key not in _WORKER_OP_CACHE: + _release_worker_models() + keep = cache_key is not None and gpu_id is not None and _keep_worker_models() try: - return _execute_node_worker(node, input_mods, gpu_id) + value = _execute_node_worker(node, input_mods, gpu_id) + except BaseException: + keep = False + raise 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: + if cache_key is not None and not keep: + _release_worker_models() + elif gpu_id is not None: cleanup_gpu(gpu_id) + value["cached_operation"] = next(iter(_WORKER_OP_CACHE), None) + value["retained_gpu_bytes"] = _retained_gpu_bytes() + return value + + +def _dispatch_release_models(_payload, _gpu_id): + _release_worker_models() + return {"retained_gpu_bytes": _retained_gpu_bytes()} def _dispatch_task(payload, gpu_id): @@ -430,6 +464,7 @@ def _dispatch_load_leaf(payload, _gpu_id): "task": _dispatch_task, "leaf_batch": _dispatch_leaf_batch, "load_leaf": _dispatch_load_leaf, + "release_models": _dispatch_release_models, } @@ -438,6 +473,11 @@ class _NodeUnit: node_id: str +@dataclass +class _ReleaseUnit: + pid: int + + @dataclass class _BatchUnit: node_ids: List[str] @@ -524,7 +564,8 @@ def __init__( self._node_attempts: Dict[str, int] = {} - self._job_units: Dict[int, Union[_NodeUnit, _BatchUnit]] = {} + self._job_units: Dict[int, Union[_NodeUnit, _BatchUnit, _ReleaseUnit]] = {} + self._worker_models: Dict[int, tuple] = {} self._job_retained_shm: Dict[int, List[str]] = {} self._leaf_shm_names: List[str] = [] self._leaf_shm_owners: Dict[str, Any] = {} @@ -712,12 +753,20 @@ def _submit_node( retained = self._retain_for_submit(parent_ids, payload) self.scheduler.begin_execution(node_id) self.scheduler.move_to_running(node_id) + cache_key = _operation_cache_key(node) job_id = self._pool.submit( "node", (node, payload), gpu_id=gpu_id, allow_gpu_worker_for_cpu=allow_gpu_worker_for_cpu, + preferred_pids=[ + pid + for pid, (_, _, key) in self._worker_models.items() + if cache_key is not None and key == cache_key + ], ) + if cache_key is not None: + self._set_worker_model(self._pool.job_pid(job_id), 0, None) self._job_units[job_id] = _NodeUnit(node_id) self._job_retained_shm[job_id] = retained @@ -740,6 +789,52 @@ def _submit_leaf_batch( self._job_units[job_id] = _BatchUnit(node_ids) self._job_retained_shm[job_id] = retained + def _set_worker_model( + self, pid: Optional[int], retained_bytes: int, cache_key: Optional[str] + ) -> None: + if pid is None: + return + if cache_key is None: + self._worker_models.pop(pid, None) + else: + self._worker_models[pid] = ( + self._pool.worker_gpu(pid), + int(retained_bytes), + cache_key, + ) + bytes_by_gpu: Dict[int, int] = {} + for gpu_id, nbytes, _ in self._worker_models.values(): + if gpu_id is not None: + bytes_by_gpu[gpu_id] = bytes_by_gpu.get(gpu_id, 0) + nbytes + set_retained = getattr(self.scheduler, "set_retained_gpu_bytes", None) + if callable(set_retained): + set_retained(bytes_by_gpu) + + def _submit_model_release(self, pid: int) -> int: + job_id = self._pool.submit_to(pid, "release_models", None) + self._job_units[job_id] = _ReleaseUnit(pid) + self._job_retained_shm[job_id] = [] + return job_id + + def _release_blocking_models(self) -> None: + pop_requests = getattr(self.scheduler, "pop_release_requests", None) + if not callable(pop_requests): + return + for gpu_id in pop_requests(): + for pid, (worker_gpu, nbytes, _) in list(self._worker_models.items()): + if worker_gpu == gpu_id and nbytes > 0 and self._pool.is_idle(pid): + self._submit_model_release(pid) + + def _release_all_worker_models(self) -> None: + """Leave a reused pool's GPUs as free as the next executor measures.""" + if self._job_units: + return + for pid in list(self._worker_models): + if self._pool.is_idle(pid): + self._submit_model_release(pid) + while self._job_units: + self._process_result(self._pool.wait()) + def _fill_pipeline(self) -> None: if not self._pool.has_idle_worker: return @@ -772,6 +867,7 @@ def _gpu_id(entry): entry, allow_gpu_worker_for_cpu=allow_gpu_worker_for_cpu, ) + self._release_blocking_models() def _get_cached_parent(self, parent_id: str): try: @@ -807,6 +903,12 @@ def _process_result(self, jr) -> None: try: if unit is None: return + if isinstance(unit, _ReleaseUnit): + # A failed release is not retried; forgetting the model at + # worst under-counts it until the worker's next node. + retained_bytes = jr.value["retained_gpu_bytes"] if jr.ok else 0 + self._set_worker_model(unit.pid, retained_bytes, None) + return if not jr.ok: self._handle_job_failure(unit, jr) return @@ -825,6 +927,11 @@ def _sync_cache_accounting(self) -> None: self.scheduler.set_cached_cpu_bytes(self._leaf_cached_bytes + cache_bytes) def _handle_job_failure(self, unit: Union[_NodeUnit, _BatchUnit], jr): + if jr.worker_died or ( + isinstance(unit, _NodeUnit) + and not self._is_task_node(self.scheduler.mapping[unit.node_id]) + ): + self._set_worker_model(jr.pid, 0, None) reason = jr.error or "unknown worker failure" if jr.cuda_oom: reason += " (CUDA out of memory)" @@ -876,6 +983,11 @@ def _handle_node_success(self, node_id: str, value: Dict[str, Any]) -> None: self._release_parents(node_id) self.scheduler.complete_node(node_id) else: + self._set_worker_model( + value.get("pid"), + value.get("retained_gpu_bytes", 0), + value.get("cached_operation"), + ) self._handle_modality_result( value["result"], node_id, @@ -1156,6 +1268,8 @@ def run(self) -> Dict[str, Any]: finally: if self._owns_pool: self._pool.shutdown() + else: + self._release_all_worker_models() self._result_cache.cleanup_all() self._cleanup_leaf_shared_memory() diff --git a/src/main/python/systemds/scuro/drsearch/node_scheduler.py b/src/main/python/systemds/scuro/drsearch/node_scheduler.py index 45f626cb939..257d62a2218 100644 --- a/src/main/python/systemds/scuro/drsearch/node_scheduler.py +++ b/src/main/python/systemds/scuro/drsearch/node_scheduler.py @@ -21,7 +21,7 @@ from __future__ import annotations import os import re -from typing import List, Dict, Optional, Any +from typing import List, Dict, Optional, Any, Set from collections import defaultdict, deque from collections import deque import torch @@ -115,7 +115,10 @@ def __init__( info["index"]: int(info["total_b"] - info["free_b"]) for info in self.gpu_memory_info }, + # Models idle workers keep between nodes (see node_executor). + "gpu_retained": {info["index"]: 0 for info in self.gpu_memory_info}, } + self.release_requests: Set[int] = set() self._cpu_reserved_nodes: Dict[str, int] = {} self.exclusive_memory_nodes: Dict[str, str] = {} self._initialized = False @@ -147,6 +150,17 @@ def update_cpu_memory_in_use(self, delta_bytes: int): def set_cached_cpu_bytes(self, cached_bytes: int) -> None: self.memory_stats["cpu_cached"] = cached_bytes + def set_retained_gpu_bytes(self, bytes_by_gpu: Dict[int, int]) -> None: + self.memory_stats["gpu_retained"] = { + gpu_id: int(bytes_by_gpu.get(gpu_id, 0)) + for gpu_id in self.memory_budget["gpu"] + } + + def pop_release_requests(self) -> Set[int]: + """GPUs whose retained worker models block a node that would fit.""" + requests, self.release_requests = self.release_requests, set() + return requests + def _cache_draining(self) -> bool: budget = self.memory_budget["cpu"] if _exclusive_over_budget_enabled() and self._total_cpu_in_use() >= budget: @@ -491,6 +505,8 @@ def not_enough_memory(self) -> bool: if gpu_mem > 0.0 and self.n_gpu > 0: gpu_id = self._gpu_with_most_free_memory(gpu_mem) if gpu_id is None: + if self._gpu_freed_by_release(gpu_mem) is not None: + continue if ( self.gpu_wait_attempts.get(node_id, 0) <= _MAX_GPU_SCHEDULE_ATTEMPTS @@ -549,6 +565,10 @@ def _check_memory_constraints( if gpu_id is None: if self._waiting_for_gpu_slot(gpu_mem, pending_gpu, pending_slots): return False, None + release_gpu = self._gpu_freed_by_release(gpu_mem, pending_gpu) + if release_gpu is not None: + self.release_requests.add(release_gpu) + 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: @@ -581,15 +601,39 @@ def _gpu_with_most_free_memory( 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) + free = self._free_gpu_bytes(gpu_id, pending_gpu) if free >= memory_needed: candidates.append((free, -used_slots, -gpu_id, gpu_id)) if not candidates: return None return max(candidates)[-1] + def _free_gpu_bytes( + self, + gpu_id: int, + pending_gpu: Dict[int, int], + include_retained: bool = True, + ) -> int: + 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 include_retained: + free -= self.memory_stats["gpu_retained"].get(gpu_id, 0) + return free + + def _gpu_freed_by_release( + self, memory_needed: int, pending_gpu: Optional[Dict[int, int]] = None + ) -> Optional[int]: + """The GPU that fits ``memory_needed`` once retained models are dropped.""" + pending_gpu = pending_gpu or {} + candidates = [ + (self._free_gpu_bytes(gpu_id, pending_gpu, include_retained=False), gpu_id) + for gpu_id in self.gpu_devices + if self.memory_stats["gpu_retained"].get(gpu_id, 0) > 0 + ] + candidates = [c for c in candidates if c[0] >= memory_needed] + return max(candidates)[1] if candidates else None + def _waiting_for_gpu_slot( self, memory_needed: int, @@ -599,10 +643,7 @@ def _waiting_for_gpu_slot( 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: + if self._free_gpu_bytes(gpu_id, pending_gpu) >= memory_needed: return True return False diff --git a/src/main/python/systemds/scuro/drsearch/worker_pool.py b/src/main/python/systemds/scuro/drsearch/worker_pool.py index ff68bef3645..9529049a238 100644 --- a/src/main/python/systemds/scuro/drsearch/worker_pool.py +++ b/src/main/python/systemds/scuro/drsearch/worker_pool.py @@ -26,7 +26,7 @@ import signal import time from dataclasses import dataclass -from typing import Any, Callable, Dict, List, Optional, Tuple +from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple import torch @@ -316,11 +316,31 @@ def has_idle_worker_for( def num_in_flight(self) -> int: return len(self._running) + def worker_gpu(self, pid: int) -> Optional[int]: + worker = self._workers.get(pid) + return None if worker is None else worker["gpu_id"] + + def is_idle(self, pid: int) -> bool: + gpu_id = self.worker_gpu(pid) + idle = self._idle_pids if gpu_id is None else self._idle_gpu_pids[gpu_id] + return pid in self._workers and pid in idle + + def job_pid(self, job_id: int) -> Optional[int]: + entry = self._running.get(job_id) + return None if entry is None else entry[0] + def _take_worker( - self, gpu_id: Optional[int], allow_gpu_worker_for_cpu: bool + self, + gpu_id: Optional[int], + allow_gpu_worker_for_cpu: bool, + preferred_pids: Sequence[int] = (), ) -> Tuple[int, Optional[int]]: if gpu_id is not None: gpu_idle = self._idle_gpu_pids.get(gpu_id, []) + for pid in preferred_pids: + if pid in gpu_idle: + gpu_idle.remove(pid) + return pid, gpu_id if gpu_idle: return gpu_idle.pop(), gpu_id if not self.gpu_devices and self._idle_pids: @@ -339,12 +359,28 @@ def submit( payload: tuple, gpu_id: Optional[int] = None, allow_gpu_worker_for_cpu: bool = False, + preferred_pids: Sequence[int] = (), ) -> int: + """Run a job on an idle worker, choosing from ``preferred_pids`` first.""" if not self.has_idle_worker_for(gpu_id, allow_gpu_worker_for_cpu): raise RuntimeError("submit() called with no idle worker available") + pid, dispatched_gpu_id = self._take_worker( + gpu_id, allow_gpu_worker_for_cpu, preferred_pids + ) + return self._enqueue(pid, kind, payload, dispatched_gpu_id) + + def submit_to(self, pid: int, kind: str, payload: Any) -> int: + """Run a job on the idle worker ``pid`` on its own device.""" + if not self.is_idle(pid): + raise RuntimeError(f"submit_to() called for busy worker {pid}") + gpu_id = self.worker_gpu(pid) + idle = self._idle_pids if gpu_id is None else self._idle_gpu_pids[gpu_id] + idle.remove(pid) + return self._enqueue(pid, kind, payload, gpu_id) + + def _enqueue(self, pid: int, kind: str, payload: Any, gpu_id: Optional[int]) -> int: job_id = next(self._job_counter) - pid, dispatched_gpu_id = self._take_worker(gpu_id, allow_gpu_worker_for_cpu) - job = _Job(job_id, kind, payload, dispatched_gpu_id) + job = _Job(job_id, kind, payload, gpu_id) self._running[job_id] = (pid, job) self._workers[pid]["job_q"].put(job) return job_id diff --git a/src/main/python/systemds/scuro/representations/image_bind.py b/src/main/python/systemds/scuro/representations/image_bind.py index 0dc448a46fb..d0e366778b3 100644 --- a/src/main/python/systemds/scuro/representations/image_bind.py +++ b/src/main/python/systemds/scuro/representations/image_bind.py @@ -19,6 +19,7 @@ # # ------------------------------------------------------------- import math +import os import numpy as np import torch @@ -37,6 +38,8 @@ from systemds.scuro.representations.unimodal import UnimodalRepresentation from systemds.scuro.representations.utils import ( OwnerAccumulator, + build_pretrained, + load_checkpoint, flatten_owned_sequences, inference_context, save_embeddings, @@ -56,6 +59,7 @@ class ImageBind(UnimodalRepresentation): _SPATIAL_CROPS = 3 _FRAMES_PER_CLIP = 2 _CROP_SIZE = 224 + _CHECKPOINT_PATH = os.path.join(".checkpoints", "imagebind_huge.pth") supports_aggregation_pushdown = True cache_in_worker = True @@ -145,7 +149,13 @@ def _ensure_model(self): "ImageBind requires the optional 'imagebind' package" ) from error if self.model is None: - self.model = imagebind_model.imagebind_huge(pretrained=True) + if os.path.exists(self._CHECKPOINT_PATH): + self.model = build_pretrained( + imagebind_model.imagebind_huge, + load_checkpoint(self._CHECKPOINT_PATH), + ) + else: + 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) diff --git a/src/main/python/systemds/scuro/representations/resnet.py b/src/main/python/systemds/scuro/representations/resnet.py index 8a30c9bd342..9b6792c8dd3 100644 --- a/src/main/python/systemds/scuro/representations/resnet.py +++ b/src/main/python/systemds/scuro/representations/resnet.py @@ -23,6 +23,7 @@ from systemds.scuro.representations.representation import RepresentationStats from systemds.scuro.representations.utils import ( LazyModelMixin, + build_torchvision_model, OwnerAccumulator, flatten_owned_sequences, get_sequence_lengths, @@ -124,7 +125,7 @@ def _build_model(self): "ResNet101": (models.resnet101, models.ResNet101_Weights), "ResNet152": (models.resnet152, models.ResNet152_Weights), }[self.model_name] - model = factory(weights=weights.DEFAULT) + model = build_torchvision_model(factory, weights.DEFAULT) model.eval() for param in model.parameters(): param.requires_grad = False 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 5d5abaac2ea..92d7a746f57 100644 --- a/src/main/python/systemds/scuro/representations/swin_video_transformer.py +++ b/src/main/python/systemds/scuro/representations/swin_video_transformer.py @@ -36,6 +36,7 @@ from systemds.scuro.dataloader.video_loader import VideoStats from systemds.scuro.representations.utils import ( LazyModelMixin, + build_torchvision_model, OwnerAccumulator, get_sequence_lengths, transformer_inference_context, @@ -104,7 +105,9 @@ def __init__( self._activation_hook = None def _build_model(self): - model = swin3d_t(weights=models.video.Swin3D_T_Weights.KINETICS400_V1) + model = build_torchvision_model( + swin3d_t, models.video.Swin3D_T_Weights.KINETICS400_V1 + ) if not self._placed: self.device = get_device_for_model(model, memory_factor=1.5) self._gpu_id = self.device.index diff --git a/src/main/python/systemds/scuro/representations/utils.py b/src/main/python/systemds/scuro/representations/utils.py index 5f3a3d56f73..67077bfcbf7 100644 --- a/src/main/python/systemds/scuro/representations/utils.py +++ b/src/main/python/systemds/scuro/representations/utils.py @@ -199,6 +199,81 @@ def model(self, model): self._model = model +@contextmanager +def _parameters_on_meta(): + """Move every parameter to the meta device as soon as it is registered. + + Weight initializers then run on meta tensors and cost nothing. Buffers stay + real because some modules compute them in ``__init__`` (e.g. ImageBind's + causal mask) and a non-persistent buffer is never restored from a + checkpoint. Whole-module ``torch.device("meta")`` would also break modules + that call ``.item()`` on tensors during construction. + """ + register_parameter = torch.nn.Module.register_parameter + + def register_on_meta(module, name, param): + register_parameter(module, name, param) + if param is not None and not param.is_meta: + module._parameters[name] = torch.nn.Parameter( + param.to("meta"), requires_grad=param.requires_grad + ) + + torch.nn.Module.register_parameter = register_on_meta + try: + yield + finally: + torch.nn.Module.register_parameter = register_parameter + + +def build_pretrained(build_fn, state_dict): + """Build ``build_fn()`` without random weight init and load ``state_dict``. + + Initializing a large model only to overwrite every weight from the + checkpoint dominates load time (about 12 s of ImageBind's 20 s). + """ + with _parameters_on_meta(): + model = build_fn() + model.load_state_dict(state_dict, strict=True, assign=True) + missing = [ + name + for name, tensor in [*model.named_parameters(), *model.named_buffers()] + if tensor.is_meta + ] + if missing: + raise RuntimeError(f"Checkpoint did not provide weights for {missing}") + return model + + +def load_checkpoint(path): + """Load a state dict on the CPU, memory-mapped when the format allows it.""" + try: + return torch.load(path, map_location="cpu", mmap=True, weights_only=True) + except RuntimeError as error: + # Checkpoints in the legacy (pre zip) format cannot be memory-mapped. + if "mmap" not in str(error): + raise + return torch.load(path, map_location="cpu", weights_only=True) + + +def torchvision_state_dict(weights): + """Load torchvision ``weights`` from the torch hub cache.""" + from urllib.parse import urlparse + + file_name = os.path.basename(urlparse(weights.url).path) + path = os.path.join(torch.hub.get_dir(), "checkpoints", file_name) + if not os.path.exists(path): + return weights.get_state_dict(progress=False, check_hash=True) + return load_checkpoint(path) + + +def build_torchvision_model(factory, weights, **kwargs): + """Equivalent to ``factory(weights=weights, **kwargs)`` without random init.""" + kwargs.setdefault("num_classes", len(weights.meta["categories"])) + return build_pretrained( + lambda: factory(weights=None, **kwargs), torchvision_state_dict(weights) + ) + + 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] diff --git a/src/main/python/systemds/scuro/representations/vgg.py b/src/main/python/systemds/scuro/representations/vgg.py index a88cfc6be9b..a6c51f67b66 100644 --- a/src/main/python/systemds/scuro/representations/vgg.py +++ b/src/main/python/systemds/scuro/representations/vgg.py @@ -44,6 +44,7 @@ from systemds.scuro.representations.representation import RepresentationStats from systemds.scuro.representations.utils import ( LazyModelMixin, + build_torchvision_model, OwnerAccumulator, flatten_owned_sequences, get_sequence_lengths, @@ -84,7 +85,7 @@ def __init__( self.batch_size = batch_size def _build_model(self): - model = models.vgg19(weights=models.VGG19_Weights.DEFAULT) + model = build_torchvision_model(models.vgg19, models.VGG19_Weights.DEFAULT) model.eval() for param in model.parameters(): param.requires_grad = False diff --git a/src/main/python/systemds/scuro/representations/x3d.py b/src/main/python/systemds/scuro/representations/x3d.py index 4d9799109c1..6caacc96fb5 100644 --- a/src/main/python/systemds/scuro/representations/x3d.py +++ b/src/main/python/systemds/scuro/representations/x3d.py @@ -38,6 +38,7 @@ from systemds.scuro.representations.unimodal import UnimodalRepresentation from systemds.scuro.representations.utils import ( LazyModelMixin, + build_torchvision_model, get_sequence_lengths, inference_context, save_embeddings, @@ -115,9 +116,11 @@ def __init__( def _build_model(self): if self.model_name == "r3d": - model = r3d_18(pretrained=True) + model = build_torchvision_model( + r3d_18, models.video.R3D_18_Weights.KINETICS400_V1 + ) else: - model = s3d(weights=models.video.S3D_Weights.DEFAULT) + model = build_torchvision_model(s3d, models.video.S3D_Weights.DEFAULT) if not self._placed: self.device = get_device_for_model(model, memory_factor=1.5) self._gpu_id = self.device.index From 210ac53936efde899fa95699ead097642c437597 Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Tue, 13 Oct 2026 09:04:23 +0200 Subject: [PATCH 08/26] float16 to uint8 in histogram to save time --- .../scuro/representations/color_histogram.py | 34 ++++++++++++++++++- 1 file changed, 33 insertions(+), 1 deletion(-) diff --git a/src/main/python/systemds/scuro/representations/color_histogram.py b/src/main/python/systemds/scuro/representations/color_histogram.py index 34147452791..c61f41691dd 100644 --- a/src/main/python/systemds/scuro/representations/color_histogram.py +++ b/src/main/python/systemds/scuro/representations/color_histogram.py @@ -34,6 +34,36 @@ NP_ARRAY_HEADER_BYTES, ) +_FLOAT16_LUTS = None + + +def _float16_luts(): + """uint8 values for every float16 bit pattern, with and without x255.""" + global _FLOAT16_LUTS + if _FLOAT16_LUTS is None: + values = np.arange(1 << 16, dtype=np.uint32).astype(np.uint16) + values = values.view(np.float16) + with np.errstate(invalid="ignore", over="ignore"): + _FLOAT16_LUTS = ( + np.clip(values * 255, 0, 255).astype(np.uint8), + np.clip(values, 0, 255).astype(np.uint8), + ) + return _FLOAT16_LUTS + + +def _float16_to_uint8(image): + """Same result as the generic float path, by table lookup. + + float16 arithmetic and min/max are emulated in numpy and made up 91% of + ColorHistogram's runtime on float16 video. + """ + bits = image.view(np.uint16) + in_unit_range = bool(image.size) and ( + bits.max() <= 0x3C00 or not np.any((bits > 0x3C00) & (bits != 0x8000)) + ) + scaled, unscaled = _float16_luts() + return (scaled if in_unit_range else unscaled)[bits] + @register_representation([ModalityType.IMAGE, ModalityType.VIDEO]) class ColorHistogram(UnimodalRepresentation): @@ -70,7 +100,9 @@ def _get_parameters(self): } def compute_histogram(self, image): - if np.issubdtype(image.dtype, np.floating): + if image.dtype == np.float16: + image = _float16_to_uint8(image) + elif 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) From 98c7d5348fd5e9cbabf30ef13583b3bcb4ef89ba Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Tue, 13 Oct 2026 13:46:55 +0200 Subject: [PATCH 09/26] add efficient model execution to other modalities --- .../systemds/scuro/representations/bert.py | 161 +++++++---- .../systemds/scuro/representations/clip.py | 264 ++++++++++++------ .../systemds/scuro/representations/resnet.py | 89 +++--- .../systemds/scuro/representations/wav2vec.py | 28 +- 4 files changed, 342 insertions(+), 200 deletions(-) diff --git a/src/main/python/systemds/scuro/representations/bert.py b/src/main/python/systemds/scuro/representations/bert.py index c1e0c8e0e4f..02382938615 100644 --- a/src/main/python/systemds/scuro/representations/bert.py +++ b/src/main/python/systemds/scuro/representations/bert.py @@ -29,6 +29,8 @@ LengthBucketBatchSampler, OwnerAccumulator, OwnedSequenceDataset, + capture_layer_outputs, + existing_layers, flatten_owned_sequences, move_batch_to_device, pin_memory_for, @@ -54,6 +56,7 @@ class BertFamily(UnimodalRepresentation): supports_aggregation_pushdown = True cache_in_worker = True + layer_parameter = "layer" MODEL_BYTES = None HIDDEN_SIZE = 768 NUM_LAYERS = 12 @@ -91,8 +94,6 @@ def __init__( 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 or {}).get("batch_size", batch_size)) @@ -212,7 +213,16 @@ def estimate_peak_memory_bytes(self, input_stats): } def transform(self, modality, aggregation=None): - transformed_modality = TransformedModality(modality, self) + layers = self.transform_layers(modality, [self.layer], aggregation) + if self.layer not in layers: + raise ValueError(f"{self.name} has no layer {self.layer!r}") + transformed_modality = layers[self.layer] + if self.output_file is not None: + save_embeddings(transformed_modality.data, self.output_file) + self.assert_output_stats(transformed_modality) + return transformed_modality + + def transform_layers(self, modality, layer_names, aggregation=None): if self.tokenizer is None: self.tokenizer = AutoTokenizer.from_pretrained( self.model_name, clean_up_tokenization_spaces=True @@ -222,53 +232,43 @@ def transform(self, modality, aggregation=None): self.model = self.model.to(self.device) self.model.eval() - self.bert_output = None - - def get_activation(name): - def hook(model, input, output): - if isinstance(output, tuple): - self.bert_output = output[0] - elif hasattr(output, "last_hidden_state"): - self.bert_output = output.last_hidden_state - else: - self.bert_output = output - - return hook + existing = set(existing_layers(self.model, layer_names)) + layer_names = [ + name + for name in dict.fromkeys(layer_names) + if name == "cls" or name in existing + ] + if not layer_names: + return {} aggregate_dim = (0,) - if self.layer != "cls": - 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"): 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( + embeddings = self._embeddings_by_layer( chunks, self.model, self.tokenizer, + layer_names, aggregation, owner_ids=owner_ids, num_owners=len(chunk_groups), grouped=aggregation is None, ) else: - embeddings = self.create_embeddings( - modality.data, self.model, self.tokenizer + embeddings = self._embeddings_by_layer( + modality.data, self.model, self.tokenizer, layer_names ) - 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 - self.assert_output_stats(transformed_modality) - return transformed_modality + transformed = {} + for name in layer_names: + transformed_modality = TransformedModality(modality, self) + transformed_modality.data_type = np.float32 + transformed_modality.aggregate_dim = aggregate_dim + transformed_modality.data = embeddings.pop(name) + transformed[name] = transformed_modality + return transformed def assert_output_stats(self, transformed_modality): if self.stats: @@ -296,6 +296,33 @@ def create_embeddings( num_owners=None, grouped=False, ): + return self._embeddings_by_layer( + data, + model, + tokenizer, + [self.layer], + aggregation, + owner_ids=owner_ids, + num_owners=num_owners, + grouped=grouped, + )[self.layer] + + def _embeddings_by_layer( + self, + data, + model, + tokenizer, + layer_names, + aggregation=None, + owner_ids=None, + num_owners=None, + grouped=False, + ): + """Encode every text once and pool each requested layer's output. + + The layer name "cls" is the CLS token of the last hidden state; any + other name is a submodule whose output is mean-pooled over the tokens. + """ texts = list(TextDataset(data)) single_owner = owner_ids is None and aggregation is not None if owner_ids is None: @@ -342,30 +369,56 @@ def collate(samples): collate_fn=collate, pin_memory=pin_memory_for(self.device), ) - accumulator = OwnerAccumulator(num_owners, len(dataset), aggregation) - with transformer_inference_context(self.device): + accumulators = { + name: OwnerAccumulator(num_owners, len(dataset), aggregation) + for name in layer_names + } + hooked = [name for name in layer_names if name != "cls"] + batch_mask = {} + + def pool_hooked(output): + return pool_transformer_output( + _hidden_state(output), batch_mask["attention_mask"] + ) + + with transformer_inference_context(self.device), capture_layer_outputs( + model, hooked, pool_hooked + ) as outputs: for inputs, batch_owner_ids, chunk_ids in dataloader: inputs = move_batch_to_device(inputs, self.device) - outputs = model(**inputs) - if self.layer == "cls": - hidden_state = outputs.last_hidden_state - else: - 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) + batch_mask["attention_mask"] = inputs["attention_mask"] + model_outputs = model(**inputs) + if "cls" in layer_names: + outputs["cls"] = pool_transformer_output( + model_outputs.last_hidden_state, + inputs["attention_mask"], + use_cls=True, + ) + for name in layer_names: + accumulators[name].update( + outputs[name], batch_owner_ids, chunk_ids + ) + + embeddings = {} + for name in layer_names: + layer_embeddings = accumulators.pop(name).finalize(grouped=grouped) + if single_owner: + layer_embeddings = layer_embeddings[0] + elif aggregation is None and not grouped: + layer_embeddings = list(layer_embeddings) + embeddings[name] = layer_embeddings return embeddings +def _hidden_state(output): + """The token states in a layer's output, whatever form the layer returns.""" + if isinstance(output, tuple): + return output[0] + if hasattr(output, "last_hidden_state"): + return output.last_hidden_state + return output + + @register_representation(ModalityType.TEXT) @register_expensive_representation(ModalityType.TEXT) class Bert(BertFamily): @@ -443,8 +496,8 @@ def __init__( "encoder.layer.9", "encoder.layer.10", "encoder.layer.11", - "pooler", - "pooler.activation", + # No "pooler": roberta-base ships no pooler weights, so it is + # randomly initialized on every load. ] } super().__init__( diff --git a/src/main/python/systemds/scuro/representations/clip.py b/src/main/python/systemds/scuro/representations/clip.py index 0e7c2951be8..5be30b42f82 100644 --- a/src/main/python/systemds/scuro/representations/clip.py +++ b/src/main/python/systemds/scuro/representations/clip.py @@ -29,6 +29,8 @@ LazyModelMixin, LengthBucketBatchSampler, OwnerAccumulator, + capture_layer_outputs, + existing_layers, OwnedSequenceDataset, flatten_owned_sequences, get_sequence_lengths, @@ -72,6 +74,7 @@ class CLIPVisual(LazyModelMixin, UnimodalRepresentation): OUTPUT_MODALITY_TYPE = ModalityType.EMBEDDING supports_aggregation_pushdown = True cache_in_worker = True + layer_parameter = "layer_name" MODEL_BYTES = CLIP_MODEL_BYTES HIDDEN_SIZE = 768 NUM_LAYERS = 12 @@ -81,7 +84,6 @@ def __init__(self, output_file=None, batch_size=32, layer_name="", params=None): parameters = self._get_parameters() super().__init__("CLIPVisual", self.OUTPUT_MODALITY_TYPE, parameters) self.params = params - self._activation_hook = None self._processor = None if params is not None: self.batch_size = int((params or {}).get("batch_size", batch_size)) @@ -255,46 +257,54 @@ def estimate_peak_memory_bytes(self, input_stats) -> dict: } def transform(self, modality, aggregation=None): - transformed_modality = TransformedModality( - modality, self, self.output_modality_type - ) + layers = self.transform_layers(modality, [self.layer_name], aggregation) + if self.layer_name not in layers: + raise ValueError(f"CLIPVisual has no layer {self.layer_name!r}") + transformed_modality = layers[self.layer_name] + if self.output_file is not None: + save_embeddings(transformed_modality.data, self.output_file) + return transformed_modality + + def transform_layers(self, modality, layer_names, 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.model = self.model.to(self.device) self.model.eval() - self.clip_output = None + layer_names = _clip_layers(self.model.vision_model, layer_names) + if not layer_names: + return {} - def get_activation(name): - def hook(model, input, output): - self.clip_output = output[0] if isinstance(output, tuple) else output - - return hook - - 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: - self._activation_hook = layer.register_forward_hook( - get_activation(name) - ) - break - - 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,) + embeddings = self._visual_embeddings_by_layer( + modality, aggregation, layer_names ) - transformed_modality.data = embeddings - return transformed_modality + transformed = {} + for name in layer_names: + 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 modality.modality_type == ModalityType.IMAGE + else (0,) + ) + transformed_modality.data = embeddings.pop(name) + transformed[name] = transformed_modality + return transformed def create_visual_embeddings(self, modality, aggregation=None): + return self._visual_embeddings_by_layer( + modality, aggregation, [self.layer_name] + )[self.layer_name] + + def _visual_embeddings_by_layer(self, modality, aggregation, layer_names): + """Embed every frame once and pool each requested layer's output. + + The empty layer name is CLIP's projected image embedding. + """ preprocess = FramePreprocessor.from_hf_processor( self.processor, dtype=self.data_type ) @@ -315,28 +325,48 @@ def create_visual_embeddings(self, modality, aggregation=None): 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) + accumulators = { + name: OwnerAccumulator(len(modality.data), len(dataset), aggregation) + for name in layer_names + } + hooked = [name for name in layer_names if name != ""] + project = "" in layer_names + + def pool_hooked(output): + return self._pool_visual_output( + output[0] if isinstance(output, tuple) else output + ) - with transformer_inference_context(self.device): + with transformer_inference_context(self.device), capture_layer_outputs( + self.model.vision_model, hooked, pool_hooked + ) as outputs: for batch in dataloader: chunk_ids = batch["id"].long() inputs = {"pixel_values": preprocess(batch["data"], self.device)} - if self.layer_name != "": - _ = self.model.vision_model(**inputs) - output = self.clip_output + if not hooked: + outputs[""] = self.model.get_image_features(**inputs) else: - output = self.model.get_image_features(**inputs) - - output = self._pool_visual_output(output) - 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: - return list(embeddings) + vision_outputs = self.model.vision_model(**inputs) + if project: + # What get_image_features computes, from this same pass. + outputs[""] = self.model.visual_projection(vision_outputs[1]) + owners = owner_by_chunk.index_select(0, chunk_ids) + for name in layer_names: + output = outputs[name] + if name == "": + output = self._pool_visual_output(output) + accumulators[name].update( + torch.flatten(output, 1), owners, chunk_ids + ) + + embeddings = {} + for name in layer_names: + layer_embeddings = accumulators.pop(name).finalize( + grouped=not is_image and aggregation is None + ) + if is_image and aggregation is None: + layer_embeddings = list(layer_embeddings) + embeddings[name] = layer_embeddings return embeddings def _pool_visual_output(self, output: torch.Tensor) -> torch.Tensor: @@ -355,6 +385,7 @@ def _pool_visual_output(self, output: torch.Tensor) -> torch.Tensor: class CLIPText(UnimodalRepresentation): supports_aggregation_pushdown = True cache_in_worker = True + layer_parameter = "layer_name" MODEL_BYTES = CLIP_MODEL_BYTES HIDDEN_SIZE = 512 NUM_LAYERS = 12 @@ -381,7 +412,6 @@ 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): @@ -494,54 +524,54 @@ def estimate_peak_memory_bytes(self, input_stats) -> dict: return {"cpu_peak_bytes": cpu_peak, "gpu_peak_bytes": gpu_peak} def transform(self, modality, aggregation=None): - transformed_modality = TransformedModality( - modality, self, self.output_modality_type - ) + layers = self.transform_layers(modality, [self.layer_name], aggregation) + if self.layer_name not in layers: + raise ValueError(f"CLIPText has no layer {self.layer_name!r}") + transformed_modality = layers[self.layer_name] + if self.output_file is not None: + save_embeddings(transformed_modality.data, self.output_file) + return transformed_modality + + def transform_layers(self, modality, layer_names, aggregation=None): if self.processor is None: self.processor = CLIPProcessor.from_pretrained(CLIP_MODEL_NAME) if self.model is None: self.model = CLIPModel.from_pretrained(CLIP_MODEL_NAME) 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] if isinstance(output, tuple) else output - - return hook - - 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: - self._activation_hook = layer.register_forward_hook( - get_activation(name) - ) - break + layer_names = _clip_layers(self.model.text_model, layer_names) + if not layer_names: + return {} aggregate_dim = None if ModalityType.TEXT.has_field(modality.metadata, "text_spans"): 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( + embeddings = self._text_embeddings_by_layer( chunks, self.model, + layer_names, 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) + embeddings = self._text_embeddings_by_layer( + modality.data, self.model, layer_names + ) - transformed_modality.data_type = np.float32 - transformed_modality.aggregate_dim = aggregate_dim - transformed_modality.data = embeddings - return transformed_modality + transformed = {} + for name in layer_names: + transformed_modality = TransformedModality( + modality, self, self.output_modality_type + ) + transformed_modality.data_type = np.float32 + transformed_modality.aggregate_dim = aggregate_dim + transformed_modality.data = embeddings.pop(name) + transformed[name] = transformed_modality + return transformed def create_text_embeddings( self, @@ -552,6 +582,30 @@ def create_text_embeddings( num_owners=None, grouped=False, ): + return self._text_embeddings_by_layer( + data, + model, + [self.layer_name], + aggregation, + owner_ids=owner_ids, + num_owners=num_owners, + grouped=grouped, + )[self.layer_name] + + def _text_embeddings_by_layer( + self, + data, + model, + layer_names, + aggregation=None, + owner_ids=None, + num_owners=None, + grouped=False, + ): + """Encode every text once and pool each requested layer's output. + + The empty layer name is CLIP's projected text embedding. + """ texts = list(TextDataset(data)) single_owner = owner_ids is None and aggregation is not None if owner_ids is None: @@ -593,22 +647,50 @@ def collate(samples): collate_fn=collate, pin_memory=pin_memory_for(self.device), ) - accumulator = OwnerAccumulator(num_owners, len(dataset), aggregation) - with transformer_inference_context(self.device): + accumulators = { + name: OwnerAccumulator(num_owners, len(dataset), aggregation) + for name in layer_names + } + hooked = [name for name in layer_names if name != ""] + project = "" in layer_names + batch_mask = {} + + def pool_hooked(output): + return pool_transformer_output( + output[0] if isinstance(output, tuple) else output, + batch_mask["attention_mask"], + ) + + with transformer_inference_context(self.device), capture_layer_outputs( + model.text_model, hooked, pool_hooked + ) as outputs: 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) - pooled = pool_transformer_output( - self.clip_output, inputs["attention_mask"] - ) + batch_mask["attention_mask"] = inputs["attention_mask"] + if not hooked: + outputs[""] = model.get_text_features(**inputs) else: - 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) + text_outputs = model.text_model(**inputs) + if project: + # What get_text_features computes, from this same pass. + outputs[""] = model.text_projection(text_outputs[1]) + for name in layer_names: + accumulators[name].update(outputs[name], batch_owner_ids, chunk_ids) + + embeddings = {} + for name in layer_names: + layer_embeddings = accumulators.pop(name).finalize(grouped=grouped) + if single_owner: + layer_embeddings = layer_embeddings[0] + elif aggregation is None and not grouped: + layer_embeddings = list(layer_embeddings) + embeddings[name] = layer_embeddings return embeddings + + +def _clip_layers(model, layer_names): + """The requested layers that exist; the empty name is the projected embedding.""" + existing = set(existing_layers(model, layer_names)) + return [ + name for name in dict.fromkeys(layer_names) if name == "" or name in existing + ] diff --git a/src/main/python/systemds/scuro/representations/resnet.py b/src/main/python/systemds/scuro/representations/resnet.py index 9b6792c8dd3..0a90541fdf7 100644 --- a/src/main/python/systemds/scuro/representations/resnet.py +++ b/src/main/python/systemds/scuro/representations/resnet.py @@ -24,6 +24,8 @@ from systemds.scuro.representations.utils import ( LazyModelMixin, build_torchvision_model, + capture_layer_outputs, + existing_layers, OwnerAccumulator, flatten_owned_sequences, get_sequence_lengths, @@ -37,7 +39,6 @@ ) from systemds.scuro.modality.transformed import TransformedModality from systemds.scuro.representations.unimodal import UnimodalRepresentation -from typing import Tuple, Any from systemds.scuro.drsearch.operator_registry import ( register_representation, register_expensive_representation, @@ -60,6 +61,7 @@ def forward(self, input_: torch.Tensor) -> torch.Tensor: class ResNet(LazyModelMixin, UnimodalRepresentation): supports_aggregation_pushdown = True cache_in_worker = True + layer_parameter = "layer_name" MODEL_BYTES = { "ResNet18": 46_796_608, "ResNet34": 87_259_072, @@ -77,8 +79,6 @@ def __init__( params=None, ): self.data_type = torch.float32 - self._activation_hook = None - self.activation = None self.gpu_id = None self.device = get_device() if params is not None: @@ -209,27 +209,19 @@ def _get_parameters(self, high_level=True): return parameters def transform(self, modality, aggregation=None): + layers = self.transform_layers(modality, [self.layer_name], aggregation) + if self.layer_name not in layers: + raise ValueError(f"{self.model_name} has no layer {self.layer_name!r}") + return layers[self.layer_name] + + def transform_layers(self, modality, layer_names, aggregation=None): if next(self.model.parameters()).dtype != self.data_type: self.model = self.model.to(self.data_type) 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 - ): - self.activation = output - - return hook - - if self.layer_name and 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_features(name) - ) - break + layer_names = existing_layers(self.model, layer_names) + if not layer_names: + return {} is_image = modality.modality_type == ModalityType.IMAGE if is_image: @@ -249,31 +241,42 @@ def hook( 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) + accumulators = { + name: OwnerAccumulator(len(modality.data), len(dataset), aggregation) + for name in layer_names + } - with inference_context(self.device): + with inference_context(self.device), capture_layer_outputs( + self.model, layer_names, _pooled_features + ) as outputs: for batch in dataloader: chunk_ids = batch["id"].long() - _ = self.model(preprocess(batch["data"], self.device)) - 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, - ) + self.model(preprocess(batch["data"], self.device)) + owners = owner_by_chunk.index_select(0, chunk_ids) + for name in layer_names: + accumulators[name].update(outputs[name], owners, chunk_ids) - embeddings = accumulator.finalize(grouped=not is_image and aggregation is None) - if is_image and aggregation is None: - embeddings = list(embeddings) + transformed = {} + for name in layer_names: + embeddings = accumulators.pop(name).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 - ) - 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 + 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 + transformed[name] = transformed_modality + return transformed + + +def _pooled_features(output: torch.Tensor) -> torch.Tensor: + if output.ndim > 2: + output = torch.nn.functional.adaptive_avg_pool2d(output, (1, 1)) + return torch.flatten(output, 1) diff --git a/src/main/python/systemds/scuro/representations/wav2vec.py b/src/main/python/systemds/scuro/representations/wav2vec.py index 2d40f541533..478239a4f41 100644 --- a/src/main/python/systemds/scuro/representations/wav2vec.py +++ b/src/main/python/systemds/scuro/representations/wav2vec.py @@ -139,8 +139,7 @@ def collate(batch): 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 + features = _extract_features(model, inputs["input_values"]) attention_mask = inputs.get("attention_mask") if attention_mask is not None and hasattr( model, "_get_feature_vector_attention_mask" @@ -177,14 +176,11 @@ def estimate_peak_memory_bytes(self, input_stats) -> dict: signal_len = 16000 signal_len = max(signal_len, 1) - hidden = 768 - stride = 320 # conv frontend effective stride - frames = max(1, int(np.ceil(signal_len / stride))) - model_resident = 420 * 1024 * 1024 # ~420 MB batch = min(max(1, self.batch_size), max(1, n)) - activation_bytes = int(batch * frames * hidden * 4 * 24) - attention_bytes = int(batch * 12 * frames * frames * 4) + # Only the CNN front end runs. Its first conv layer (512 channels, + # stride 5) dominates: measured peaks are ~1 KB per input sample. + activation_bytes = int(batch * signal_len * 1024) io_temp = int(signal_len * 4 * 4) + 16 * 1024 * 1024 output_bytes = n * 512 * 4 @@ -193,10 +189,18 @@ def estimate_peak_memory_bytes(self, input_stats) -> dict: (model_resident + activation_bytes + io_temp + output_bytes) * 1.25 ) - gpu_peak = int( - (model_resident + activation_bytes + attention_bytes + batch * io_temp) - * 1.25 - ) + gpu_peak = int((model_resident + activation_bytes + batch * io_temp) * 1.25) cpu_peak = max(cpu_peak, 600 * 1024 * 1024) return {"cpu_peak_bytes": cpu_peak, "gpu_peak_bytes": gpu_peak} + + +def _extract_features(model, input_values): + """``Wav2Vec2Model(...).extract_features`` without running the encoder. + + That output is the layer-normed CNN front end, which the transformer + encoder does not change, so computing only the front end gives the same + features at a fraction of the cost. + """ + features = model.feature_extractor(input_values).transpose(1, 2) + return model.feature_projection.layer_norm(features) From c04fa4660864317c802d49179201747a7d9c37f3 Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Tue, 13 Oct 2026 13:47:41 +0200 Subject: [PATCH 10/26] cache all layers for hpo --- .../scuro/drsearch/hyperparameter_tuner.py | 39 +++- .../systemds/scuro/drsearch/process_cache.py | 138 ++++++++++++++- .../scuro/drsearch/representation_dag.py | 26 ++- .../scuro/drsearch/unimodal_optimizer.py | 3 +- .../scuro/modality/unimodal_modality.py | 166 +++++++++++------- .../scuro/representations/unimodal.py | 14 ++ .../systemds/scuro/representations/utils.py | 37 ++++ .../systemds/scuro/representations/x3d.py | 80 +++++---- 8 files changed, 392 insertions(+), 111 deletions(-) diff --git a/src/main/python/systemds/scuro/drsearch/hyperparameter_tuner.py b/src/main/python/systemds/scuro/drsearch/hyperparameter_tuner.py index 7805ac66b3f..f99b1eb05d2 100644 --- a/src/main/python/systemds/scuro/drsearch/hyperparameter_tuner.py +++ b/src/main/python/systemds/scuro/drsearch/hyperparameter_tuner.py @@ -587,6 +587,26 @@ def tune_dag_representation( print(f"[hpo-cache] {task.model.name}: {self._trial_cache.summary()}") self._trial_cache.clear() self._trial_cache = None + self.execution_cache.set_layer_sweeps(None) + + def _layer_sweeps(self, dag, param_specs) -> Dict[str, List[Any]]: + """Per node, the layers a search over its layer parameter can select. + + An encoder computes all of them in one pass the first time any is + needed, so trying another layer reuses that pass instead of running + the encoder again. SCURO_HPO_LAYER_SWEEP=0 turns this off. + """ + if os.environ.get("SCURO_HPO_LAYER_SWEEP", "1") == "0": + return {} + sweeps = {} + for spec in param_specs: + node_id, _, param = spec["name"].rpartition("-") + if spec["type"] != "categorical" or len(spec["domain"]) < 2: + continue + operation = dag.get_node_by_id(node_id).operation + if param == getattr(operation, "layer_parameter", None): + sweeps[node_id] = list(spec["domain"]) + return sweeps def _tune_dag_representation( self, dag, root_node_id, task, max_evals=None, mm_opt=False @@ -622,6 +642,18 @@ def visit_node(node_id): self._get_cached_modalities_for_task(task, modality_ids) if mm_opt else None ) + param_specs = [] + tied_params = {} + if hyperparams: + node_parameters = { + node_id: (dag.get_node_by_id(node_id).parameters or {}) + for node_id in hyperparams + } + param_specs = self._build_param_specs(hyperparams, node_parameters) + tied_params = _tied_window_params(dag) + param_specs = [s for s in param_specs if s["name"] not in tied_params] + self.execution_cache.set_layer_sweeps(self._layer_sweeps(dag, param_specs)) + baseline_params, baseline_raw_scores = self.evaluate_dag_config( dag, {}, @@ -638,13 +670,6 @@ def visit_node(node_id): if not hyperparams: all_results = [baseline] else: - node_parameters = { - node_id: (dag.get_node_by_id(node_id).parameters or {}) - for node_id in hyperparams - } - param_specs = self._build_param_specs(hyperparams, node_parameters) - tied_params = _tied_window_params(dag) - param_specs = [s for s in param_specs if s["name"] not in tied_params] discrete_size = self._estimate_discrete_search_size(param_specs) n_calls = min(discrete_size, max_evals) if max_evals else discrete_size all_results = self._search_best_configs( diff --git a/src/main/python/systemds/scuro/drsearch/process_cache.py b/src/main/python/systemds/scuro/drsearch/process_cache.py index 361db144ff0..357190df0ef 100644 --- a/src/main/python/systemds/scuro/drsearch/process_cache.py +++ b/src/main/python/systemds/scuro/drsearch/process_cache.py @@ -19,14 +19,16 @@ # # ------------------------------------------------------------- import hashlib +import logging import os import threading from collections import OrderedDict -from typing import Any, Callable, Hashable, Optional, Tuple +from typing import Any, Callable, Dict, Hashable, Iterable, List, Optional, Tuple import numpy as np _MISSING = object() +logger = logging.getLogger(__name__) def freeze_cache_value(value: Any) -> Hashable: @@ -160,6 +162,49 @@ def get_or_compute(self, key: Hashable, compute: Callable[[], Any]) -> Any: pending.set() pending.wait() + def get_or_compute_group( + self, + key: Hashable, + group_key: Hashable, + compute_group: Callable[[], Dict[Hashable, Any]], + ) -> Any: + """Like get_or_compute, but one computation yields several entries. + + ``compute_group`` returns {key: value} including ``key``; every entry is + stored. A caller that misses a different key of the same group while it + is being computed waits for it instead of starting a second computation. + """ + if self.max_entries == 0: + return compute_group()[key] + + while True: + with self._lock: + value = self._values.pop(key, _MISSING) + if value is not _MISSING: + self._values[key] = value + return value + pending = self._pending.get(group_key) + if pending is None: + pending = threading.Event() + self._pending[group_key] = pending + owner = True + else: + owner = False + + if owner: + try: + values = compute_group() + for other_key, other_value in values.items(): + if other_key != key: + self.put(other_key, other_value) + self.put(key, values[key]) + return values[key] + finally: + with self._lock: + self._pending.pop(group_key, None) + pending.set() + pending.wait() + def clear(self) -> None: with self._lock: self._values.clear() @@ -178,6 +223,7 @@ def __init__( self.representations = BoundedProcessCache(representation_entries) self.should_cache_representation = should_cache_representation self._modality_keys = {} + self._layer_sweeps = {} self._lock = threading.RLock() def modality_key(self, modality) -> Tuple[str, str]: @@ -251,7 +297,97 @@ def get_or_compute_representation(self, operation, key, compute): return compute() return self.representations.get_or_compute(key, compute) + def set_layer_sweeps(self, sweeps: Optional[Dict[str, Iterable[Any]]]) -> None: + """Set, per DAG node id, the layers to compute together on a miss.""" + with self._lock: + self._layer_sweeps = { + node_id: list(layers) for node_id, layers in (sweeps or {}).items() + } + + def layer_sweep(self, node_id: str) -> Optional[List[Any]]: + with self._lock: + return self._layer_sweeps.get(node_id) + + def get_or_compute_layer( + self, + operation, + parameters, + input_keys, + layer_names: List[Any], + compute_layers: Callable[[List[Any]], Dict[Any, Any]], + compute: Callable[[], Any], + ): + """Return one layer's representation, computing all ``layer_names`` at once. + + On a miss, ``compute_layers(layers)`` runs the encoder once for every + layer and returns {layer: representation}; each one is cached under the + key that a configuration selecting that layer looks up. If that pass + fails or lacks the requested layer, ``compute()`` computes it alone. + """ + parameters = dict(parameters or {}) + key = self.representation_key(operation, parameters, input_keys) + layer_parameter = getattr(operation, "layer_parameter", None) + if layer_parameter is None or layer_parameter not in parameters: + return self.get_or_compute_representation(operation, key, compute) + if ( + self.should_cache_representation is not None + and not self.should_cache_representation(operation) + ): + return compute() + + layer = parameters[layer_parameter] + layers = list(dict.fromkeys([layer, *layer_names])) + keys = { + name: self.representation_key( + operation, {**parameters, layer_parameter: name}, input_keys + ) + for name in layers + } + shared = {k: v for k, v in parameters.items() if k != layer_parameter} + group_key = ( + "layers", + layer_parameter, + self.representation_key(operation, shared, input_keys), + ) + + def compute_group(): + try: + values = compute_layers(layers) + except Exception as error: + logger.warning( + "Computing %d layers of %s in one pass failed (%s); " + "computing %s=%r alone", + len(layers), + operation.__name__, + error, + layer_parameter, + layer, + ) + values = {} + _release_cuda_cache() + computed = { + keys[name]: value for name, value in values.items() if name in keys + } + if key not in computed: + # Keep the other layers even if this one cannot be computed. + for other_key, value in computed.items(): + self.representations.put(other_key, value) + computed[key] = compute() + return computed + + return self.representations.get_or_compute_group(key, group_key, compute_group) + def clear(self) -> None: self.representations.clear() with self._lock: self._modality_keys.clear() + self._layer_sweeps.clear() + + +def _release_cuda_cache() -> None: + try: + import torch + except ImportError: + return + if torch.cuda.is_available(): + torch.cuda.empty_cache() diff --git a/src/main/python/systemds/scuro/drsearch/representation_dag.py b/src/main/python/systemds/scuro/drsearch/representation_dag.py index 3d3a48d553e..860b075561e 100644 --- a/src/main/python/systemds/scuro/drsearch/representation_dag.py +++ b/src/main/python/systemds/scuro/drsearch/representation_dag.py @@ -435,7 +435,31 @@ def run_operation(): result = input_mods[0].combine(input_mods[1:], fusion_op) return result - if execution_cache is not None and is_unimodal: + def run_layers(layer_names): + node_operation = node.operation(params=node.parameters) + if gpu_id is not None and hasattr(node_operation, "gpu_id"): + node_operation.gpu_id = gpu_id + return input_mods[0].apply_representation_layers( + node_operation, + layer_names, + aggregation=pushdown_aggregation_for_node(node.parameters), + ) + + layer_sweep = ( + execution_cache.layer_sweep(node_id) + if execution_cache is not None and is_unimodal and rep_cache is None + else None + ) + if layer_sweep and hasattr(input_mods[0], "apply_representation_layers"): + result = execution_cache.get_or_compute_layer( + node.operation, + node.parameters, + [execution_key(input_id) for input_id in node.inputs], + layer_sweep, + run_layers, + run_operation, + ) + elif execution_cache is not None and is_unimodal: result = execution_cache.get_or_compute_representation( node.operation, execution_key(node_id), run_operation ) diff --git a/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py b/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py index 8b717e5ced3..afa22a3f38b 100644 --- a/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py +++ b/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py @@ -702,7 +702,8 @@ def add_aggregation_operator(self, builder, dags): for dag in dags: if self._aggregation_needed(dag): agg_op = AggregatedRepresentation( - target_dimensions=self.expected_dimensions + target_dimensions=self.expected_dimensions, + aggregate_leading=True, ) agg_node_id = builder.create_operation_node( agg_op.__class__, diff --git a/src/main/python/systemds/scuro/modality/unimodal_modality.py b/src/main/python/systemds/scuro/modality/unimodal_modality.py index 7a40cd59360..1b6173d8aab 100644 --- a/src/main/python/systemds/scuro/modality/unimodal_modality.py +++ b/src/main/python/systemds/scuro/modality/unimodal_modality.py @@ -212,32 +212,70 @@ def apply_representations( if len(aggregations) != len(representations): raise ValueError("aggregations must match representations") - representation_specs = list( - zip(representation_keys, representations, aggregations) + units = [ + ( + (representation_key,), + representation, + partial( + _keyed_transform, + representation_key, + representation, + self, + rep_aggregation, + ), + ) + for representation_key, representation, rep_aggregation in zip( + representation_keys, representations, aggregations + ) + ] + return self._apply_transform_units(units, parallel) + + def apply_representation_layers( + self, representation, layer_names, aggregation=None + ): + """Apply one encoder for several of its layers in one pass over the data. + + Returns {layer name: transformed modality} for the layers the encoder has; + the missing ones are recorded in ``failed_representations``. + """ + layer_names = tuple(dict.fromkeys(layer_names)) + transform = partial( + representation.transform_layers, self, list(layer_names), aggregation + ) + return self._apply_transform_units( + [(layer_names, representation, transform)], parallel=False ) + + def _apply_transform_units(self, units, parallel): + """Run each unit over every data chunk and gather its outputs by key. + + A unit is (keys, representation, transform), where one call of transform + returns {key: transformed chunk} for its keys. + """ transformed_modalities_per_representation = {} padding_per_representation = {} original_lengths_per_representation = {} failed_representations = {} - 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_key] = ( - transformed_modality - ) - padding_per_representation[representation_key] = False - original_lengths_per_representation[representation_key] = [] + for keys, representation, _ in units: + for representation_key in keys: + transformed_modality = TransformedModality( + self, representation.name, representation.output_modality_type + ) + transformed_modality.data = [] + transformed_modality.metadata = [] + transformed_modalities_per_representation[representation_key] = ( + transformed_modality + ) + padding_per_representation[representation_key] = False + original_lengths_per_representation[representation_key] = [] start = ( time.time() ) # TODO: should be repalced in unimodal_representation.transform with ( - ThreadPoolExecutor(max_workers=len(representations)) - if parallel and len(representations) > 1 + ThreadPoolExecutor(max_workers=len(units)) + if parallel and len(units) > 1 else nullcontext() ) as executor: time_s = time.time() @@ -249,27 +287,18 @@ def apply_representations( for _ in chunks: representations_futures = {} serial_transforms = [] - for ( - representation_key, - representation, - rep_aggregation, - ) in representation_specs: - if representation_key in failed_representations: + for keys, _, transform in units: + pending_keys = [ + key for key in keys if key not in failed_representations + ] + if not pending_keys: continue if executor is None: - serial_transforms.append( - ( - representation_key, - partial( - representation.transform, self, rep_aggregation - ), - ) - ) + serial_transforms.append((pending_keys, transform)) else: - future = executor.submit( - representation.transform, self, rep_aggregation + representations_futures[executor.submit(transform)] = ( + pending_keys ) - representations_futures[future] = representation_key completed = ( ( (representations_futures[future], future.result) @@ -278,47 +307,56 @@ def apply_representations( if executor is not None else serial_transforms ) - for representation_key, transform in completed: + for pending_keys, transform in completed: try: - transformed_chunk = transform() + transformed_chunks = transform() except Exception as e: - failed_representations[representation_key] = e + for representation_key in pending_keys: + 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_key - ].append(int(shape[0])) + for representation_key in pending_keys: + transformed_chunk = transformed_chunks.get(representation_key) + if transformed_chunk is None: + failed_representations[representation_key] = ValueError( + f"No output for {representation_key!r}" + ) + 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_key + ].append(int(shape[0])) if self.data_loader.is_chunked: print(f"Time for transforming data chunks: {time.time() - time_s}") for representation_name in failed_representations: transformed_modalities_per_representation.pop(representation_name, None) - if representations and not transformed_modalities_per_representation: + if units 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_key], - original_lengths_per_representation[representation_key], - padding_per_representation[representation_key], - ) - transformed_modalities_per_representation[ - representation_key - ].transform_time += (time.time() - start) - transformed_modalities_per_representation[ - representation_key - ].self_contained = representation.self_contained + for keys, representation, _ in units: + for representation_key in keys: + if representation_key in failed_representations: + continue + self._apply_padding( + transformed_modalities_per_representation[representation_key], + original_lengths_per_representation[representation_key], + padding_per_representation[representation_key], + ) + transformed_modalities_per_representation[ + representation_key + ].transform_time += (time.time() - start) + transformed_modalities_per_representation[ + representation_key + ].self_contained = representation.self_contained self.failed_representations = failed_representations return transformed_modalities_per_representation @@ -384,3 +422,7 @@ def _apply_padding(self, modality, original_lengths, pad_dim_one): ) modality.data = padded_embeddings modality.update_metadata() + + +def _keyed_transform(representation_key, representation, modality, aggregation): + return {representation_key: representation.transform(modality, aggregation)} diff --git a/src/main/python/systemds/scuro/representations/unimodal.py b/src/main/python/systemds/scuro/representations/unimodal.py index 3ad787c6773..d871cef8b8a 100644 --- a/src/main/python/systemds/scuro/representations/unimodal.py +++ b/src/main/python/systemds/scuro/representations/unimodal.py @@ -24,6 +24,11 @@ class UnimodalRepresentation(Representation): + # Parameter that selects which layer's activations an encoder returns. An + # encoder that sets it implements transform_layers, so a search over that + # parameter can obtain every layer from one pass over the data. + layer_parameter = None + def __init__( self, name: str, output_modality_type, parameters=None, self_contained=True ): @@ -46,6 +51,15 @@ def transform(self, data, aggregation=None): # TODO: check if there is a way to time the transformation in here (needed for chunked execution) raise f"Not implemented for {self.name}" + def transform_layers(self, data, layer_names, aggregation=None): + """Return {layer name: transformed modality} from one pass over the data. + + Layers that the model does not have are left out of the result. + """ + raise NotImplementedError( + f"transform_layers is not implemented for {self.name}" + ) + class PixelRepresentation(UnimodalRepresentation): def __init__(self): diff --git a/src/main/python/systemds/scuro/representations/utils.py b/src/main/python/systemds/scuro/representations/utils.py index 67077bfcbf7..76d339256ad 100644 --- a/src/main/python/systemds/scuro/representations/utils.py +++ b/src/main/python/systemds/scuro/representations/utils.py @@ -327,6 +327,43 @@ def inference_context(device): transformer_inference_context = inference_context +def existing_layers(model, layer_names): + """Return the requested layer names that name a submodule of ``model``. + + The empty name, which ``named_modules`` gives the model itself, is not a layer. + """ + modules = dict(model.named_modules()) + return [name for name in dict.fromkeys(layer_names) if name and name in modules] + + +@contextmanager +def capture_layer_outputs(model, layer_names, reduce): + """Hook every named layer and keep ``reduce(output)`` of the latest forward. + + One forward pass then yields the features of all layers. Reducing inside + the hook keeps only the pooled vectors on the device instead of every full + activation map until the pass ends; no hooked layer's output is modified + in place later in the pass, so this equals reducing after the pass. + """ + modules = dict(model.named_modules()) + outputs = {} + handles = [] + + def hook_for(name): + def hook(_module, _input, output): + outputs[name] = reduce(output) + + return hook + + try: + for name in layer_names: + handles.append(modules[name].register_forward_hook(hook_for(name))) + yield outputs + finally: + for handle in handles: + handle.remove() + + class OwnerAccumulator: """Collect chunk vectors on-device and restore or aggregate their owners.""" diff --git a/src/main/python/systemds/scuro/representations/x3d.py b/src/main/python/systemds/scuro/representations/x3d.py index 6caacc96fb5..c0838186cfe 100644 --- a/src/main/python/systemds/scuro/representations/x3d.py +++ b/src/main/python/systemds/scuro/representations/x3d.py @@ -39,6 +39,8 @@ from systemds.scuro.representations.utils import ( LazyModelMixin, build_torchvision_model, + capture_layer_outputs, + existing_layers, get_sequence_lengths, inference_context, save_embeddings, @@ -69,6 +71,7 @@ def forward(self, input_: torch.Tensor) -> torch.Tensor: @register_expensive_representation([ModalityType.VIDEO]) class X3D(LazyModelMixin, FrameBudgetedVideoModel, UnimodalRepresentation): cache_in_worker = True + layer_parameter = "layer_name" MODEL_BYTES = {"r3d": 133_524_448, "s3d": 33_365_672} _MEMORY_PROFILES = { "r3d": VideoMemoryProfile( @@ -111,8 +114,6 @@ def __init__( self.layer_name = layer self.batch_size = batch_size self.max_batch_frames = max_batch_frames - self._activation_hook = None - self.activation = None def _build_model(self): if self.model_name == "r3d": @@ -227,60 +228,61 @@ def _collate_videos(samples): return {"id": video_ids, "data": stack_or_list(videos), "lengths": lengths} def transform(self, modality, aggregation=None): + layers = self.transform_layers(modality, [self.layer_name], aggregation) + if self.layer_name not in layers: + raise ValueError(f"{self.model_name} has no layer {self.layer_name!r}") + transformed_modality = layers[self.layer_name] + if self.output_file is not None: + save_embeddings(transformed_modality.data, self.output_file) + return transformed_modality + + def transform_layers(self, modality, layer_names, 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 - ): - self.activation = output - - return hook - - if self.layer_name and 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_features(name) - ) - break + layer_names = existing_layers(self.model, layer_names) + if not layer_names: + return {} dataset = FrameDataset(modality.data, video=True) preprocess = FramePreprocessor(dtype=self.data_type) lengths = get_sequence_lengths(modality.data, modality.metadata) - embeddings = [None] * len(dataset) + embeddings = {name: [None] * len(dataset) for name in layer_names} def forward(model, batch): - video_ids = batch["id"].long() + video_ids = batch["id"].long().cpu().tolist() frames = preprocess(batch["data"], next(model.parameters()).device) for i, length in enumerate(batch["lengths"].tolist()): frames[i, length:] = 0 - _ = model(frames.permute(0, 2, 1, 3, 4)) - 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 + # The CPU fallback runs a copy of the model, so hook the one in use. + with capture_layer_outputs(model, layer_names, _pooled_features) as outputs: + model(frames.permute(0, 2, 1, 3, 4)) + for name in layer_names: + vectors = outputs[name].detach().float().cpu().numpy() + for video_id, vector in zip(video_ids, vectors): + embeddings[name][video_id] = vector with inference_context(self.device): self._run_batches( dataset, lengths, forward, collate_fn=self._collate_videos ) - 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 = embeddings - transformed_modality.data_type = np.float32 - return transformed_modality + transformed = {} + for name in layer_names: + transformed_modality = TransformedModality( + modality, self, self.output_modality_type + ) + transformed_modality.data = embeddings.pop(name) + transformed_modality.data_type = np.float32 + transformed[name] = transformed_modality + return transformed + + +def _pooled_features(output) -> torch.Tensor: + if isinstance(output, tuple): + output = output[0] + if output.ndim > 2: + output = torch.nn.functional.adaptive_avg_pool2d(output, (1, 1)) + return torch.flatten(output, 1) class I3D(FrameBudgetedVideoModel, UnimodalRepresentation): From 3cfa8c469041f3f416204d58dc3a02b9385c4e93 Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Tue, 13 Oct 2026 13:47:53 +0200 Subject: [PATCH 11/26] imagbind licensing --- .../scuro/representations/image_bind.py | 204 ++++++++++++------ 1 file changed, 135 insertions(+), 69 deletions(-) diff --git a/src/main/python/systemds/scuro/representations/image_bind.py b/src/main/python/systemds/scuro/representations/image_bind.py index d0e366778b3..9d33fd352e7 100644 --- a/src/main/python/systemds/scuro/representations/image_bind.py +++ b/src/main/python/systemds/scuro/representations/image_bind.py @@ -18,15 +18,15 @@ # under the License. # # ------------------------------------------------------------- +import html import math import os +from fractions import Fraction 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 systemds.scuro.drsearch.operator_registry import ( register_representation, @@ -53,13 +53,34 @@ [ModalityType.VIDEO, ModalityType.AUDIO, ModalityType.TEXT] ) class ImageBind(UnimodalRepresentation): + """ImageBind (Girdhar et al., CVPR 2023) embeddings for video, audio and text. + + Input preparation is implemented here; only the pretrained network comes + from the optional ``imagebind`` package. Its code and weights are licensed + CC BY-NC-SA 4.0, so this representation is restricted to non-commercial use. + """ + _EMBEDDING_DIM = 1024 _MODEL_PARAMETER_COUNT = 1_200_000_000 + _CLIP_SECONDS = 2 _CLIPS_PER_VIDEO = 5 _SPATIAL_CROPS = 3 _FRAMES_PER_CLIP = 2 _CROP_SIZE = 224 + # OpenAI CLIP pixel statistics, which ImageBind's vision encoder inherits. + _PIXEL_MEAN = (0.48145466, 0.4578275, 0.40821073) + _PIXEL_STD = (0.26862954, 0.26130258, 0.27577711) + _AUDIO_SAMPLE_RATE = 16000 + _CLIPS_PER_AUDIO = 3 + _MEL_BINS = 128 + _MEL_FRAMES = 204 + _MEL_MEAN = -4.268 + _MEL_STD = 9.138 + _TEXT_CONTEXT_LENGTH = 77 + # ImageBind's text encoder uses the OpenAI CLIP BPE vocabulary. + _TOKENIZER_NAME = "openai/clip-vit-base-patch32" _CHECKPOINT_PATH = os.path.join(".checkpoints", "imagebind_huge.pth") + _tokenizer = None supports_aggregation_pushdown = True cache_in_worker = True @@ -139,9 +160,8 @@ def estimate_peak_memory_bytes(self, input_stats) -> dict: return {"cpu_peak_bytes": int(cpu_peak), "gpu_peak_bytes": int(gpu_peak)} def _ensure_model(self): - global data, imagebind_model, IBModalityType + global 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: @@ -179,24 +199,46 @@ def _sampling_rate(metadata, modality_type): ) return sampling_rate - def _transform_audio_data(self, samples, metadata): - """Adapt ImageBind audio preprocessing to Scuro-loaded waveforms. + @staticmethod + def _clip_windows(duration, num_clips, clip_seconds): + """(start, end) seconds of num_clips evenly spaced windows over duration. - The clip sampling, mel conversion, and normalization match ImageBind's - load_and_transform_audio_data; only path-based loading is replaced. + The first window starts at 0 and the last one ends with the signal; + signals shorter than a window repeat [0, clip_seconds). Fractions keep + the sample and frame boundaries free of float rounding. """ - 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) + clip_seconds = Fraction(clip_seconds) + last_start = Fraction(max(duration - clip_seconds, 0)) + stride = last_start / max(num_clips - 1, 1) + return [(i * stride, i * stride + clip_seconds) for i in range(num_clips)] + + def _log_mel(self, waveform): + """Normalized (1, mel bins, frames) Kaldi filterbank of a centred clip.""" + fbank = torchaudio.compliance.kaldi.fbank( + waveform, + sample_frequency=self._AUDIO_SAMPLE_RATE, + num_mel_bins=self._MEL_BINS, + frame_length=25, + frame_shift=10, + window_type="hanning", + dither=0.0, + use_energy=False, + htk_compat=True, + ).T + # The encoder expects a fixed number of frames: zero-pad or truncate. + missing = self._MEL_FRAMES - fbank.shape[1] + if missing > 0: + fbank = torch.nn.functional.pad(fbank, (0, missing)) + fbank = fbank[:, : self._MEL_FRAMES] + return ((fbank - self._MEL_MEAN) / self._MEL_STD).unsqueeze(0) + + def _transform_audio_data(self, samples, metadata): + """(batch, clips, 1, mel bins, frames) log-mel inputs for the audio encoder.""" + sample_rate = self._AUDIO_SAMPLE_RATE audio_outputs = [] 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. + # Copy, since clips are centred in place below. waveform = torch.as_tensor(np.asarray(sample)).clone().float() if waveform.ndim == 1: waveform = waveform.unsqueeze(0) @@ -214,52 +256,46 @@ def _transform_audio_data(self, samples, metadata): waveform, orig_freq=original_rate, new_freq=sample_rate ) - timepoints = data.get_clip_timepoints( - clip_sampler, waveform.size(1) / sample_rate + windows = self._clip_windows( + waveform.size(1) / sample_rate, + self._CLIPS_PER_AUDIO, + self._CLIP_SECONDS, ) 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)) + for start, end in windows: + clip = waveform[:, int(start * sample_rate) : int(end * sample_rate)] + # Centring the view in place means overlapping windows of short + # audio see already centred samples. The released ImageBind + # pipeline behaves the same way, so embeddings stay comparable. + clip -= clip.mean() + clips.append(self._log_mel(clip)) 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. + def _three_crops(self, clip): + """Square crops at the start, centre and end of a clip's longer side.""" + height, width = clip.shape[-2:] + size = self._CROP_SIZE + top = math.ceil((height - size) / 2) + left = math.ceil((width - size) / 2) + if height > width: + offsets = [(0, left), (top, left), (height - size, left)] + else: + offsets = [(top, 0), (top, left), (top, width - size)] + return [clip[..., y : y + size, x : x + size] for y, x in offsets] - 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), - ), - ] - ) + def _transform_video_data(self, samples, metadata): + """(batch, clips * crops, C, T, H, W) inputs for the vision encoder.""" + subsample = pv_transforms.UniformTemporalSubsample(self._FRAMES_PER_CLIP) + rescale = pv_transforms.ShortSideScale(self._CROP_SIZE) + mean = torch.tensor(self._PIXEL_MEAN)[:, None, None, None] + std = torch.tensor(self._PIXEL_STD)[:, None, None, None] 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. + # Scuro stores decoded video as (T, H, W, C), already scaled to + # [0, 1] when it has a float dtype; the encoder wants (C, T, H, W). frames = np.asarray(sample) if frames.ndim != 4 or frames.shape[-1] != 3: raise ValueError( @@ -272,30 +308,60 @@ def _transform_video_data(self, samples, metadata): 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 + windows = self._clip_windows( + video.shape[1] / sampling_rate, + self._CLIPS_PER_VIDEO, + self._CLIP_SECONDS, ) - clips = [] - for start, end in timepoints: - # Match the ceil-based [start, end) indexing used by the - # Decord-backed EncodedVideo loader in ImageBind. + crops = [] + for start, end in windows: + # Frames whose timestamps fall in [start, end). 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)) + clip = rescale(subsample(video[:, start_frame:end_frame])) + crops.extend(self._three_crops((clip - mean) / std)) + video_outputs.append(torch.stack(crops)) return torch.stack(video_outputs).to(self.device) + @classmethod + def _get_tokenizer(cls): + if cls._tokenizer is None: + from transformers import CLIPTokenizer + + cls._tokenizer = CLIPTokenizer.from_pretrained(cls._TOKENIZER_NAME) + return cls._tokenizer + + @staticmethod + def _clean_text(text): + try: + import ftfy + + text = ftfy.fix_text(text) + except ImportError: + pass + return html.unescape(html.unescape(text)).strip() + + def _tokenize(self, texts): + """(len(texts), 77) token ids: start, up to 75 text tokens, end, zeros.""" + tokenizer = self._get_tokenizer() + text_ids = tokenizer( + [self._clean_text(text) for text in texts], + add_special_tokens=False, + truncation=True, + max_length=self._TEXT_CONTEXT_LENGTH - 2, + )["input_ids"] + tokens = torch.zeros( + (len(text_ids), self._TEXT_CONTEXT_LENGTH), dtype=torch.long + ) + for row, ids in enumerate(text_ids): + ids = [tokenizer.bos_token_id, *ids, tokenizer.eos_token_id] + tokens[row, : len(ids)] = torch.tensor(ids) + return tokens.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 - ) + return IBModalityType.TEXT, self._tokenize(samples) if modality_type == ModalityType.AUDIO: return IBModalityType.AUDIO, self._transform_audio_data(samples, metadata) if modality_type == ModalityType.VIDEO: From c947a079dce7f243026faf5b45cffa54695efb03 Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Tue, 13 Oct 2026 20:21:14 +0200 Subject: [PATCH 12/26] more efficient audio loading --- .../systemds/scuro/dataloader/audio_loader.py | 42 ++++++++++++++----- .../systemds/scuro/representations/mfcc.py | 4 +- 2 files changed, 34 insertions(+), 12 deletions(-) diff --git a/src/main/python/systemds/scuro/dataloader/audio_loader.py b/src/main/python/systemds/scuro/dataloader/audio_loader.py index b5ddcb07677..f24f6e53b5c 100644 --- a/src/main/python/systemds/scuro/dataloader/audio_loader.py +++ b/src/main/python/systemds/scuro/dataloader/audio_loader.py @@ -20,12 +20,11 @@ # ------------------------------------------------------------- from dataclasses import dataclass import math -import os from typing import List, Optional, Union import librosa import numpy as np -from systemds.scuro.dataloader.base_loader import BaseLoader +from systemds.scuro.dataloader.base_loader import BaseLoader, LazyFileSequence from systemds.scuro.modality.type import ModalityType @@ -51,34 +50,55 @@ def __init__( chunk_size: Optional[int] = None, normalize: bool = True, load=True, + sampling_rate: Optional[int] = None, ): super().__init__( source_path, indices, data_type, chunk_size, ModalityType.AUDIO ) self.normalize = normalize self.load_data_from_file = load + self.sampling_rate = sampling_rate files = self.get_file_names(self.indices) self.stats = self.cached_stats( files, - {"dtype": str(data_type), "normalize": normalize}, + { + "dtype": str(data_type), + "normalize": normalize, + "sampling_rate": sampling_rate, + }, lambda: self.get_stats(source_path), ) - def extract(self, file: str, index: Optional[Union[str, List[str]]] = None): + def _load(self, indices: List[str]): + if self.data is not None and len(self.data) == len(indices): + return self.data, self.metadata + + file_names = self.get_file_names(indices) + if isinstance(file_names, str): + file_names = [file_names] + decoded = LazyFileSequence(file_names, self._decode_file).materialize() + self.data = [audio for audio, _ in decoded] + self.metadata = [metadata for _, metadata in decoded] + return self.data, self.metadata + + def _decode_file(self, file: str): self.file_sanity_check(file) if not self.load_data_from_file: - import numpy as np - audio = np.array([0]) sr = 1000 else: - audio, sr = librosa.load(file, dtype=self._data_type) + audio, sr = librosa.load( + file, sr=self.sampling_rate, dtype=self._data_type + ) if self.normalize: audio = librosa.util.normalize(audio) - self.metadata.append(self.modality_type.create_metadata(sr, audio)) + return audio, self.modality_type.create_metadata(sr, audio) + def extract(self, file: str, index: Optional[Union[str, List[str]]] = None): + audio, metadata = self._decode_file(file) + self.metadata.append(metadata) self.data.append(audio) def get_stats(self, source_path: str): @@ -92,10 +112,12 @@ def get_stats(self, source_path: str): self.file_sanity_check(file) try: info = soundfile.info(file) - sr = 22050 + sr = self.sampling_rate or info.samplerate length = math.ceil(info.frames * sr / info.samplerate) except (RuntimeError, TypeError): - audio, sr = librosa.load(file, dtype=self._data_type) + audio, sr = librosa.load( + file, sr=self.sampling_rate, dtype=self._data_type + ) length = audio.shape[0] num_instances += 1 sampling_rate = max(sampling_rate, sr) diff --git a/src/main/python/systemds/scuro/representations/mfcc.py b/src/main/python/systemds/scuro/representations/mfcc.py index 2994a37efeb..f97e51dbc53 100644 --- a/src/main/python/systemds/scuro/representations/mfcc.py +++ b/src/main/python/systemds/scuro/representations/mfcc.py @@ -79,8 +79,8 @@ def transform(self, modality, aggregation=None): modality, self, self.output_modality_type ) result = [] - sr = modality.metadata[0]["frequency"] if modality.metadata else 22050 - for sample in modality.data: + for i, sample in enumerate(modality.data): + sr = modality.metadata[i]["frequency"] if modality.metadata else None computed_feature = self.compute_feature(sample, sr=sr) result.append(computed_feature) transformed_modality.data = result From b631fecddcc4244d8ca783226f084445671fa845 Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Tue, 13 Oct 2026 20:22:01 +0200 Subject: [PATCH 13/26] add chunk cache for text modality --- .../systemds/scuro/representations/bert.py | 240 ++++++++++++++---- .../systemds/scuro/representations/utils.py | 94 +++++++ 2 files changed, 278 insertions(+), 56 deletions(-) diff --git a/src/main/python/systemds/scuro/representations/bert.py b/src/main/python/systemds/scuro/representations/bert.py index 02382938615..748426835ac 100644 --- a/src/main/python/systemds/scuro/representations/bert.py +++ b/src/main/python/systemds/scuro/representations/bert.py @@ -18,6 +18,7 @@ # under the License. # # ------------------------------------------------------------- +from contextlib import contextmanager from dataclasses import dataclass import numpy as np from systemds.scuro.modality.transformed import TransformedModality @@ -29,7 +30,9 @@ LengthBucketBatchSampler, OwnerAccumulator, OwnedSequenceDataset, + canonical_batches, capture_layer_outputs, + chunk_embedding_cache, existing_layers, flatten_owned_sequences, move_batch_to_device, @@ -52,6 +55,9 @@ os.environ["TOKENIZERS_PARALLELISM"] = "false" +# Padded lengths of cached chunk batches are multiples of this many tokens. +CHUNK_LENGTH_STEP = 16 + class BertFamily(UnimodalRepresentation): supports_aggregation_pushdown = True @@ -318,10 +324,10 @@ def _embeddings_by_layer( num_owners=None, grouped=False, ): - """Encode every text once and pool each requested layer's output. + """Encode the texts once and pool each requested layer's output. - The layer name "cls" is the CLS token of the last hidden state; any - other name is a submodule whose output is mean-pooled over the tokens. + With SCURO_TEXT_CHUNK_CACHE=1, texts already encoded by this process + for these layers are taken from the chunk cache instead. """ texts = list(TextDataset(data)) single_owner = owner_ids is None and aggregation is not None @@ -329,50 +335,44 @@ def _embeddings_by_layer( 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) + accumulators = { + name: OwnerAccumulator(num_owners, len(texts), aggregation) + for name in layer_names + } - length_encoding = tokenizer( - texts, - padding=False, - truncation=True, - max_length=self.max_seq_length, - return_attention_mask=True, - ) - attention_mask = length_encoding.get("attention_mask") - if isinstance(attention_mask, torch.Tensor): - lengths = attention_mask.sum(dim=1).tolist() + cache = chunk_embedding_cache() + encoder_key = self._chunk_cache_key(model) if cache is not None else None + if encoder_key is None: + self._encode(texts, owner_ids, model, tokenizer, layer_names, accumulators) else: - lengths = [sum(mask) for mask in attention_mask] - - def collate(samples): - batch_texts, batch_owner_ids, chunk_ids = zip(*samples) - inputs = tokenizer( - list(batch_texts), - return_offsets_mapping=True, - return_tensors="pt", - padding=True, - return_attention_mask=True, - truncation=True, - 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), + self._encode_cached( + texts, + owner_ids, + model, + tokenizer, + layer_names, + accumulators, + cache, + encoder_key, ) - dataloader = DataLoader( - dataset, - batch_sampler=LengthBucketBatchSampler(lengths, self.batch_size), - collate_fn=collate, - pin_memory=pin_memory_for(self.device), - ) - accumulators = { - name: OwnerAccumulator(num_owners, len(dataset), aggregation) - for name in layer_names - } + embeddings = {} + for name in layer_names: + layer_embeddings = accumulators.pop(name).finalize(grouped=grouped) + if single_owner: + layer_embeddings = layer_embeddings[0] + elif aggregation is None and not grouped: + layer_embeddings = list(layer_embeddings) + embeddings[name] = layer_embeddings + return embeddings + + @contextmanager + def _layer_pass(self, model, layer_names): + """Yield run(inputs) -> {layer: pooled vectors} from one forward pass. + + The layer name "cls" is the CLS token of the last hidden state; any + other name is a submodule whose output is mean-pooled over the tokens. + """ hooked = [name for name in layer_names if name != "cls"] batch_mask = {} @@ -384,8 +384,8 @@ def pool_hooked(output): with transformer_inference_context(self.device), capture_layer_outputs( model, hooked, pool_hooked ) as outputs: - for inputs, batch_owner_ids, chunk_ids in dataloader: - inputs = move_batch_to_device(inputs, self.device) + + def run(inputs): batch_mask["attention_mask"] = inputs["attention_mask"] model_outputs = model(**inputs) if "cls" in layer_names: @@ -394,20 +394,148 @@ def pool_hooked(output): inputs["attention_mask"], use_cls=True, ) + return {name: outputs[name] for name in layer_names} + + yield run + + def _encode(self, texts, owner_ids, model, tokenizer, layer_names, accumulators): + """Encode all texts in length-bucketed batches padded to their longest.""" + encoding = tokenizer( + texts, + padding=False, + truncation=True, + max_length=self.max_seq_length, + return_attention_mask=True, + ) + features = [ + {key: values[i] for key, values in encoding.items()} + for i in range(len(texts)) + ] + lengths = [len(ids) for ids in encoding["input_ids"]] + dataset = OwnedSequenceDataset(features, owner_ids) + + def collate(samples): + batch_features, batch_owner_ids, chunk_ids = zip(*samples) + inputs = tokenizer.pad( + list(batch_features), padding=True, return_tensors="pt" + ) + return ( + dict(inputs), + torch.tensor(batch_owner_ids, dtype=torch.long), + torch.tensor(chunk_ids, dtype=torch.long), + ) + + dataloader = DataLoader( + dataset, + batch_sampler=LengthBucketBatchSampler(lengths, self.batch_size), + collate_fn=collate, + pin_memory=pin_memory_for(self.device), + ) + with self._layer_pass(model, layer_names) as run: + for inputs, batch_owner_ids, chunk_ids in dataloader: + pooled = run(move_batch_to_device(inputs, self.device)) for name in layer_names: - accumulators[name].update( - outputs[name], batch_owner_ids, chunk_ids - ) + accumulators[name].update(pooled[name], batch_owner_ids, chunk_ids) - embeddings = {} + def _chunk_cache_key(self, model): + """What a cached chunk vector depends on besides its layer and text. + + Only a model loaded by name is identified; others are not cached. + """ + model_name = getattr(getattr(model, "config", None), "_name_or_path", "") + if not model_name: + return None + device = torch.device(self.device) + if device.type == "cuda": + hardware = ( + torch.cuda.get_device_name(device), + torch.cuda.get_device_capability(device), + ) + else: + hardware = (device.type,) + return ( + type(self).__name__, + model_name, + self.max_seq_length, + self.batch_size, + CHUNK_LENGTH_STEP, + *hardware, + ) + + def _encode_cached( + self, + texts, + owner_ids, + model, + tokenizer, + layer_names, + accumulators, + cache, + encoder_key, + ): + """Encode only the texts not cached for every layer, then gather all. + + Texts are encoded in fixed-shape batches, so a vector does not depend + on which texts were missing alongside it: a cached vector is exactly + what encoding the text again would give. + """ + vectors = {} + missing = [] + for text in dict.fromkeys(texts): + cached = [cache.get((encoder_key, name, text)) for name in layer_names] + if any(vector is None for vector in cached): + missing.append(text) + else: + for name, vector in zip(layer_names, cached): + vectors[name, text] = vector + + if missing: + encoding = tokenizer( + missing, + padding=False, + truncation=True, + max_length=self.max_seq_length, + return_attention_mask=True, + ) + # Each batch ends with an empty text, so its attention mask always + # has padding: without any, transformers drops the mask and SDPA + # switches kernels, which changes the other rows' numerics. + padding_row = dict( + tokenizer("", truncation=True, max_length=self.max_seq_length) + ) + batches = canonical_batches( + [len(ids) for ids in encoding["input_ids"]], + max(1, self.batch_size - 1), + CHUNK_LENGTH_STEP, + self.max_seq_length, + ) + with self._layer_pass(model, layer_names) as run: + for padded_length, rows, num_real in batches: + inputs = tokenizer.pad( + [ + {key: values[i] for key, values in encoding.items()} + for i in rows + ] + + [padding_row], + padding="max_length", + max_length=padded_length, + return_tensors="pt", + ) + pooled = run(move_batch_to_device(dict(inputs), self.device)) + for name in layer_names: + batch_vectors = pooled[name][:num_real].float().cpu() + for row, index in enumerate(rows[:num_real]): + vector = batch_vectors[row].clone() + vectors[name, missing[index]] = vector + cache.put((encoder_key, name, missing[index]), vector) + + if not texts: + return + owners = torch.tensor(list(owner_ids), dtype=torch.long) + chunk_ids = torch.arange(len(texts), dtype=torch.long) for name in layer_names: - layer_embeddings = accumulators.pop(name).finalize(grouped=grouped) - if single_owner: - layer_embeddings = layer_embeddings[0] - elif aggregation is None and not grouped: - layer_embeddings = list(layer_embeddings) - embeddings[name] = layer_embeddings - return embeddings + stacked = torch.stack([vectors[name, text] for text in texts]) + accumulators[name].update(stacked.to(self.device), owners, chunk_ids) def _hidden_state(output): diff --git a/src/main/python/systemds/scuro/representations/utils.py b/src/main/python/systemds/scuro/representations/utils.py index 76d339256ad..9ae5129eb15 100644 --- a/src/main/python/systemds/scuro/representations/utils.py +++ b/src/main/python/systemds/scuro/representations/utils.py @@ -20,7 +20,9 @@ # ------------------------------------------------------------- import os import pickle +import threading from bisect import bisect_right +from collections import OrderedDict from collections.abc import Sequence from contextlib import contextmanager @@ -463,6 +465,98 @@ def finalize(self, grouped=False): return [values[owner_ids == owner] for owner in range(self.num_owners)] +class ChunkEmbeddingCache: + """Process-wide LRU of pooled text-chunk vectors, bounded in bytes. + + Text-context configurations split the same texts into largely the same + chunks, so across a search the encoder sees most chunks many times. A + vector is only reusable if it does not depend on the batch it was encoded + in; under mixed precision it does depend on the batch shape, so encoders + that use this cache encode in fixed-shape batches (see ``canonical_batches``). + """ + + def __init__(self, max_bytes): + self.max_bytes = int(max_bytes) + self._entries = OrderedDict() + self._bytes = 0 + self._lock = threading.Lock() + + def get(self, key): + with self._lock: + value = self._entries.get(key) + if value is not None: + self._entries.move_to_end(key) + return value + + def put(self, key, value): + size = value.element_size() * value.nelement() + if size > self.max_bytes: + return + with self._lock: + previous = self._entries.pop(key, None) + if previous is not None: + self._bytes -= previous.element_size() * previous.nelement() + self._entries[key] = value + self._bytes += size + while self._bytes > self.max_bytes: + _, evicted = self._entries.popitem(last=False) + self._bytes -= evicted.element_size() * evicted.nelement() + + def clear(self): + with self._lock: + self._entries.clear() + self._bytes = 0 + + def __len__(self): + return len(self._entries) + + +_CHUNK_EMBEDDING_CACHE = None +_CHUNK_EMBEDDING_CACHE_LOCK = threading.Lock() + + +def chunk_embedding_cache(): + """The process's chunk cache, or None unless SCURO_TEXT_CHUNK_CACHE=1. + + SCURO_TEXT_CHUNK_CACHE_MB bounds it (default 512 MB per process). + """ + global _CHUNK_EMBEDDING_CACHE + if os.environ.get("SCURO_TEXT_CHUNK_CACHE", "0") != "1": + return None + with _CHUNK_EMBEDDING_CACHE_LOCK: + if _CHUNK_EMBEDDING_CACHE is None: + max_mb = float(os.environ.get("SCURO_TEXT_CHUNK_CACHE_MB", "512")) + _CHUNK_EMBEDDING_CACHE = ChunkEmbeddingCache(max_mb * 1024 * 1024) + return _CHUNK_EMBEDDING_CACHE + + +def canonical_batches(lengths, batch_size, length_step, max_length=None): + """Group sequences into batches of exactly ``batch_size`` rows each. + + Each batch holds sequences of one padded length, ``lengths`` rounded up to + a multiple of ``length_step`` (at most ``max_length``). A partial batch is + filled up by repeating + its last index. Returns [(padded length, indices, number of real rows)]. + + With every batch of the same shape, a sequence's output depends only on + its own tokens and padded length, not on which others share its batch. + """ + buckets = {} + for index, length in sorted(enumerate(lengths), key=lambda item: item[1]): + padded = max(1, -(-int(length) // length_step)) * length_step + if max_length is not None: + padded = min(padded, int(max_length)) + buckets.setdefault(padded, []).append(index) + batches = [] + for padded, indices in sorted(buckets.items()): + for start in range(0, len(indices), batch_size): + rows = indices[start : start + batch_size] + batches.append( + (padded, rows + [rows[-1]] * (batch_size - len(rows)), len(rows)) + ) + return batches + + def dense_instance_batch(data): if isinstance(data, np.ndarray): if data.ndim == 2 and np.issubdtype(data.dtype, np.number): From 6983aba2f016089aa51babd266e346eac377995f Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Wed, 14 Oct 2026 06:57:34 +0200 Subject: [PATCH 14/26] improve audio reps --- .../systemds/scuro/drsearch/process_cache.py | 60 +++++++++- .../scuro/drsearch/representation_dag.py | 4 + .../scuro/representations/mel_spectrogram.py | 79 ++++++++++---- .../systemds/scuro/representations/mfcc.py | 45 ++++---- .../scuro/representations/spectral.py | 103 ++++++++++++++++++ .../scuro/representations/spectrogram.py | 18 +-- .../scuro/representations/unimodal.py | 4 + .../systemds/scuro/representations/wav2vec.py | 99 +++++++++++++---- .../representations/window_aggregation.py | 5 +- 9 files changed, 343 insertions(+), 74 deletions(-) create mode 100644 src/main/python/systemds/scuro/representations/spectral.py diff --git a/src/main/python/systemds/scuro/drsearch/process_cache.py b/src/main/python/systemds/scuro/drsearch/process_cache.py index 357190df0ef..558d6acc111 100644 --- a/src/main/python/systemds/scuro/drsearch/process_cache.py +++ b/src/main/python/systemds/scuro/drsearch/process_cache.py @@ -110,9 +110,20 @@ def data_fingerprint(*values: Any) -> str: class BoundedProcessCache: - def __init__(self, max_entries: int): + def __init__( + self, + max_entries: int, + max_bytes: Optional[int] = None, + sizeof: Optional[Callable[[Any], int]] = None, + ): + """LRU cache holding at most ``max_entries`` values and, if ``max_bytes`` + is set, at most that many bytes as measured by ``sizeof``.""" self.max_entries = max(0, int(max_entries)) + self.max_bytes = None if max_bytes is None else max(0, int(max_bytes)) + self.sizeof = sizeof or nbytes self._values = OrderedDict() + self._sizes = {} + self._total_bytes = 0 self._pending = {} self._lock = threading.RLock() @@ -127,11 +138,22 @@ def get(self, key: Hashable, default=None): def put(self, key: Hashable, value: Any) -> None: if self.max_entries == 0: return + size = self.sizeof(value) if self.max_bytes is not None else 0 with self._lock: - self._values.pop(key, None) + self._discard(key) + if self.max_bytes is not None and size > self.max_bytes: + return self._values[key] = value - while len(self._values) > self.max_entries: - self._values.popitem(last=False) + self._sizes[key] = size + self._total_bytes += size + while len(self._values) > self.max_entries or ( + self.max_bytes is not None and self._total_bytes > self.max_bytes + ): + self._discard(next(iter(self._values))) + + def _discard(self, key: Hashable) -> None: + if self._values.pop(key, _MISSING) is not _MISSING: + self._total_bytes -= self._sizes.pop(key) def get_or_compute(self, key: Hashable, compute: Callable[[], Any]) -> Any: if self.max_entries == 0: @@ -208,19 +230,48 @@ def get_or_compute_group( def clear(self) -> None: with self._lock: self._values.clear() + self._sizes.clear() + self._total_bytes = 0 def __len__(self) -> int: with self._lock: return len(self._values) +def nbytes(value: Any) -> int: + """Bytes held by the arrays in ``value`` (nested lists, tuples and dicts).""" + if isinstance(value, np.ndarray): + if value.dtype.hasobject: + return value.nbytes + sum(nbytes(item) for item in value.flat) + return value.nbytes + if isinstance(value, (list, tuple)): + return sum(nbytes(item) for item in value) + if isinstance(value, dict): + return sum(nbytes(item) for item in value.values()) + return int(getattr(value, "nbytes", 0)) + + +def _intermediate_cache_bytes() -> int: + return int(float(os.environ.get("SCURO_HPO_INTERMEDIATE_CACHE_MB", "2048")) * 2**20) + + class HPOExecutionCache: def __init__( self, representation_entries: int = 128, should_cache_representation: Optional[Callable[[Any], bool]] = None, + intermediate_bytes: Optional[int] = None, ): self.representations = BoundedProcessCache(representation_entries) + # Results that representations with different parameters share, such as + # the mel spectrogram under every MFCC n_mfcc/dct_type. A representation + # reads it through its ``intermediate_cache`` attribute and keys entries + # by the parameters and data they depend on. + if intermediate_bytes is None: + intermediate_bytes = _intermediate_cache_bytes() + self.intermediates = BoundedProcessCache( + 4096 if intermediate_bytes > 0 else 0, max_bytes=intermediate_bytes + ) self.should_cache_representation = should_cache_representation self._modality_keys = {} self._layer_sweeps = {} @@ -379,6 +430,7 @@ def compute_group(): def clear(self) -> None: self.representations.clear() + self.intermediates.clear() with self._lock: self._modality_keys.clear() self._layer_sweeps.clear() diff --git a/src/main/python/systemds/scuro/drsearch/representation_dag.py b/src/main/python/systemds/scuro/drsearch/representation_dag.py index 860b075561e..d400ccbb199 100644 --- a/src/main/python/systemds/scuro/drsearch/representation_dag.py +++ b/src/main/python/systemds/scuro/drsearch/representation_dag.py @@ -409,6 +409,10 @@ def run_operation(): node_operation = node.operation(params=node.parameters) if gpu_id is not None and hasattr(node_operation, "gpu_id"): node_operation.gpu_id = gpu_id + if execution_cache is not None and hasattr( + node_operation, "intermediate_cache" + ): + node_operation.intermediate_cache = execution_cache.intermediates if len(input_mods) == 1: # It's a unimodal operation if isinstance(node_operation, Context): diff --git a/src/main/python/systemds/scuro/representations/mel_spectrogram.py b/src/main/python/systemds/scuro/representations/mel_spectrogram.py index 5683b8f2117..663174a23b4 100644 --- a/src/main/python/systemds/scuro/representations/mel_spectrogram.py +++ b/src/main/python/systemds/scuro/representations/mel_spectrogram.py @@ -18,18 +18,19 @@ # under the License. # # ------------------------------------------------------------- -import librosa import numpy as np 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.spectral import mel_spectrogram from systemds.scuro.representations.unimodal import UnimodalRepresentation from systemds.scuro.drsearch.operator_registry import ( register_representation, register_context_representation_operator, ) +from systemds.scuro.drsearch.process_cache import data_fingerprint from systemds.scuro.utils.static_variables import ( NP_ARRAY_HEADER_BYTES, PY_LIST_HEADER_BYTES, @@ -45,6 +46,51 @@ ) +def power_mel_spectrogram(instance, sr, n_mels, hop_length, n_fft): + """Power mel spectrogram of one signal, shaped (..., n_mels, frames).""" + S = mel_spectrogram( + np.asarray(instance, dtype=np.float32), + 22050 if sr is None else sr, + n_fft, + hop_length, + n_mels, + ) + return np.swapaxes(S, -2, -1) + + +def mel_spectrograms(modality, n_mels, hop_length, n_fft, cache=None): + """Power mel spectrogram of every instance of ``modality``. + + MFCC and MelSpectrogram configurations with equal (n_fft, hop_length, + n_mels) need the same spectrograms, and they are most of an MFCC's cost. + With ``cache`` (a BoundedProcessCache) they are computed once per data + chunk and reused by every such configuration. Callers must not modify + the returned arrays. + """ + samples = modality.data + srs = [ + modality.metadata[i]["frequency"] if modality.metadata else None + for i in range(len(samples)) + ] + + def compute(): + return [ + power_mel_spectrogram(sample, sr, n_mels, hop_length, n_fft) + for sample, sr in zip(samples, srs) + ] + + if cache is None: + return compute() + key = ( + "mel_spectrogram", + int(n_fft), + int(hop_length), + int(n_mels), + data_fingerprint(samples, srs), + ) + return cache.get_or_compute(key, compute) + + @register_representation(ModalityType.AUDIO) @register_context_representation_operator(ModalityType.AUDIO) class MelSpectrogram(UnimodalRepresentation): @@ -70,29 +116,24 @@ 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"] - computed_feature = self.compute_feature(sample, sr) - result.append(computed_feature) - + spectrograms = mel_spectrograms( + modality, + self.n_mels, + self.hop_length, + self.n_fft, + self.intermediate_cache, + ) + result = [np.swapaxes(S, -2, -1) for S in spectrograms] + if self.intermediate_cache is not None: + # Do not hand out views of cached spectrograms. + result = [S.copy() for S in result] transformed_modality.data = result return transformed_modality def compute_feature(self, instance, sr=None): - if sr is None: - sr = 22050 - S = librosa.feature.melspectrogram( - y=np.asarray(instance), - sr=sr, - n_mels=self.n_mels, - hop_length=self.hop_length, - n_fft=self.n_fft, + S = power_mel_spectrogram( + instance, sr, self.n_mels, self.hop_length, self.n_fft ) - if instance.ndim == 1: - return S.T - return np.swapaxes(S, -2, -1) def compute_features_batched(self, data, sr=None): diff --git a/src/main/python/systemds/scuro/representations/mfcc.py b/src/main/python/systemds/scuro/representations/mfcc.py index f97e51dbc53..3e93de8d970 100644 --- a/src/main/python/systemds/scuro/representations/mfcc.py +++ b/src/main/python/systemds/scuro/representations/mfcc.py @@ -26,6 +26,11 @@ from systemds.scuro.representations.representation import RepresentationStats from systemds.scuro.representations.unimodal import UnimodalRepresentation +from systemds.scuro.representations.mel_spectrogram import ( + mel_spectrograms, + power_mel_spectrogram, +) +from systemds.scuro.representations.spectral import power_to_db from systemds.scuro.drsearch.operator_registry import ( register_representation, register_context_representation_operator, @@ -78,25 +83,31 @@ 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"] if modality.metadata else None - computed_feature = self.compute_feature(sample, sr=sr) - result.append(computed_feature) - transformed_modality.data = result + spectrograms = mel_spectrograms( + modality, + self.n_mels, + self.hop_length, + self.n_fft, + self.intermediate_cache, + ) + transformed_modality.data = [self.mfcc_from_mel(S) for S in spectrograms] return transformed_modality def compute_feature(self, instance, sr=None): - if sr is None: - sr = 22050 + return self.mfcc_from_mel( + power_mel_spectrogram( + instance, sr, self.n_mels, self.hop_length, self.n_fft + ) + ) + + def mfcc_from_mel(self, S): + """Normalized MFCCs, frames first, from a power mel spectrogram S. + + Equivalent to librosa.feature.mfcc on the signal; S is not modified. + A batch of spectrograms is processed independently per spectrogram. + """ mfcc = librosa.feature.mfcc( - y=np.asarray(instance, dtype=np.float32), - sr=sr, - n_mfcc=self.n_mfcc, - dct_type=self.dct_type, - hop_length=self.hop_length, - n_mels=self.n_mels, - n_fft=self.n_fft, + S=power_to_db(S), n_mfcc=self.n_mfcc, dct_type=self.dct_type ) if mfcc.ndim == 2: mean = mfcc.mean(keepdims=True) @@ -106,10 +117,6 @@ def compute_feature(self, instance, sr=None): std = mfcc.std(axis=(-2, -1), keepdims=True) mfcc -= mean mfcc /= np.maximum(std, 1e-8) - - if instance.ndim == 1: - return mfcc.T - return np.swapaxes(mfcc, -2, -1) def compute_features_batched(self, data, sr=None): diff --git a/src/main/python/systemds/scuro/representations/spectral.py b/src/main/python/systemds/scuro/representations/spectral.py new file mode 100644 index 00000000000..648907ed3f3 --- /dev/null +++ b/src/main/python/systemds/scuro/representations/spectral.py @@ -0,0 +1,103 @@ +# ------------------------------------------------------------- +# +# 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. +# +# ------------------------------------------------------------- +"""Vectorized STFT helpers shared by the audio representations. + +librosa rebuilds the window function and the mel filterbank on every call, +and its stft has a large fixed per-call cost. Context operators call the +audio representations once per batch of windows, so that overhead dominated +their runtime. These helpers cache the window and filterbank and transform +any leading batch shape with one FFT call. Results match librosa's defaults +(Hann window, center=True, zero padding) up to float rounding. +""" + +import functools + +import librosa +import numpy as np +import scipy.fft +from numpy.lib.stride_tricks import sliding_window_view + +# Bound on the framed-signal elements materialized at once, so long signals +# do not need a full (frames x n_fft) copy next to the spectrogram. +_MAX_BLOCK_ELEMENTS = 1 << 22 + + +@functools.lru_cache(maxsize=None) +def hann_window(n_fft, dtype): + window = librosa.filters.get_window("hann", n_fft, fftbins=True).astype(dtype) + window.setflags(write=False) + return window + + +@functools.lru_cache(maxsize=64) +def mel_basis(sr, n_fft, n_mels): + basis = librosa.filters.mel(sr=sr, n_fft=n_fft, n_mels=n_mels) + basis.setflags(write=False) + return basis + + +def _framed_power(frames, window): + spectrum = scipy.fft.rfft(frames * window, axis=-1) + return np.square(spectrum.real) + np.square(spectrum.imag) + + +def power_spectrogram(y, n_fft, hop_length): + """|STFT|^2 over the last axis of y, shaped (..., frames, 1 + n_fft // 2).""" + y = np.asarray(y) + if not np.issubdtype(y.dtype, np.floating): + y = y.astype(np.float32) + pad = n_fft // 2 + y = np.pad(y, [(0, 0)] * (y.ndim - 1) + [(pad, pad)]) + frames = sliding_window_view(y, n_fft, axis=-1)[..., ::hop_length, :] + window = hann_window(n_fft, y.dtype) + + num_frames = frames.shape[-2] + batch_size = int(np.prod(frames.shape[:-2], dtype=np.int64)) + block = max(1, _MAX_BLOCK_ELEMENTS // max(1, batch_size * n_fft)) + if block >= num_frames: + return _framed_power(frames, window) + + power = np.empty((*frames.shape[:-1], n_fft // 2 + 1), dtype=y.dtype) + for start in range(0, num_frames, block): + power[..., start : start + block, :] = _framed_power( + frames[..., start : start + block, :], window + ) + return power + + +def mel_spectrogram(y, sr, n_fft, hop_length, n_mels): + """Power mel spectrogram over the last axis of y, shaped (..., frames, n_mels).""" + power = power_spectrogram(y, n_fft, hop_length) + return power @ mel_basis(sr, n_fft, n_mels).T + + +def power_to_db(S, amin=1e-10, top_db=80.0): + """librosa.power_to_db (ref=1.0) with top_db applied per spectrogram. + + librosa clips against the maximum of the whole array, which would make + each window's features depend on the other windows of its batch. Here + the last two axes form one spectrogram. + """ + log_spec = 10.0 * np.log10(np.maximum(amin, S)) + if top_db is not None: + peak = log_spec.max(axis=(-2, -1), keepdims=True) + np.maximum(log_spec, peak - top_db, out=log_spec) + return log_spec diff --git a/src/main/python/systemds/scuro/representations/spectrogram.py b/src/main/python/systemds/scuro/representations/spectrogram.py index e4822386d77..c8db8d34d9a 100644 --- a/src/main/python/systemds/scuro/representations/spectrogram.py +++ b/src/main/python/systemds/scuro/representations/spectrogram.py @@ -18,13 +18,13 @@ # under the License. # # ------------------------------------------------------------- -import librosa import numpy as np 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.spectral import power_spectrogram, power_to_db from systemds.scuro.representations.unimodal import UnimodalRepresentation from systemds.scuro.drsearch.operator_registry import ( register_representation, @@ -71,14 +71,18 @@ def transform(self, modality, aggregation=None): return transformed_modality def compute_feature(self, instance): - data = np.array(instance) - spectrogram = librosa.stft( - y=np.abs(data), hop_length=self.hop_length, n_fft=self.n_fft + """dB magnitude spectrogram, shaped (..., frames, 1 + n_fft // 2). + + Equal to librosa.amplitude_to_db(np.abs(librosa.stft(...))) per + signal; leading axes are a batch of independent signals. + """ + power = power_spectrogram( + np.abs(np.asarray(instance)), self.n_fft, self.hop_length ) - if data.ndim == 1: - return librosa.amplitude_to_db(np.abs(spectrogram)).T + return power_to_db(power) - return librosa.amplitude_to_db(np.abs(spectrogram)).transpose(0, 2, 1) + def compute_features_batched(self, data): + return self.compute_feature(np.asarray(data)) def estimate_peak_memory_bytes(self, input_stats) -> dict: # TODO: validate this function diff --git a/src/main/python/systemds/scuro/representations/unimodal.py b/src/main/python/systemds/scuro/representations/unimodal.py index d871cef8b8a..5c61629099b 100644 --- a/src/main/python/systemds/scuro/representations/unimodal.py +++ b/src/main/python/systemds/scuro/representations/unimodal.py @@ -28,6 +28,10 @@ class UnimodalRepresentation(Representation): # encoder that sets it implements transform_layers, so a search over that # parameter can obtain every layer from one pass over the data. layer_parameter = None + # Set by the executor (e.g. HPOExecutionCache.intermediates) for + # representations that share an intermediate result with other + # configurations; None computes it every time. + intermediate_cache = None def __init__( self, name: str, output_modality_type, parameters=None, self_contained=True diff --git a/src/main/python/systemds/scuro/representations/wav2vec.py b/src/main/python/systemds/scuro/representations/wav2vec.py index 478239a4f41..1b46d60fc36 100644 --- a/src/main/python/systemds/scuro/representations/wav2vec.py +++ b/src/main/python/systemds/scuro/representations/wav2vec.py @@ -18,10 +18,12 @@ # under the License. # # ------------------------------------------------------------- +import functools + import numpy as np from transformers import Wav2Vec2Processor, Wav2Vec2Model -import librosa import torch +import torchaudio from systemds.scuro.modality.type import ModalityType from systemds.scuro.modality.transformed import TransformedModality @@ -36,7 +38,6 @@ LengthBucketBatchSampler, OwnerAccumulator, OwnedSequenceDataset, - move_batch_to_device, pin_memory_for, pool_transformer_output, transformer_inference_context, @@ -47,6 +48,8 @@ transformers_logging.set_verbosity_error() +_TARGET_SAMPLING_RATE = 16000 + @register_representation(ModalityType.AUDIO) @register_expensive_representation(ModalityType.AUDIO) @@ -98,30 +101,22 @@ def transform(self, modality, aggregation=None): transformed_modality = TransformedModality( modality, self, self.output_modality_type ) - samples = [ - librosa.resample( - np.asarray(sample), - orig_sr=modality.metadata[owner_id]["frequency"], - target_sr=16000, - ) - for owner_id, sample in enumerate(modality.data) + # Resampling and normalization run on self.device, so the raw + # signals are what crosses to the GPU. + sampling_rates = [ + int(modality.metadata[owner_id]["frequency"]) + for owner_id in range(len(modality.data)) + ] + dataset = OwnedSequenceDataset(modality.data) + lengths = [ + _resampled_length(len(sample), sr) + for sample, sr in zip(modality.data, sampling_rates) ] - 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.from_numpy(np.asarray(a, dtype=np.float32)) for a in audio], torch.tensor(owner_ids, dtype=torch.long), torch.tensor(chunk_ids, dtype=torch.long), ) @@ -137,8 +132,10 @@ def collate(batch): 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) + for audio, owner_ids, chunk_ids in dataloader: + inputs = self._prepare_inputs( + audio, [sampling_rates[i] for i in owner_ids.tolist()] + ) features = _extract_features(model, inputs["input_values"]) attention_mask = inputs.get("attention_mask") if attention_mask is not None and hasattr( @@ -160,6 +157,38 @@ def collate(batch): transformed_modality.data_type = np.float32 return transformed_modality + def _prepare_inputs(self, audio, sampling_rates): + """Wav2Vec2Processor's padded, normalized batch, computed on self.device. + + Each signal is resampled to 16 kHz, zero-padded to the longest one + and normalized to zero mean and unit variance over its own samples. + """ + non_blocking = pin_memory_for(self.device) + waveforms = [ + _resample(waveform.to(self.device, non_blocking=non_blocking), sr) + for waveform, sr in zip(audio, sampling_rates) + ] + lengths = torch.tensor([len(w) for w in waveforms], device=self.device) + input_values = torch.nn.utils.rnn.pad_sequence(waveforms, batch_first=True) + attention_mask = ( + torch.arange(input_values.shape[1], device=self.device)[None, :] + < lengths[:, None] + ) + feature_extractor = self.processor.feature_extractor + if feature_extractor.do_normalize: + mask = attention_mask.to(input_values.dtype) + count = lengths.clamp(min=1)[:, None].to(input_values.dtype) + mean = (input_values * mask).sum(dim=1, keepdim=True) / count + var = (((input_values - mean) * mask) ** 2).sum(dim=1, keepdim=True) / count + input_values = (input_values - mean) / torch.sqrt(var + 1e-7) + input_values = input_values.masked_fill( + ~attention_mask, feature_extractor.padding_value + ) + return { + "input_values": input_values, + "attention_mask": attention_mask.long(), + } + def get_output_stats(self, input_stats) -> RepresentationStats: num_instances = getattr(input_stats, "num_instances", 0) embedding_dim = 512 @@ -195,6 +224,30 @@ def estimate_peak_memory_bytes(self, input_stats) -> dict: return {"cpu_peak_bytes": cpu_peak, "gpu_peak_bytes": gpu_peak} +def _resampled_length(length, sampling_rate): + return -(-int(length) * _TARGET_SAMPLING_RATE // int(sampling_rate)) + + +@functools.lru_cache(maxsize=32) +def _resampler(sampling_rate, device): + # librosa's former "kaiser_best" filter, the closest torchaudio match to + # the soxr_hq resampler librosa uses by default. + return torchaudio.transforms.Resample( + sampling_rate, + _TARGET_SAMPLING_RATE, + resampling_method="sinc_interp_kaiser", + lowpass_filter_width=64, + rolloff=0.9475937167399596, + beta=14.769656459379492, + ).to(device) + + +def _resample(waveform, sampling_rate): + if sampling_rate == _TARGET_SAMPLING_RATE: + return waveform + return _resampler(sampling_rate, waveform.device)(waveform) + + def _extract_features(model, input_values): """``Wav2Vec2Model(...).extract_features`` without running the encoder. diff --git a/src/main/python/systemds/scuro/representations/window_aggregation.py b/src/main/python/systemds/scuro/representations/window_aggregation.py index 60b5b467ff3..73843cca1e3 100644 --- a/src/main/python/systemds/scuro/representations/window_aggregation.py +++ b/src/main/python/systemds/scuro/representations/window_aggregation.py @@ -415,7 +415,8 @@ def execute(self, modality): target_length = max(original_lengths) padded_features = np.zeros( - (len(windowed_data), target_length, *windowed_data[0].shape[1:]) + (len(windowed_data), target_length, *windowed_data[0].shape[1:]), + dtype=np.result_type(*windowed_data), ) for i, features in enumerate(windowed_data): if padded_features.ndim == 3: @@ -435,7 +436,7 @@ def execute(self, modality): windowed_data = padded_features data_type = modality.metadata[0]["data_layout"]["type"] if data_type != "str": - windowed_data = windowed_data.astype(data_type) + windowed_data = windowed_data.astype(data_type, copy=False) self.assert_output_stats(windowed_data) return windowed_data From 214fbec0e4563c385d2d0ded8d26e0d8dc4f8c6f Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Wed, 14 Oct 2026 08:29:03 +0200 Subject: [PATCH 15/26] fix seeds --- .../scuro/drsearch/hyperparameter_tuner.py | 5 +++-- .../scuro/drsearch/multimodal_ga_optimizer.py | 4 +++- src/main/python/systemds/scuro/drsearch/task.py | 5 +++-- .../systemds/scuro/representations/lstm.py | 13 ++----------- .../scuro/representations/mlp_averaging.py | 2 +- .../multimodal_attention_fusion.py | 13 ++----------- .../systemds/scuro/utils/static_variables.py | 17 ++++++++++++++++- src/main/python/systemds/scuro/utils/utils.py | 4 +++- 8 files changed, 33 insertions(+), 30 deletions(-) diff --git a/src/main/python/systemds/scuro/drsearch/hyperparameter_tuner.py b/src/main/python/systemds/scuro/drsearch/hyperparameter_tuner.py index f99b1eb05d2..d006e2f87f9 100644 --- a/src/main/python/systemds/scuro/drsearch/hyperparameter_tuner.py +++ b/src/main/python/systemds/scuro/drsearch/hyperparameter_tuner.py @@ -19,6 +19,7 @@ # # ------------------------------------------------------------- from typing import Dict, List, Tuple, Any, Optional +from systemds.scuro.utils.static_variables import get_global_seed import inspect import os import numpy as np @@ -416,7 +417,7 @@ def __init__( checkpoint_every: Optional[int] = None, checkpoint_dir: Optional[str] = None, resume: bool = False, - random_state: int = 42, + random_state: Optional[int] = None, exhaustive_threshold: int = 256, local_search_patience: int = 3, optuna_sampler: str = "tpe", # "tpe" | "random" | "bayes" @@ -446,7 +447,7 @@ def __init__( self.logger = logging.getLogger(__name__) self.checkpoint_every = checkpoint_every self.resume = resume - self.random_state = random_state + self.random_state = get_global_seed() if random_state is None else random_state self.exhaustive_threshold = max(1, exhaustive_threshold) self.local_search_patience = max(1, local_search_patience) self._rng = random.Random(self.random_state) diff --git a/src/main/python/systemds/scuro/drsearch/multimodal_ga_optimizer.py b/src/main/python/systemds/scuro/drsearch/multimodal_ga_optimizer.py index 66d86023e58..1f591c58d79 100644 --- a/src/main/python/systemds/scuro/drsearch/multimodal_ga_optimizer.py +++ b/src/main/python/systemds/scuro/drsearch/multimodal_ga_optimizer.py @@ -34,6 +34,7 @@ from dataclasses import dataclass, field from itertools import chain from typing import Any, Dict, List, Optional, Tuple, Union +from systemds.scuro.utils.static_variables import get_global_seed import numpy as np from deap import base, tools @@ -388,7 +389,7 @@ def __init__( generations: int = 20, crossover_probability: float = 0.7, mutation_probability: float = 0.4, - random_seed: int = 42, + random_seed: Optional[int] = None, maximize_metric: bool = True, elite_size: int = 2, max_workers: int = 1, @@ -476,6 +477,7 @@ def __init__( self.hall_of_fame_size = max(1, hall_of_fame_size) self.hall_of_fame: Dict[str, List[FusionSearchResult]] = {} self._hof_fitness: Dict[str, List[Tuple[float, ...]]] = {} + random_seed = get_global_seed() if random_seed is None else random_seed self.rng = random.Random(random_seed) self._eval_counter = 0 self._current_generation = -1 diff --git a/src/main/python/systemds/scuro/drsearch/task.py b/src/main/python/systemds/scuro/drsearch/task.py index 33fe1b43fe8..2c8aa4617b2 100644 --- a/src/main/python/systemds/scuro/drsearch/task.py +++ b/src/main/python/systemds/scuro/drsearch/task.py @@ -22,6 +22,7 @@ import os import time from typing import List +from systemds.scuro.utils.static_variables import get_global_seed from systemds.scuro.models.model import Model from systemds.scuro.models.training_stats import TrainingWorkAccumulator import numpy as np @@ -95,7 +96,7 @@ def __init__( performance_measures=["accuracy"], groups=None, allow_group_overlap=False, - cv_seed=11, + cv_seed=None, ): """ Parent class for the prediction task that is performed on top of the aligned representation @@ -115,7 +116,7 @@ def __init__( self.kfold = kfold self.groups = None if groups is None else np.asarray(groups) self.allow_group_overlap = allow_group_overlap - self.cv_seed = cv_seed + self.cv_seed = get_global_seed() if cv_seed is None else cv_seed self.measure_performance = measure_performance self.inference_time = [] self.training_time = [] diff --git a/src/main/python/systemds/scuro/representations/lstm.py b/src/main/python/systemds/scuro/representations/lstm.py index 46e0842c0c1..9578b9efe80 100644 --- a/src/main/python/systemds/scuro/representations/lstm.py +++ b/src/main/python/systemds/scuro/representations/lstm.py @@ -18,14 +18,13 @@ # under the License. # # ------------------------------------------------------------- -import os -import random import torch from torch import nn from torch.utils.data import DataLoader, TensorDataset from typing import List, Dict, Any, Optional from systemds.scuro.utils.static_variables import get_device +from systemds.scuro.utils.utils import set_random_seeds import numpy as np from systemds.scuro.modality.modality import Modality @@ -97,16 +96,8 @@ def __init__( self.model_state = None self.is_multilabel = False - self._set_random_seeds() + set_random_seeds() - def _set_random_seeds(self, seed=42): - os.environ["PYTHONHASHSEED"] = str(seed) - random.seed(seed) - np.random.seed(seed) - torch.manual_seed(seed) - torch.cuda.manual_seed(seed) - torch.backends.cudnn.deterministic = True - torch.backends.cudnn.benchmark = False @staticmethod def _shape_description(modality: Modality) -> str: diff --git a/src/main/python/systemds/scuro/representations/mlp_averaging.py b/src/main/python/systemds/scuro/representations/mlp_averaging.py index 1e330b0eb25..0f3c3c4ae3f 100644 --- a/src/main/python/systemds/scuro/representations/mlp_averaging.py +++ b/src/main/python/systemds/scuro/representations/mlp_averaging.py @@ -140,7 +140,7 @@ def estimate_peak_memory_bytes(self, input_stats: RepresentationStats) -> dict: return {"cpu_peak_bytes": cpu_peak, "gpu_peak_bytes": gpu_peak} def execute(self, data): - set_random_seeds(42) + set_random_seeds() input_dim = data.shape[1] if input_dim <= self.output_dim: diff --git a/src/main/python/systemds/scuro/representations/multimodal_attention_fusion.py b/src/main/python/systemds/scuro/representations/multimodal_attention_fusion.py index b92817bc531..8f4c4db30de 100644 --- a/src/main/python/systemds/scuro/representations/multimodal_attention_fusion.py +++ b/src/main/python/systemds/scuro/representations/multimodal_attention_fusion.py @@ -19,8 +19,6 @@ # # ------------------------------------------------------------- -import os -import random import torch import torch.nn as nn import torch.nn.functional as F @@ -35,6 +33,7 @@ stats_bytes, ) from systemds.scuro.utils.static_variables import get_device +from systemds.scuro.utils.utils import set_random_seeds @register_fusion_operator() @@ -95,16 +94,8 @@ def __init__( self.model_state = None self.is_multilabel = False - self._set_random_seeds() + set_random_seeds() - def _set_random_seeds(self, seed=42): - os.environ["PYTHONHASHSEED"] = str(seed) - random.seed(seed) - np.random.seed(seed) - torch.manual_seed(seed) - torch.cuda.manual_seed(seed) - torch.backends.cudnn.deterministic = True - torch.backends.cudnn.benchmark = False def _prepare_data(self, modalities: List[Modality]): inputs = {} diff --git a/src/main/python/systemds/scuro/utils/static_variables.py b/src/main/python/systemds/scuro/utils/static_variables.py index 0da186348e7..c163496f6f3 100644 --- a/src/main/python/systemds/scuro/utils/static_variables.py +++ b/src/main/python/systemds/scuro/utils/static_variables.py @@ -18,6 +18,7 @@ # under the License. # # ------------------------------------------------------------- +import os import numpy as np import torch from typing import Tuple @@ -28,7 +29,21 @@ DEBUG = False -global_rng = np.random.default_rng(42) +# Read from the environment so spawned worker processes inherit the parent's seed. +SEED = int(os.environ.get("SCURO_SEED", 42)) +global_rng = np.random.default_rng(SEED) + + +def set_seed(seed: int): + """Set the root seed that all scuro seeds (CV folds, search, model init) default to.""" + global SEED, global_rng + SEED = seed + os.environ["SCURO_SEED"] = str(seed) + global_rng = np.random.default_rng(seed) + + +def get_global_seed() -> int: + return SEED def get_seed(): diff --git a/src/main/python/systemds/scuro/utils/utils.py b/src/main/python/systemds/scuro/utils/utils.py index fc4a5df8b52..5d8ff374ef8 100644 --- a/src/main/python/systemds/scuro/utils/utils.py +++ b/src/main/python/systemds/scuro/utils/utils.py @@ -22,9 +22,11 @@ import torch import random import numpy as np +from systemds.scuro.utils.static_variables import get_global_seed -def set_random_seeds(seed=42): +def set_random_seeds(seed=None): + seed = get_global_seed() if seed is None else seed os.environ["PYTHONHASHSEED"] = str(seed) random.seed(seed) np.random.seed(seed) From cc8e0020d6c8244b6e578a095755b8723147a075 Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Wed, 14 Oct 2026 08:46:48 +0200 Subject: [PATCH 16/26] fix lstm --- src/main/python/systemds/scuro/representations/lstm.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/main/python/systemds/scuro/representations/lstm.py b/src/main/python/systemds/scuro/representations/lstm.py index 9578b9efe80..a5c1476860e 100644 --- a/src/main/python/systemds/scuro/representations/lstm.py +++ b/src/main/python/systemds/scuro/representations/lstm.py @@ -296,8 +296,8 @@ def __init__( self.classifier = nn.Linear(hidden_dim * 2, num_classes) def forward(self, x): - lstm_out, _ = self.lstm(x) - last_output = lstm_out[:, -1, :] + _, (h_n, _) = self.lstm(x) + last_output = torch.cat((h_n[-2], h_n[-1]), dim=1) dropped = self.dropout(last_output) output = self.classifier(dropped) From 7516f2d395271da53802543b35b6b500b6d8e17f Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Wed, 14 Oct 2026 09:04:59 +0200 Subject: [PATCH 17/26] apply reduction only to vectors --- .../scuro/drsearch/unimodal_optimizer.py | 24 +++++++++++++++---- 1 file changed, 19 insertions(+), 5 deletions(-) diff --git a/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py b/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py index afa22a3f38b..bd6956e298e 100644 --- a/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py +++ b/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py @@ -533,11 +533,10 @@ def _merge_results(self, local_results): local_results.wall_clock_s.get(modality_id, 0.0) ) - def add_dimensionality_reduction_operators(self, builder, current_node_id): + def add_dimensionality_reduction_operators(self, builder, current_node_id, rep_dag): dags = [] - modality_type = ( - builder.get_node(current_node_id).operation().output_modality_type - ) + rep_operation = builder.get_node(current_node_id).operation + modality_type = rep_operation().output_modality_type if modality_type is not ModalityType.EMBEDDING: return None @@ -545,6 +544,21 @@ def add_dimensionality_reduction_operators(self, builder, current_node_id): dimensionality_reduction_operators = ( self._get_dimensionality_reduction_operators(modality_type) ) + + if self._dag_output_has_time_axis(rep_dag): + for dimensionality_reduction_op in dimensionality_reduction_operators: + self.pruned.append( + { + "operation": dimensionality_reduction_op.__name__, + "input": rep_operation.__name__, + "reason": ( + "representation output is a sequence; dimensionality " + "reduction only applies to rank-1 vectors" + ), + } + ) + return None + for dimensionality_reduction_op in dimensionality_reduction_operators: dimensionality_reduction_node_id = builder.create_operation_node( dimensionality_reduction_op, @@ -573,7 +587,7 @@ def _build_modality_dag( dags.append(rep_dag) dimensionality_reduction_dags = self.add_dimensionality_reduction_operators( - builder, current_node_id + builder, current_node_id, rep_dag ) if dimensionality_reduction_dags is not None: dags.extend(dimensionality_reduction_dags) From 5f50a9cedb710998f7e9c35e49c4a63e63152964 Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Thu, 15 Oct 2026 08:17:12 +0200 Subject: [PATCH 18/26] find ranks of data --- .../python/systemds/scuro/drsearch/algebra.py | 22 ++++++++++++++++++- .../multimodal_attention_fusion.py | 2 +- 2 files changed, 22 insertions(+), 2 deletions(-) diff --git a/src/main/python/systemds/scuro/drsearch/algebra.py b/src/main/python/systemds/scuro/drsearch/algebra.py index 880e7afe7ea..eba7a018a80 100644 --- a/src/main/python/systemds/scuro/drsearch/algebra.py +++ b/src/main/python/systemds/scuro/drsearch/algebra.py @@ -144,7 +144,7 @@ def _output_type(op_cls): class Violation(NamedTuple): severity: str # error | implicit - kind: str # input_type | arity | hidden_aggregation | zero_padding + kind: str # input_type | arity | rank | hidden_aggregation | zero_padding node_id: str operator: str detail: str @@ -176,6 +176,24 @@ def check_fusion(node): if len(set(widths)) > 1: flag("implicit", "zero_padding", node, f"operand widths {widths}") + def rank(nid): + """seq or vector, None if unknown. Raw modalities declare a max_length.""" + node = nodes[nid] + if node.inputs and node.operation is not None: + declared = operator_signature(node.operation, node.parameters).output_rank + if declared != "any": + return "vector" if declared == "vector" else "seq" + s = stats(nid) + if s is None or not s.output_shape: + return None + seq = len(s.output_shape) > 1 or getattr(s, "max_length", None) is not None + return "seq" if seq else "vector" + + def check_rank(node, sig): + got = rank(node.inputs[0]) + if got is not None and (got == "vector") != (sig.input_rank == "vector"): + flag("error", "rank", node, f"{got} input, expects {sig.input_rank}") + def out_type(nid): if nid in types: return types[nid] @@ -193,6 +211,8 @@ def out_type(nid): for t in in_types: if t is not None and t not in sig.input_types: flag("error", "input_type", node, f"{t.name} input") + if sig.arity == 1 and len(node.inputs) == 1 and sig.input_rank != "any": + check_rank(node, sig) if sig.kind == "fuse" and leaf_stats is not None: check_fusion(node) types[nid] = sig.output_type or in_types[0] diff --git a/src/main/python/systemds/scuro/representations/multimodal_attention_fusion.py b/src/main/python/systemds/scuro/representations/multimodal_attention_fusion.py index 8f4c4db30de..80d9fa1b946 100644 --- a/src/main/python/systemds/scuro/representations/multimodal_attention_fusion.py +++ b/src/main/python/systemds/scuro/representations/multimodal_attention_fusion.py @@ -54,7 +54,7 @@ def __init__( ): parameters = { "hidden_dim": [32, 128, 256, 384, 512, 768], - "num_heads": [2, 4, 8, 12], + "num_heads": [2, 4, 8], "dropout": [0.0, 0.1, 0.2, 0.3, 0.4], # "batch_size": [8, 16, 32, 64, 128], "num_epochs": [10, 20, 50, 100, 150, 200], From 5896fd14df1801260547b1bec598f1c4423ac140 Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Thu, 15 Oct 2026 22:00:31 +0200 Subject: [PATCH 19/26] reording --- .../python/systemds/scuro/drsearch/algebra.py | 29 ++++++++++--------- .../scuro/drsearch/representation_dag.py | 12 +++----- 2 files changed, 20 insertions(+), 21 deletions(-) diff --git a/src/main/python/systemds/scuro/drsearch/algebra.py b/src/main/python/systemds/scuro/drsearch/algebra.py index eba7a018a80..3baf859a519 100644 --- a/src/main/python/systemds/scuro/drsearch/algebra.py +++ b/src/main/python/systemds/scuro/drsearch/algebra.py @@ -439,7 +439,8 @@ def compute(nid, laws): if not inputs: return ("leaf", modality_id, representation_index) - children = [sig(i, laws & EXACT_LAWS if _fitted(op) else laws) for i in inputs] + below = laws & EXACT_LAWS if _fitted(op) else laws - {"concat_perm"} + children = [sig(i, below) for i in inputs] params = dict(params) pushed = ( params.pop("_pushdown_aggregation", None) @@ -482,7 +483,7 @@ def compute(nid, laws): def count_fusion_plans(n, fusion_ops, laws, invariant) -> dict: heads = _fusion_heads(fusion_ops, laws) f = len(heads) - reorders = sum(_reorders(op_laws, invariant) for *_, op_laws in heads) + reorders = sum(_reorders(op_laws, False) for *_, op_laws in heads) ordered = factorial(2 * n - 2) // factorial(n - 1) * f ** (n - 1) if n == 1: canonical = 1 @@ -523,19 +524,21 @@ def _fitted(op) -> bool: def _commutative_plans(n, heads, invariant) -> int: plans = {(1, True): 1, (1, False): 1} + + def count(m, free, invariant): + total = 0 + for *_, op, _, op_laws in heads: + below = free and not _fitted(op) + splits = sum( + comb(m, k) * plans[k, below] * plans[m - k, below] for k in range(1, m) + ) + total += splits // (2 if free and _reorders(op_laws, invariant) else 1) + return total + for m in range(2, n + 1): for free in (True, False): - plans[m, free] = 0 - for *_, op, _, op_laws in heads: - below = free and not _fitted(op) - splits = sum( - comb(m, k) * plans[k, below] * plans[m - k, below] - for k in range(1, m) - ) - plans[m, free] += splits // ( - 2 if free and _reorders(op_laws, invariant) else 1 - ) - return plans[n, True] + plans[m, free] = count(m, free, False) + return count(n, True, invariant) if n > 1 else 1 def _reorders(op_laws, invariant) -> bool: diff --git a/src/main/python/systemds/scuro/drsearch/representation_dag.py b/src/main/python/systemds/scuro/drsearch/representation_dag.py index d400ccbb199..5567c45554e 100644 --- a/src/main/python/systemds/scuro/drsearch/representation_dag.py +++ b/src/main/python/systemds/scuro/drsearch/representation_dag.py @@ -44,7 +44,6 @@ count_violations, explicit_fusion_operands, explicit_prep_on, - order_invariant, typecheck, ) from systemds.scuro.representations.fusion import Fusion @@ -738,7 +737,6 @@ def __init__(self, algebra: Optional[AlgebraConfig] = None, tasks=None, leaves=( self.algebra = algebra or AlgebraConfig() self.leaves = leaves self.algebra_stats = Counter() - self._order_invariant = order_invariant(self.algebra, tasks) self._seen_ordered = set() self._seen_legacy = set() self.global_nodes: List[RepresentationNode] = [] @@ -751,19 +749,17 @@ def __init__(self, algebra: Optional[AlgebraConfig] = None, tasks=None, leaves=( def _compute_node_signature( self, operation: Any, inputs: List[str], parameters: Dict[str, Any] = None ) -> Hashable: - ip = [self.node_to_signature[inp] for inp in inputs] - input_sigs = tuple(sorted(ip, key=repr)) if inputs else () + input_sigs = tuple(self.node_to_signature[inp] for inp in inputs) op_cls = operation().name params_items = tuple(sorted((parameters or {}).items())) - if self.algebra.cse_law_aware and not self._may_reorder(operation, parameters): - input_sigs = tuple(ip) + if self._may_reorder(operation, parameters): + input_sigs = tuple(sorted(input_sigs, key=repr)) return ("op", op_cls, params_items, input_sigs) def _may_reorder(self, operation: Any, parameters: Dict[str, Any]) -> bool: if not issubclass(operation, Fusion): return False - laws = operation.algebraic_laws(parameters or {}) - return "comm" in laws or ("concat_perm" in laws and self._order_invariant) + return "comm" in operation.algebraic_laws(parameters or {}) def _count_law_merges(self, signature: Hashable, inputs: List[str]): ip = tuple(self.node_to_signature[inp] for inp in inputs) From fc34f97c730c3b276358becf097ca7cc284126cc Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Thu, 15 Oct 2026 22:19:44 +0200 Subject: [PATCH 20/26] fix padding --- .../systemds/scuro/modality/transformed.py | 5 ++++ .../scuro/modality/unimodal_modality.py | 4 +++ .../scuro/representations/aggregate.py | 29 +++++++++++++++++++ .../aggregated_representation.py | 3 +- .../representations/window_aggregation.py | 11 ++----- 5 files changed, 42 insertions(+), 10 deletions(-) diff --git a/src/main/python/systemds/scuro/modality/transformed.py b/src/main/python/systemds/scuro/modality/transformed.py index 87cc3e91434..5806e3d4098 100644 --- a/src/main/python/systemds/scuro/modality/transformed.py +++ b/src/main/python/systemds/scuro/modality/transformed.py @@ -25,6 +25,7 @@ from systemds.scuro.modality.joined import JoinedModality from systemds.scuro.modality.modality import Modality from systemds.scuro.representations.window_aggregation import WindowAggregation +from systemds.scuro.representations.aggregate import set_valid_windows import time import sys @@ -129,6 +130,7 @@ def window_aggregation(self, window_size, aggregation): ) start = time.time() transformed_modality.data = w.execute(self) + set_valid_windows(transformed_modality, w.valid_windows) transformed_modality.transform_time += time.time() - start return transformed_modality @@ -138,6 +140,9 @@ def context(self, context_operator): ) start = time.time() transformed_modality.data = context_operator.execute(self) + set_valid_windows( + transformed_modality, getattr(context_operator, "valid_windows", None) + ) transformed_modality.transform_time += time.time() - start return transformed_modality diff --git a/src/main/python/systemds/scuro/modality/unimodal_modality.py b/src/main/python/systemds/scuro/modality/unimodal_modality.py index 1b6173d8aab..df2c3adb733 100644 --- a/src/main/python/systemds/scuro/modality/unimodal_modality.py +++ b/src/main/python/systemds/scuro/modality/unimodal_modality.py @@ -29,6 +29,7 @@ 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.aggregate import set_valid_windows from systemds.scuro.representations.representation import ( CONTAINER_ARRAY, RepresentationStats, @@ -178,6 +179,9 @@ def context(self, context_operator): transformed_modality.data = d else: transformed_modality.data = self.data + set_valid_windows( + transformed_modality, getattr(context_operator, "valid_windows", None) + ) transformed_modality.transform_time += time.time() - start return transformed_modality diff --git a/src/main/python/systemds/scuro/representations/aggregate.py b/src/main/python/systemds/scuro/representations/aggregate.py index 2c54913445c..bf1ae6517e9 100644 --- a/src/main/python/systemds/scuro/representations/aggregate.py +++ b/src/main/python/systemds/scuro/representations/aggregate.py @@ -22,6 +22,32 @@ from systemds.scuro.representations import utils +VALID_WINDOWS_KEY = "valid_windows" + + +def set_valid_windows(modality, valid_windows=None): + """Record the valid window counts per instance; None clears stale ones.""" + metadata = modality.metadata + if not metadata: + return + if valid_windows is not None and len(valid_windows) != len(metadata): + valid_windows = None + for i, md in enumerate(metadata): + if not isinstance(md, dict): + continue + md = {k: v for k, v in md.items() if k != VALID_WINDOWS_KEY} + if valid_windows is not None: + md[VALID_WINDOWS_KEY] = valid_windows[i] + metadata[i] = md + + +def _drop_padded_windows(instance, metadata, i): + md = metadata[i] if metadata is not None and i < len(metadata) else None + counts = md.get(VALID_WINDOWS_KEY) if isinstance(md, dict) else None + if counts is not None and len(instance) == counts[1]: + return instance[: counts[0]] + return instance + class Aggregation: @staticmethod @@ -104,6 +130,8 @@ def get_current_parameters(self): def execute(self, modality, aggregate_dim=(0,), squeeze_singleton=True): data = [] max_len = 0 + axes = aggregate_dim if isinstance(aggregate_dim, tuple) else (aggregate_dim,) + metadata = getattr(modality, "metadata", None) if 0 in axes else None for i, instance in enumerate(modality.data): data.append([]) if isinstance(instance, np.ndarray) or isinstance(instance, list): @@ -115,6 +143,7 @@ def execute(self, modality, aggregate_dim=(0,), squeeze_singleton=True): ): aggregated_data = instance.flatten() else: + instance = _drop_padded_windows(instance, metadata, i) aggregated_data = self._aggregation_func(instance, aggregate_dim) else: aggregated_data = [] diff --git a/src/main/python/systemds/scuro/representations/aggregated_representation.py b/src/main/python/systemds/scuro/representations/aggregated_representation.py index b37b1eceda9..35ea8c6eff6 100644 --- a/src/main/python/systemds/scuro/representations/aggregated_representation.py +++ b/src/main/python/systemds/scuro/representations/aggregated_representation.py @@ -24,7 +24,7 @@ Representation, RepresentationStats, ) -from systemds.scuro.representations.aggregate import Aggregation +from systemds.scuro.representations.aggregate import Aggregation, set_valid_windows import time import numpy as np @@ -124,6 +124,7 @@ def transform(self, modality): ) aggregated_modality.data = aggregated_data + set_valid_windows(aggregated_modality) end = time.perf_counter() aggregated_modality.transform_time += end - start diff --git a/src/main/python/systemds/scuro/representations/window_aggregation.py b/src/main/python/systemds/scuro/representations/window_aggregation.py index 73843cca1e3..b80906c7ead 100644 --- a/src/main/python/systemds/scuro/representations/window_aggregation.py +++ b/src/main/python/systemds/scuro/representations/window_aggregation.py @@ -379,6 +379,7 @@ def estimate_peak_memory_bytes(self, input_stats: RepresentationStats) -> dict: return {"cpu_peak_bytes": cpu_peak, "gpu_peak_bytes": 0} def execute(self, modality): + self.valid_windows = None batch = self._dense_batch(modality) if batch is not None: windowed_data = self.window_aggregate_single_level_batched(batch) @@ -424,15 +425,7 @@ def execute(self, modality): else: padded_features[i, : features.shape[0]] = features - # attention_masks = np.zeros((len(windowed_data), target_length)) - # for i, length in enumerate(original_lengths): - # actual_length = min(length, target_length) - # attention_masks[i, :actual_length] = 1 - - # ModalityType(modality.modality_type).add_field_for_instances( - # modality.metadata, "attention_masks", attention_masks - # ) - + self.valid_windows = [(n, target_length) for n in original_lengths] windowed_data = padded_features data_type = modality.metadata[0]["data_layout"]["type"] if data_type != "str": From db7f851c2fe3268bf4ce4f2ccee41e222faab10b Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Thu, 15 Oct 2026 22:50:26 +0200 Subject: [PATCH 21/26] fit reps in multimodal optimization --- .../scuro/drsearch/multimodal_ga_optimizer.py | 2 +- .../scuro/drsearch/multimodal_optimizer.py | 4 +-- .../scuro/drsearch/split_aware_evaluator.py | 31 ++++++++++++++++--- .../scuro/drsearch/unimodal_optimizer.py | 2 ++ 4 files changed, 31 insertions(+), 8 deletions(-) diff --git a/src/main/python/systemds/scuro/drsearch/multimodal_ga_optimizer.py b/src/main/python/systemds/scuro/drsearch/multimodal_ga_optimizer.py index 1f591c58d79..35e5c6a1e30 100644 --- a/src/main/python/systemds/scuro/drsearch/multimodal_ga_optimizer.py +++ b/src/main/python/systemds/scuro/drsearch/multimodal_ga_optimizer.py @@ -261,7 +261,7 @@ def _evaluate_genome_body( share_subtrees: bool = False, ) -> Tuple[Optional[Tuple[float, ...]], Optional[Dict[str, Any]]]: start = time.time() - if dag_requires_fitting(dag): + if dag_requires_fitting(dag, modalities): node_cache = _subtree_node_cache(dag, task) if share_subtrees else None if node_cache is not None: store = node_cache.store diff --git a/src/main/python/systemds/scuro/drsearch/multimodal_optimizer.py b/src/main/python/systemds/scuro/drsearch/multimodal_optimizer.py index 22476bce60c..a4f6a701efe 100644 --- a/src/main/python/systemds/scuro/drsearch/multimodal_optimizer.py +++ b/src/main/python/systemds/scuro/drsearch/multimodal_optimizer.py @@ -72,7 +72,7 @@ def _evaluate_dag_worker(dag_pickle, task_pickle, modalities_pickle, debug=False f"[DEBUG][worker] pid={os.getpid()} evaluating dag_root={getattr(dag, 'root_node_id', None)} task={getattr(task.model, 'name', None)}" ) - if dag_requires_fitting(dag): + if dag_requires_fitting(dag, modalities_for_dag): evaluation = SplitAwareDAGEvaluator( dag, modalities_for_dag, task ).evaluate() @@ -491,7 +491,7 @@ def _evaluate_dag(self, dag: RepresentationDag, task: Task) -> "OptimizationResu self.k_best_representations[task.model.name].values() ) ) - if dag_requires_fitting(dag): + if dag_requires_fitting(dag, modalities_for_dag): evaluation = SplitAwareDAGEvaluator( dag, modalities_for_dag, task ).evaluate() diff --git a/src/main/python/systemds/scuro/drsearch/split_aware_evaluator.py b/src/main/python/systemds/scuro/drsearch/split_aware_evaluator.py index 05b70d03af6..aead3e9dfaf 100644 --- a/src/main/python/systemds/scuro/drsearch/split_aware_evaluator.py +++ b/src/main/python/systemds/scuro/drsearch/split_aware_evaluator.py @@ -53,11 +53,17 @@ class SplitEvaluationResult: task_time: float -def dag_requires_fitting(dag: RepresentationDag) -> bool: +def dag_requires_fitting(dag: RepresentationDag, modalities=()) -> bool: + fold_ids = { + str(m.modality_id) for m in modalities if getattr(m, "fold_source", None) + } return any( - node.operation is not None - and getattr(node.operation, "fit_requirement", FitRequirement.NONE) - is not FitRequirement.NONE + ( + node.operation is not None + and getattr(node.operation, "fit_requirement", FitRequirement.NONE) + is not FitRequirement.NONE + ) + or (not node.inputs and str(node.modality_id) in fold_ids) for node in dag.nodes ) @@ -235,6 +241,12 @@ def execute_node(node_id): result = tuple( self._leaf_input(partition, node) for partition in partitions ) + source = getattr( + self._leaf_input(self._full, node).modality, "fold_source", None + ) + if source is not None: + result = self._refit_leaf(source, result) + fitted.add(node_id) cache[node_id] = result return result @@ -302,12 +314,21 @@ def execute_node(node_id): return execute_node(self.dag.root_node_id) + def _refit_leaf(self, source, partitions): + source_dag, raw_modality = source + evaluator = SplitAwareDAGEvaluator(source_dag, raw_modality, self.task) + (full,) = evaluator._indexed_modalities() + return evaluator._execute_partitions( + *((full.subset(item.instance_ids),) for item in partitions) + ) + def _indexed_modalities(self): full = [] for modality in self.modalities: self.materialize(modality) full.append(IndexedModality(modality, range(len(modality.data)))) - return tuple(full) + self._full = tuple(full) + return self._full def fit_transform(self, fit_indices, transform_indices) -> IndexedModality: full = self._indexed_modalities() diff --git a/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py b/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py index bd6956e298e..431dc6666f2 100644 --- a/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py +++ b/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py @@ -1149,12 +1149,14 @@ def get_k_best_results( cache = [] for result in results: if dag_requires_fitting(result.dag): + SplitAwareDAGEvaluator.materialize(modality) transformed = SplitAwareDAGEvaluator( result.dag, modality, task ).fit_transform( task.train_indices, range(len(modality.data)), ) + transformed.modality.fold_source = (result.dag, modality) cache.append(transformed.modality) else: cache.append(result.dag.execute([modality])) From 236db1fd16f0f76230c5b8904a02b966a2cdc5bc Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Thu, 15 Oct 2026 23:07:32 +0200 Subject: [PATCH 22/26] swap same subtree --- .../scuro/drsearch/multimodal_ga_optimizer.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/src/main/python/systemds/scuro/drsearch/multimodal_ga_optimizer.py b/src/main/python/systemds/scuro/drsearch/multimodal_ga_optimizer.py index 35e5c6a1e30..97281b2d2a6 100644 --- a/src/main/python/systemds/scuro/drsearch/multimodal_ga_optimizer.py +++ b/src/main/python/systemds/scuro/drsearch/multimodal_ga_optimizer.py @@ -1315,11 +1315,15 @@ def _crossover_genomes( c1, c2 = copy.deepcopy(g1), copy.deepcopy(g2) if c1.leaves == c2.leaves: - paths1 = self._internal_paths(c1.tree) - paths2 = self._internal_paths(c2.tree) - if paths1 and paths2: - path1 = self.rng.choice(paths1) - path2 = self.rng.choice(paths2) + pairs = [ + (path1, path2) + for path1 in self._internal_paths(c1.tree) + for path2 in self._internal_paths(c2.tree) + if sorted(_collect_leaf_indices(_get_subtree(c1.tree, path1))) + == sorted(_collect_leaf_indices(_get_subtree(c2.tree, path2))) + ] + if pairs: + path1, path2 = self.rng.choice(pairs) subtree1 = _get_subtree(c1.tree, path1) subtree2 = _get_subtree(c2.tree, path2) c1.tree = _replace_subtree(c1.tree, path1, subtree2) From 0c479dc3e69c18119f9c8877b61065940a9f9d92 Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Thu, 15 Oct 2026 23:22:01 +0200 Subject: [PATCH 23/26] fix upper bound for windowing stats --- .../representations/window_aggregation.py | 24 +++++++++++++++---- 1 file changed, 19 insertions(+), 5 deletions(-) diff --git a/src/main/python/systemds/scuro/representations/window_aggregation.py b/src/main/python/systemds/scuro/representations/window_aggregation.py index b80906c7ead..115fb433d79 100644 --- a/src/main/python/systemds/scuro/representations/window_aggregation.py +++ b/src/main/python/systemds/scuro/representations/window_aggregation.py @@ -251,8 +251,9 @@ def assert_output_stats(self, windowed_data): assert ( windowed_data.shape[1] == self.stats.output_shape[0] ), f"Output shape: {windowed_data.shape}, Expected shape: {self.stats.output_shape}" - assert ( - windowed_data.shape[2:] == self.stats.output_shape[1:] + feat_shape, bound = windowed_data.shape[2:], self.stats.output_shape[1:] + assert len(feat_shape) == len(bound) and all( + s <= b for s, b in zip(feat_shape, bound) ), f"Output shape: {windowed_data.shape}, Expected shape: {self.stats.output_shape}" @staticmethod @@ -683,14 +684,27 @@ def _window_sizes(self, signal_length: int) -> np.ndarray: sizes[-1] += length - int(sizes.sum()) return sizes + def _max_window_size(self, max_length: int) -> int: + # Upper bound on every window of every instance with at most max_length + # elements. Window i < q-1 holds more than w_i (L - q) elements, so the + # last one (which also takes the rounding remainder) holds fewer than + # q + w_last (L - q); and no window exceeds L - q + 1. + num_windows = self._effective_num_windows(max_length) + length = max(int(max_length), num_windows) + weights = np.geomspace(4, 256, num=num_windows) + last_share = weights[-1] / np.sum(weights) + geometric_bound = int(np.ceil(num_windows + last_share * (length - num_windows))) + return max(1, min(length - num_windows + 1, geometric_bound)) + def _feature_shape(self, in_shape): if isinstance(self.aggregation_function, Aggregation): return in_shape[1:] length = in_shape[0] if in_shape else 0 # Window sizes grow geometrically and _pad_stack pads every window's - # features to the largest one, so size them from the largest window - # (about 4x the mean), not the mean. - max_window_size = int(self._window_sizes(length).max()) if in_shape else 1 + # features to the largest one, so size them from an upper bound on the + # largest window (about 4x the mean) over all lengths up to L, not the + # mean: a shorter instance can have a longer last window. + max_window_size = self._max_window_size(length) if in_shape else 1 return self._per_window_feature_shape(max_window_size) def get_output_stats(self, input_stats: RepresentationStats) -> tuple: From ef836f2e752e7156074cb6112e209b6937712778 Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Thu, 15 Oct 2026 23:29:55 +0200 Subject: [PATCH 24/26] update memory of failed nodes --- .../python/systemds/scuro/drsearch/node_executor.py | 2 ++ .../python/systemds/scuro/drsearch/node_scheduler.py | 4 ++++ .../python/systemds/scuro/drsearch/worker_pool.py | 7 +++++-- .../python/systemds/scuro/utils/memory_utility.py | 12 ++++++++++++ 4 files changed, 23 insertions(+), 2 deletions(-) diff --git a/src/main/python/systemds/scuro/drsearch/node_executor.py b/src/main/python/systemds/scuro/drsearch/node_executor.py index 753683ede67..ab9d1017742 100644 --- a/src/main/python/systemds/scuro/drsearch/node_executor.py +++ b/src/main/python/systemds/scuro/drsearch/node_executor.py @@ -940,6 +940,8 @@ def _handle_job_failure(self, unit: Union[_NodeUnit, _BatchUnit], jr): requeued = self._requeue_or_give_up(node_id, reason) if not requeued: self._release_parents(node_id) + elif isinstance(unit, _NodeUnit) and jr.gpu_peak_bytes > 0: + self.scheduler.raise_gpu_estimate(node_id, jr.gpu_peak_bytes) def _handle_node_success(self, node_id: str, value: Dict[str, Any]) -> None: node = self.scheduler.mapping[node_id] diff --git a/src/main/python/systemds/scuro/drsearch/node_scheduler.py b/src/main/python/systemds/scuro/drsearch/node_scheduler.py index 257d62a2218..5498ad98c7a 100644 --- a/src/main/python/systemds/scuro/drsearch/node_scheduler.py +++ b/src/main/python/systemds/scuro/drsearch/node_scheduler.py @@ -323,6 +323,10 @@ def requeue_node(self, node_id: str) -> None: self._release_execution_memory(node_id, self.mapping[node_id].gpu_id) self._candidates.add(node_id) + def raise_gpu_estimate(self, node_id: str, observed_bytes: int) -> None: + cpu_mem, gpu_mem = self.node_resources[node_id] + self.node_resources[node_id] = (cpu_mem, max(gpu_mem, int(observed_bytes))) + 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] diff --git a/src/main/python/systemds/scuro/drsearch/worker_pool.py b/src/main/python/systemds/scuro/drsearch/worker_pool.py index 9529049a238..f54578aca1f 100644 --- a/src/main/python/systemds/scuro/drsearch/worker_pool.py +++ b/src/main/python/systemds/scuro/drsearch/worker_pool.py @@ -30,7 +30,7 @@ import torch -from systemds.scuro.utils.memory_utility import is_cuda_oom +from systemds.scuro.utils.memory_utility import cuda_oom_peak_bytes, is_cuda_oom _THREAD_ENV_VARS = ( "OMP_NUM_THREADS", @@ -102,6 +102,7 @@ class _JobResult: error: Optional[str] = None cuda_oom: bool = False worker_died: bool = False + gpu_peak_bytes: int = -1 def _worker_main( @@ -129,13 +130,15 @@ def _worker_main( 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: + cuda_oom = is_cuda_oom(e) result_q.put( _JobResult( job.job_id, False, os.getpid(), error=f"{type(e).__name__}: {e}", - cuda_oom=is_cuda_oom(e), + cuda_oom=cuda_oom, + gpu_peak_bytes=cuda_oom_peak_bytes(e) if cuda_oom else -1, ) ) diff --git a/src/main/python/systemds/scuro/utils/memory_utility.py b/src/main/python/systemds/scuro/utils/memory_utility.py index a28f14b5004..5b80431add0 100644 --- a/src/main/python/systemds/scuro/utils/memory_utility.py +++ b/src/main/python/systemds/scuro/utils/memory_utility.py @@ -19,6 +19,7 @@ # # ------------------------------------------------------------- import os +import re import resource import sys import threading @@ -255,6 +256,17 @@ def is_cuda_oom(exc: BaseException) -> bool: ) +def cuda_oom_peak_bytes(exc: BaseException) -> int: + if not torch.cuda.is_available(): + return -1 + peak = torch.cuda.max_memory_allocated() + m = re.search(r"tried to allocate ([\d.]+) (bytes|kib|mib|gib)", str(exc).lower()) + if m: + unit = ("bytes", "kib", "mib", "gib").index(m.group(2)) + peak += int(float(m.group(1)) * 1024**unit) + return int(peak) + + def cleanup_gpu(gpu_id: Optional[int]) -> None: if gpu_id is None or not torch.cuda.is_available(): return From 68c1e39a75dbdc62cb532b42ffc220868606789e Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Wed, 28 Oct 2026 13:26:55 +0100 Subject: [PATCH 25/26] fix errors --- .../systemds/scuro/dataloader/audio_loader.py | 4 +- .../scuro/drsearch/hyperparameter_tuner.py | 28 ++- .../scuro/drsearch/multimodal_ga_optimizer.py | 33 +++- .../systemds/scuro/drsearch/node_executor.py | 183 ++++++++++++++++-- .../systemds/scuro/drsearch/node_scheduler.py | 43 ++-- .../scuro/drsearch/operator_registry.py | 10 +- .../python/systemds/scuro/drsearch/ranking.py | 30 +-- .../scuro/drsearch/split_aware_evaluator.py | 37 ++-- .../python/systemds/scuro/drsearch/task.py | 80 +++++++- .../scuro/drsearch/unimodal_optimizer.py | 160 +++++++++++---- .../systemds/scuro/drsearch/worker_pool.py | 12 +- .../systemds/scuro/modality/transformed.py | 2 + .../systemds/scuro/representations/clip.py | 1 + .../systemds/scuro/representations/fusion.py | 33 +++- .../scuro/representations/image_bind.py | 135 +++++++++++-- .../systemds/scuro/representations/lstm.py | 8 +- .../multimodal_attention_fusion.py | 6 +- .../scuro/representations/openface.py | 16 ++ .../timeseries_representations.py | 4 +- .../scuro/representations/unimodal.py | 10 +- .../systemds/scuro/representations/utils.py | 11 +- .../representations/window_aggregation.py | 4 +- .../systemds/scuro/utils/memory_utility.py | 9 +- .../systemds/scuro/utils/static_variables.py | 46 ++++- .../systemds/scuro/utils/torch_dataset.py | 19 +- .../tests/scuro/test_algebra_canonical.py | 17 ++ .../python/tests/scuro/test_algebra_cse.py | 15 +- .../python/tests/scuro/test_algebra_laws.py | 13 +- .../tests/scuro/test_algebra_plan_counts.py | 4 +- .../python/tests/scuro/test_algebra_stats.py | 3 +- .../tests/scuro/test_algebra_typecheck.py | 28 ++- .../python/tests/scuro/test_data_loaders.py | 47 ++++- src/main/python/tests/scuro/test_hp_tuner.py | 39 ++++ .../python/tests/scuro/test_modality_pad.py | 6 +- .../scuro/test_multimodal_ga_optimizer.py | 118 +++++++++++ .../tests/scuro/test_operator_registry.py | 14 ++ src/main/python/tests/scuro/test_scheduler.py | 2 +- .../test_transformer_text_aggregation.py | 79 +++++--- .../tests/scuro/test_unimodal_optimizer.py | 166 ++++++++++++++++ .../tests/scuro/test_window_operations.py | 51 +++++ .../test_window_representation_batching.py | 2 + 41 files changed, 1324 insertions(+), 204 deletions(-) diff --git a/src/main/python/systemds/scuro/dataloader/audio_loader.py b/src/main/python/systemds/scuro/dataloader/audio_loader.py index f24f6e53b5c..7ed59fc1150 100644 --- a/src/main/python/systemds/scuro/dataloader/audio_loader.py +++ b/src/main/python/systemds/scuro/dataloader/audio_loader.py @@ -87,9 +87,7 @@ def _decode_file(self, file: str): audio = np.array([0]) sr = 1000 else: - audio, sr = librosa.load( - file, sr=self.sampling_rate, dtype=self._data_type - ) + audio, sr = librosa.load(file, sr=self.sampling_rate, dtype=self._data_type) if self.normalize: audio = librosa.util.normalize(audio) diff --git a/src/main/python/systemds/scuro/drsearch/hyperparameter_tuner.py b/src/main/python/systemds/scuro/drsearch/hyperparameter_tuner.py index d006e2f87f9..14658714d4b 100644 --- a/src/main/python/systemds/scuro/drsearch/hyperparameter_tuner.py +++ b/src/main/python/systemds/scuro/drsearch/hyperparameter_tuner.py @@ -53,6 +53,7 @@ trial_node_signatures, ) from systemds.scuro.drsearch.process_cache import HPOExecutionCache +from systemds.scuro.drsearch.ranking import tie_broken from systemds.scuro.drsearch.operator_registry import is_expensive_representation _CUDA_OOM_ERRORS = tuple( @@ -144,6 +145,29 @@ def _expand_aggregation_param_specs( return specs +def _limit_context_words(node: RepresentationNode, dag, params: Dict[str, Any]): + """Drop max_words values longer than the consuming encoder's text window. + + A text context operator's chunks are only seen in full by an encoder up to + its max_context_words; longer chunks are truncated, so those trials repeat + a shorter configuration. + """ + values = params.get("max_words") + if dag is None or not isinstance(values, list): + return params + limits = [ + getattr(consumer.operation, "max_context_words", None) + for consumer in dag.nodes + if node.node_id in consumer.inputs + ] + limits = [limit for limit in limits if limit is not None] + if not limits: + return params + limit = min(limits) + params["max_words"] = [value for value in values if value <= limit] or [min(values)] + return params + + def _is_window_operation(op: Any) -> bool: if not inspect.isclass(op): return False @@ -777,6 +801,7 @@ def __get_params_for_node( params[name] = ( op.filter_parameter_domain(name, values, input_stats) or values ) + params = _limit_context_words(node, dag, params) if node.parameters: if inspect.isclass(node.parameters.get("aggregation_function")): params["aggregation_function"] = node.parameters["aggregation_function"] @@ -1368,8 +1393,9 @@ def _get_metric_value(result): ) return score + candidates = optimization_results[task.model.name] best_results = sorted( - optimization_results[task.model.name], + [candidates[i] for i in tie_broken(candidates)], key=_get_metric_value, reverse=self.maximize_metric, )[:k] diff --git a/src/main/python/systemds/scuro/drsearch/multimodal_ga_optimizer.py b/src/main/python/systemds/scuro/drsearch/multimodal_ga_optimizer.py index 97281b2d2a6..f754116de93 100644 --- a/src/main/python/systemds/scuro/drsearch/multimodal_ga_optimizer.py +++ b/src/main/python/systemds/scuro/drsearch/multimodal_ga_optimizer.py @@ -34,7 +34,11 @@ from dataclasses import dataclass, field from itertools import chain from typing import Any, Dict, List, Optional, Tuple, Union -from systemds.scuro.utils.static_variables import get_global_seed +from systemds.scuro.utils.static_variables import ( + cuda_allowed, + deterministic_enabled, + get_global_seed, +) import numpy as np from deap import base, tools @@ -344,6 +348,10 @@ def _evaluate_genome_body( def _dispatch_genome_evaluation(payload, _gpu_id): dag, task, modalities, objective_specs, *share = payload + allowed = os.sched_getaffinity(0) + if cuda_allowed(): + cores = sorted(allowed) + os.sched_setaffinity(0, {cores[os.getpid() % len(cores)]}) try: try: fitness, result = _evaluate_genome_body( @@ -359,6 +367,7 @@ def _dispatch_genome_evaluation(payload, _gpu_id): "_evaluation_error": traceback.format_exc() } finally: + os.sched_setaffinity(0, allowed) _release_cuda_memory() @@ -372,7 +381,6 @@ def __init__(self, values, fitness_type): class MultimodalDeapOptimizer: - # For instances made without __init__, as regret_analysis/exhaustive_enumerator.py does. algebra = AlgebraConfig() def __init__( @@ -533,6 +541,7 @@ def _optimize( for task in self.tasks: task_name = task.model.name self._current_task_name = task_name + self.rng = self._task_rng(task_name) self.optimization_results.setdefault(task_name, []) self.evaluation_errors.setdefault(task_name, 0) self.oversized_candidates.setdefault(task_name, 0) @@ -601,6 +610,20 @@ def _optimize( return self.optimization_results + def _task_rng(self, task_name: str) -> random.Random: + return random.Random(f"{self.random_seed}:{task_name}") + + def _select_tournament(self, population: List[Any], k: int, tournsize: int): + """deap.tools.selTournament, but drawing from self.rng instead of the + global random module.""" + return [ + max( + (self.rng.choice(population) for _ in range(tournsize)), + key=lambda ind: ind.fitness, + ) + for _ in range(k) + ] + def _make_individual(self, genome: DagGenome): return _FusionIndividual([genome], self._fitness_type) @@ -676,7 +699,7 @@ def _breed_offspring( retries = 0 tournsize = max(1, min(3, len(population))) while len(next_population) < self.population_size and retries < retry_budget: - p1, p2 = tools.selTournament(population, 2, tournsize=tournsize) + p1, p2 = self._select_tournament(population, 2, tournsize) if self.rng.random() < self.crossover_probability: g1, g2 = self._crossover_genomes(p1[0], p2[0]) @@ -770,6 +793,8 @@ def _idle_lane(self, uses_gpu: bool): for gpu_id in pool.gpu_worker_devices: if pool.has_idle_worker_for(gpu_id): return gpu_id + if deterministic_enabled() and pool.gpu_worker_devices: + return False return None if pool.has_idle_cpu_worker else False if pool.has_idle_worker_for(None, allow_gpu_worker_for_cpu=True): return None @@ -1344,7 +1369,7 @@ def _crossover_genomes( ) return c1, c2 - shared_paths = set(c1.fusion_ops) & set(c2.fusion_ops) + shared_paths = sorted(set(c1.fusion_ops) & set(c2.fusion_ops)) for path in shared_paths: if self.rng.random() < 0.5: c1.fusion_ops[path], c2.fusion_ops[path] = ( diff --git a/src/main/python/systemds/scuro/drsearch/node_executor.py b/src/main/python/systemds/scuro/drsearch/node_executor.py index ab9d1017742..4647c23fcaf 100644 --- a/src/main/python/systemds/scuro/drsearch/node_executor.py +++ b/src/main/python/systemds/scuro/drsearch/node_executor.py @@ -18,6 +18,7 @@ # under the License. # # ------------------------------------------------------------- +import math import multiprocessing as mp import os import time @@ -43,6 +44,13 @@ RepresentationNode, ) from systemds.scuro.drsearch.task import PerformanceMeasure +from systemds.scuro.models.fused_mlp import ( + candidate_bytes, + candidate_slots, + fused_gpu_bytes, + fused_training_enabled, + supports_fused_training, +) from systemds.scuro.drsearch.worker_pool import ( PersistentWorkerPool, create_mp_context, @@ -73,6 +81,8 @@ from systemds.scuro.utils.static_variables import DEBUG _MAX_NODE_RETRIES = int(os.environ.get("SCURO_MAX_NODE_RETRIES", "3")) +_FUSED_CONTEXT_BYTES = 512 * 2**20 +_FUSED_WAIT_S = float(os.environ.get("SCURO_FUSED_WAIT_S", "30")) def _materialize_lazy_leaves() -> bool: @@ -263,6 +273,51 @@ def _run_node_op(): } +def _task_data(modality, aggregation): + if aggregation is None: + return modality.data + return aggregation.operation(params=aggregation.parameters).transform(modality).data + + +def _execute_task_batch_worker(task: Any, items: List[tuple], gpu_id: int): + start_time = time.perf_counter() + device = torch.device(f"cuda:{gpu_id}") + torch.cuda.set_device(device) + datas, failed = [], {} + for node_id, modality, aggregation in items: + try: + datas.append((node_id, _task_data(modality, aggregation))) + except Exception as e: + failed[node_id] = f"{type(e).__name__}: {e}" + + allowed = os.sched_getaffinity(0) + cores = sorted(allowed) + os.sched_setaffinity(0, {cores[os.getpid() % len(cores)]}) + try: + scores = task.run_many([data for _, data in datas], device) + finally: + os.sched_setaffinity(0, allowed) + cleanup_gpu(gpu_id) + share = (time.perf_counter() - start_time) / max(1, len(datas)) + results = {} + for i, ((node_id, _), result) in enumerate(zip(datas, scores)): + if isinstance(result, Exception): + failed[node_id] = f"{type(result).__name__}: {result}" + continue + results[node_id] = { + "scores": result, + "task_time": share, + "task_timing": dict(task.last_run_timing), + "memory": None, + "peak_bytes": -1, + "gpu_peak_bytes": -1, + "start_time": start_time + i * share, + "end_time": start_time + (i + 1) * share, + "pid": os.getpid(), + } + return {"results": results, "failed_nodes": failed} + + def _execute_task_worker( task_node_id: str, task: Any, @@ -282,15 +337,7 @@ def _execute_task_worker( def _run_task(): start = time.perf_counter() - if aggregation is not None: - data = ( - aggregation.operation(params=aggregation.parameters) - .transform(modality) - .data - ) - else: - data = modality.data - scores = task.run(data) + scores = task.run(_task_data(modality, aggregation)) end = time.perf_counter() return scores, end - start @@ -449,6 +496,11 @@ def _dispatch_task(payload, gpu_id): return _execute_task_worker(task_node_id, task, modality, gpu_id, aggregation) +def _dispatch_task_batch(payload, gpu_id): + task, items = payload + return _execute_task_batch_worker(task, items, gpu_id) + + def _dispatch_leaf_batch(payload, gpu_id): nodes, modality = payload return _execute_leaf_batch_worker(nodes, modality, gpu_id) @@ -462,6 +514,7 @@ def _dispatch_load_leaf(payload, _gpu_id): _WORKER_DISPATCH = { "node": _dispatch_node, "task": _dispatch_task, + "task_batch": _dispatch_task_batch, "leaf_batch": _dispatch_leaf_batch, "load_leaf": _dispatch_load_leaf, "release_models": _dispatch_release_models, @@ -483,6 +536,12 @@ class _BatchUnit: node_ids: List[str] +@dataclass +class _TaskBatchUnit(_BatchUnit): + gpu_id: int = 0 + gpu_bytes: int = 0 + + @dataclass class ResultEntry: val_score: PerformanceMeasure = None @@ -572,6 +631,8 @@ def __init__( self._leaf_original_data: Dict[str, Any] = {} self._leaf_unshared_originals: List[tuple] = [] self._task_results: Dict[str, ResultEntry] = {} + self._fused_shapes: Dict[tuple, tuple] = {} + self._fused_ready_since: Dict[str, float] = {} self._owns_pool = worker_pool is None if worker_pool is None: @@ -615,7 +676,7 @@ def _requeue_or_give_up(self, node_id: str, reason: str) -> bool: return True def _release_parents(self, node_id: str) -> None: - for parent_id in self.scheduler.get_valid_parents(node_id): + for parent_id in set(self.scheduler.get_valid_parents(node_id)): self._result_cache.dec_ref(parent_id) def _retain_for_submit(self, parent_ids: List[str], payload: Any) -> List[str]: @@ -835,10 +896,90 @@ def _release_all_worker_models(self) -> None: while self._job_units: self._process_result(self._pool.wait()) + def _fused_task_nodes(self, ready) -> List[str]: + if not self.scheduler.gpu_devices or not fused_training_enabled(): + return [] + fused = [] + for node_id in ready: + if isinstance(node_id, list): + continue + node = self.scheduler.mapping[node_id] + if not self._is_task_node(node): + continue + task = self._tasks[int(node.parameters.get("_task_idx", 0))] + if not supports_fused_training(getattr(task, "model", None)): + continue + if self._fused_batch_shape(node)[1] <= fused_gpu_bytes(): + fused.append(node_id) + return fused + + def _fused_batch_shape(self, node) -> tuple: + task_idx = int(node.parameters.get("_task_idx", 0)) + stats = self.scheduler.node_stats.get(node.inputs[0]) + width = math.prod(getattr(stats, "output_shape", None) or (1,)) + key = (task_idx, int(width)) + if key not in self._fused_shapes: + task = self._tasks[task_idx] + rows = len(set(task.train_indices) | set(task.test_indices or [])) + per_candidate = candidate_bytes(task.model, key[1], rows, task.kfold) + slots = candidate_slots( + task.model, key[1], rows, task.kfold, per_candidate=per_candidate + ) + self._fused_shapes[key] = (key, per_candidate, slots) + return self._fused_shapes[key] + + def _submit_fused_task_batches(self, node_ids: List[str]) -> None: + idle = all(isinstance(u, _TaskBatchUnit) for u in self._job_units.values()) + now = time.perf_counter() + groups = {} + for node_id in node_ids: + self._fused_ready_since.setdefault(node_id, now) + shape = self._fused_batch_shape(self.scheduler.mapping[node_id]) + groups.setdefault(shape, []).append(node_id) + for ((task_idx, _), per_candidate, slots), members in sorted(groups.items()): + task = self._tasks[task_idx] + gpu_bytes = slots * per_candidate + _FUSED_CONTEXT_BYTES + waited = now - min(self._fused_ready_since[n] for n in members) + for start in range(0, len(members), slots): + if len(members) - start < slots and not idle and waited < _FUSED_WAIT_S: + break + gpu_id = self.scheduler.reserve_gpu( + gpu_bytes, + lambda g: self._pool.has_idle_worker_for(g), + force=not self._job_units, + ) + if gpu_id is None: + return + batch = members[start : start + slots] + items, retained = [], [] + for node_id in batch: + node = self.scheduler.mapping[node_id] + parent_ids = self.scheduler.get_valid_parents(node_id) + payload = ( + self._get_cached_parent(parent_ids[0]) + if parent_ids + else self._modalities[0] + ) + self._task_results[node_id] = ResultEntry( + dag=self._get_dag_from_node_ids(node_id), + representation_time=payload.transform_time, + ) + retained += self._retain_for_submit(parent_ids, payload) + items.append((node_id, payload, node.aggregation)) + self._fused_ready_since.pop(node_id, None) + self.scheduler.begin_execution(node_id) + self.scheduler.move_to_running(node_id) + job_id = self._pool.submit("task_batch", (task, items), gpu_id=gpu_id) + self._job_units[job_id] = _TaskBatchUnit(batch, gpu_id, gpu_bytes) + self._job_retained_shm[job_id] = retained + def _fill_pipeline(self) -> None: if not self._pool.has_idle_worker: return ready = self.scheduler.get_runnable().copy() + fused = self._fused_task_nodes(ready) + if fused: + ready = [e for e in ready if isinstance(e, list) or e not in set(fused)] def _gpu_id(entry): node_id = entry[0] if isinstance(entry, list) else entry @@ -867,6 +1008,8 @@ def _gpu_id(entry): entry, allow_gpu_worker_for_cpu=allow_gpu_worker_for_cpu, ) + if fused: + self._submit_fused_task_batches(fused) self._release_blocking_models() def _get_cached_parent(self, parent_id: str): @@ -904,15 +1047,19 @@ def _process_result(self, jr) -> None: if unit is None: return if isinstance(unit, _ReleaseUnit): - # A failed release is not retried; forgetting the model at - # worst under-counts it until the worker's next node. retained_bytes = jr.value["retained_gpu_bytes"] if jr.ok else 0 self._set_worker_model(unit.pid, retained_bytes, None) return + if isinstance(unit, _TaskBatchUnit): + self.scheduler.release_gpu(unit.gpu_id, unit.gpu_bytes) if not jr.ok: self._handle_job_failure(unit, jr) return - if isinstance(unit, _BatchUnit): + if isinstance(unit, _TaskBatchUnit): + for node_id, value in jr.value["results"].items(): + self._handle_node_success(node_id, value) + self._requeue_failed(jr.value["failed_nodes"]) + elif isinstance(unit, _BatchUnit): self._handle_batch_success(jr.value) else: self._handle_node_success(unit.node_id, jr.value) @@ -1020,7 +1167,10 @@ def _handle_batch_success(self, value: Dict[str, Any]) -> None: shm_bytes=info.get("shm_bytes", 0), ) - for node_id, reason in value.get("failed_nodes", {}).items(): + self._requeue_failed(value.get("failed_nodes", {})) + + def _requeue_failed(self, failed_nodes: Dict[str, str]) -> None: + for node_id, reason in failed_nodes.items(): if not self._requeue_or_give_up(node_id, reason): self._release_parents(node_id) @@ -1265,8 +1415,9 @@ def run(self) -> Dict[str, Any]: self._fill_pipeline() if not self._job_units: continue - jr = self._pool.wait() - self._process_result(jr) + jr = self._pool.wait(_FUSED_WAIT_S if self._fused_ready_since else None) + if jr is not None: + self._process_result(jr) finally: if self._owns_pool: self._pool.shutdown() diff --git a/src/main/python/systemds/scuro/drsearch/node_scheduler.py b/src/main/python/systemds/scuro/drsearch/node_scheduler.py index 5498ad98c7a..77ddc9f9557 100644 --- a/src/main/python/systemds/scuro/drsearch/node_scheduler.py +++ b/src/main/python/systemds/scuro/drsearch/node_scheduler.py @@ -36,7 +36,7 @@ ) from systemds.scuro.drsearch.worker_pool import default_gpu_slots_per_device from systemds.scuro.utils.memory_utility import gpu_memory_info -from systemds.scuro.utils.static_variables import DEBUG +from systemds.scuro.utils.static_variables import DEBUG, deterministic_enabled _MAX_GPU_SCHEDULE_ATTEMPTS = int(os.environ.get("SCURO_MAX_GPU_SCHEDULE_ATTEMPTS", "3")) @@ -241,13 +241,6 @@ def get_runnable(self) -> List[RepresentationNode]: def _get_runnable_nodes(self) -> List[str]: runnable_nodes = list(self._candidates) - - # Upward rank always puts a producer ahead of its consumers, so without - # a counterweight every large intermediate is materialized before the - # tasks that would release it. Once the cache holds a large share of - # the budget (or, with exclusive admission, the budget is exceeded), - # nodes whose completion does not grow the cache (task roots, or last - # children releasing at least their own output) run first. draining_first = self._cache_draining() def _score(node_id: str): @@ -323,6 +316,27 @@ def requeue_node(self, node_id: str) -> None: self._release_execution_memory(node_id, self.mapping[node_id].gpu_id) self._candidates.add(node_id) + def reserve_gpu(self, memory_needed: int, usable, force: bool = False): + for gpu_id in sorted( + self.gpu_devices, key=lambda g: -self._free_gpu_bytes(g, {}) + ): + if not usable(gpu_id): + continue + if self.gpu_slots_in_use[gpu_id] >= self.gpu_slots_per_device: + continue + if not force and self._free_gpu_bytes(gpu_id, {}) < memory_needed: + continue + self.gpu_slots_in_use[gpu_id] += 1 + self.memory_stats["gpu_in_use"][gpu_id] += memory_needed + return gpu_id + return None + + def release_gpu(self, gpu_id: int, memory_needed: int) -> None: + self.gpu_slots_in_use[gpu_id] = max(0, self.gpu_slots_in_use[gpu_id] - 1) + self.memory_stats["gpu_in_use"][gpu_id] = max( + 0, self.memory_stats["gpu_in_use"][gpu_id] - memory_needed + ) + def raise_gpu_estimate(self, node_id: str, observed_bytes: int) -> None: cpu_mem, gpu_mem = self.node_resources[node_id] self.node_resources[node_id] = (cpu_mem, max(gpu_mem, int(observed_bytes))) @@ -387,12 +401,8 @@ def is_finished(self) -> bool: return False - def get_valid_parents(self, node_id: str) -> bool: - parents = [] - for parent_id in self.parents[node_id]: - if parent_id not in self.leaves: - parents.append(parent_id) - return parents + def get_valid_parents(self, node_id: str) -> List[str]: + return [pid for pid in self.mapping[node_id].inputs if pid not in self.leaves] def get_children(self, node_id: str) -> List[str]: return list(self.children[node_id]) @@ -573,6 +583,11 @@ def _check_memory_constraints( if release_gpu is not None: self.release_requests.add(release_gpu) return False, None + if deterministic_enabled() and any( + gpu_mem <= self.memory_budget["gpu"][gpu] + for gpu in self.gpu_devices + ): + 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: diff --git a/src/main/python/systemds/scuro/drsearch/operator_registry.py b/src/main/python/systemds/scuro/drsearch/operator_registry.py index f8a83b28e59..f62d5192832 100644 --- a/src/main/python/systemds/scuro/drsearch/operator_registry.py +++ b/src/main/python/systemds/scuro/drsearch/operator_registry.py @@ -43,7 +43,14 @@ "tfidf", "word2vec", ), - ModalityType.IMAGE: ("clip", "color_histogram", "openface", "resnet", "vgg"), + ModalityType.IMAGE: ( + "clip", + "color_histogram", + "image_bind", + "openface", + "resnet", + "vgg", + ), ModalityType.VIDEO: ( "openface", "clip", @@ -72,7 +79,6 @@ "average", "concatenation", "hadamard", - "lstm", "max", "multimodal_attention_fusion", "sum", diff --git a/src/main/python/systemds/scuro/drsearch/ranking.py b/src/main/python/systemds/scuro/drsearch/ranking.py index 7b92ab0b130..ee28359a045 100644 --- a/src/main/python/systemds/scuro/drsearch/ranking.py +++ b/src/main/python/systemds/scuro/drsearch/ranking.py @@ -35,6 +35,18 @@ def _dag_size(entry) -> int: return 0 +def _tie_key(entry): + try: + pipeline_id = entry.dag.pipeline_id() + except Exception: + pipeline_id = "" + return _dag_size(entry), pipeline_id + + +def tie_broken(entries) -> list: + return sorted(range(len(entries)), key=lambda i: _tie_key(entries[i])) + + def rank_by_robustness( entries: Iterable, *, @@ -113,7 +125,7 @@ def variation_of(entry): else: keys = [(True, 0.0, float(s)) for s in robustness] - local_indices = sorted(range(len(entries)), key=lambda i: keys[i], reverse=True) + local_indices = sorted(tie_broken(entries), key=lambda i: keys[i], reverse=True) sorted_entries = [entries[i] for i in local_indices] sorted_indices = [original_indices[i] for i in local_indices] @@ -183,19 +195,7 @@ def safe_normalize(values, vmin, vmax): else: setattr(entry, score_attr, score) - sorted_entries = sorted( - entries, - key=lambda e: e.tradeoff_score if hasattr(e, "tradeoff_score") else 0.0, - reverse=True, - ) - - sorted_indices = [ - i - for i, _ in sorted( - enumerate(entries), - key=lambda pair: pair[1].tradeoff_score if pair is not None else None, - reverse=True, - ) - ] + sorted_indices = sorted(tie_broken(entries), key=lambda i: scores[i], reverse=True) + sorted_entries = [entries[i] for i in sorted_indices] return sorted_entries, sorted_indices diff --git a/src/main/python/systemds/scuro/drsearch/split_aware_evaluator.py b/src/main/python/systemds/scuro/drsearch/split_aware_evaluator.py index aead3e9dfaf..9e47f1a960b 100644 --- a/src/main/python/systemds/scuro/drsearch/split_aware_evaluator.py +++ b/src/main/python/systemds/scuro/drsearch/split_aware_evaluator.py @@ -19,11 +19,13 @@ # # ------------------------------------------------------------- import copy +import os import time from dataclasses import dataclass import numpy as np +from systemds.scuro.drsearch.trial_result_cache import TrialResultCache, _freeze from systemds.scuro.models.training_stats import TrainingWorkAccumulator from systemds.scuro.drsearch.representation_dag import ( RepresentationDag, @@ -53,6 +55,17 @@ class SplitEvaluationResult: task_time: float +_REFIT_CACHE = None + + +def _refit_cache() -> TrialResultCache: + global _REFIT_CACHE + if _REFIT_CACHE is None: + gb = float(os.environ.get("SCURO_LEAF_REFIT_CACHE_GB", "4")) + _REFIT_CACHE = TrialResultCache(int(gb * 1024**3), min_compute_seconds=0) + return _REFIT_CACHE + + def dag_requires_fitting(dag: RepresentationDag, modalities=()) -> bool: fold_ids = { str(m.modality_id) for m in modalities if getattr(m, "fold_source", None) @@ -149,16 +162,6 @@ def _apply_stateless(self, operation, inputs, node): @staticmethod def _align_partition_shapes(result): - """Zero-pad every partition's output to one common instance shape. - - Stateless operators pad to the largest instance of whatever they are - given (e.g. DynamicWindow pads each window to the partition's longest - clip), so executing them per partition yields different feature widths - for the fit and evaluation partitions, and a model fitted on one cannot - score the other. Padding to the elementwise maximum across partitions - reproduces the shapes of executing on all of their instances at once, - which is how the search computes representations. Labels play no part. - """ if len(result) < 2 or not all(isinstance(r, IndexedModality) for r in result): return result arrays = [] @@ -316,11 +319,23 @@ def execute_node(node_id): def _refit_leaf(self, source, partitions): source_dag, raw_modality = source + key = ( + _freeze(source_dag.to_spec()), + str(raw_modality.modality_id), + getattr(self.task.model, "name", None), + tuple(item.instance_ids for item in partitions), + ) + cached = _refit_cache().get(key) + if cached is not None: + return cached + start = time.perf_counter() evaluator = SplitAwareDAGEvaluator(source_dag, raw_modality, self.task) (full,) = evaluator._indexed_modalities() - return evaluator._execute_partitions( + result = evaluator._execute_partitions( *((full.subset(item.instance_ids),) for item in partitions) ) + _refit_cache().put(key, result, time.perf_counter() - start) + return result def _indexed_modalities(self): full = [] diff --git a/src/main/python/systemds/scuro/drsearch/task.py b/src/main/python/systemds/scuro/drsearch/task.py index 2c8aa4617b2..6cc1ec37b5c 100644 --- a/src/main/python/systemds/scuro/drsearch/task.py +++ b/src/main/python/systemds/scuro/drsearch/task.py @@ -22,8 +22,14 @@ import os import time from typing import List -from systemds.scuro.utils.static_variables import get_global_seed +from systemds.scuro.utils.static_variables import get_global_seed, get_train_seed from systemds.scuro.models.model import Model +from systemds.scuro.models.fused_mlp import ( + candidate_slots, + fit_folds, + fused_training_enabled, + supports_fused_training, +) from systemds.scuro.models.training_stats import TrainingWorkAccumulator import numpy as np from sklearn.model_selection import ( @@ -40,6 +46,11 @@ ) +def _as_features(X): + X = np.asarray(X, dtype=np.float32) + return X.reshape(len(X), -1) + + class PerformanceMeasure: def __init__(self, name, metrics, higher_is_better=True): self.average_scores = None @@ -199,6 +210,8 @@ def estimate_peak_memory_bytes(self, input_stats): uses_gpu = getattr(self.model, "uses_gpu", None) if uses_gpu is None: uses_gpu = model_peak_memory_gpu > 0 + if fused_training_enabled() and supports_fused_training(self.model): + uses_gpu = False gpu_peak_bytes = int(model_peak_memory_gpu * 1.4) if uses_gpu else 0 if uses_gpu and gpu_peak_bytes <= 0: @@ -494,6 +507,71 @@ def run(self, data): self.test_scores.compute_averages(), ] + def run_many(self, datas, device): + model = self.create_model() + rows = list( + dict.fromkeys(list(self.train_indices) + list(self.test_indices or [])) + ) + local = {row: i for i, row in enumerate(rows)} + splits = [ + ([local[i] for i in fit], [local[i] for i in val]) + for fit, val in zip(self.cv_train_indices, self.cv_val_indices) + ] + test = [local[i] for i in self.test_indices or []] + labels = np.asarray(self._gather_by_indices(self.labels, rows)) + prepare = getattr(model, "prepare_features", None) or _as_features + + results, by_width = [None] * len(datas), {} + for c, data in enumerate(datas): + try: + X = prepare(self._gather_by_indices(data, rows)) + if not np.isfinite(X).all(): + raise ValueError("non-finite input features") + by_width.setdefault(X.shape[1], []).append((c, X)) + except Exception as e: + results[c] = e + + seed = getattr(model, "seed", None) + seed = get_train_seed(get_global_seed() if seed is None else seed) + outputs = {} + start = time.perf_counter() + for width, group in by_width.items(): + slots = candidate_slots(model, width, len(rows), self.kfold) + for i in range(0, len(group), slots): + chunk = group[i : i + slots] + fitted = fit_folds( + model, [X for _, X in chunk], labels, splits, seed, device + ) + outputs.update(zip([c for c, _ in chunk], fitted)) + per_fit_s = (time.perf_counter() - start) / max(1, len(outputs) * self.kfold) + + for c, out in outputs.items(): + if not np.isfinite(out).all(): + results[c] = ValueError("training diverged: non-finite model output") + continue + self._reset_params() + for k, (fit, val) in enumerate(splits): + for measure, idx in ( + (self.train_scores, fit), + (self.val_scores, val), + (self.test_scores, test), + ): + measure.add_scores( + model.scores_from_output(out[k][idx], labels[idx])[0] + ) + self.last_run_timing = { + "train_time_per_fold_s": [per_fit_s] * self.kfold, + "train_time_mean_s": per_fit_s, + "n_test_instances": len(test), + "fused_batch_candidates": len(outputs), + } + results[c] = [ + self.train_scores.compute_averages(), + self.val_scores.compute_averages(), + self.test_scores.compute_averages(), + ] + return results + def fit_once_and_time_inference( self, data, latency_repeats: int = 200, latency_warmup: int = 20 ): diff --git a/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py b/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py index 431dc6666f2..f9331288698 100644 --- a/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py +++ b/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py @@ -29,7 +29,13 @@ from functools import lru_cache from systemds.scuro.modality.type import ModalityType -from systemds.scuro.drsearch.node_executor import NodeExecutor, ResultEntry +from systemds.scuro.drsearch.node_executor import ( + NodeExecutor, + ResultEntry, + _lazy_leaf_budget_fraction, +) +from systemds.scuro.dataloader.base_loader import LazyFileSequence +from systemds.scuro.utils.memory_utility import cpu_memory_budget_bytes from systemds.scuro.drsearch.split_aware_evaluator import ( SplitAwareDAGEvaluator, dag_requires_fitting, @@ -48,6 +54,7 @@ from systemds.scuro.drsearch.algebra import AlgebraConfig from systemds.scuro.drsearch.operator_registry import Registry from systemds.scuro.representations.window_aggregation import WindowAggregation +from systemds.scuro.representations.lstm import LSTM from systemds.scuro.utils.checkpointing import CheckpointManager from systemds.scuro.drsearch.representation_dag import ( RepresentationDag, @@ -413,8 +420,14 @@ def _process_modality(self, modality, skip_remaining: int = 0, scheduler=None): if fitted_dags: fitted_task_dags = self._expand_dags_with_task_roots(fitted_dags) for dag in fitted_dags: + eval_dag, eval_modality, encode_time = self._encode_below_lstm( + dag, modality + ) for task_idx, task in enumerate(self.tasks): - evaluation = SplitAwareDAGEvaluator(dag, modality, task).evaluate() + evaluation = SplitAwareDAGEvaluator( + eval_dag, eval_modality, task + ).evaluate() + evaluation.representation_time += encode_time result_dag = next( task_dag for task_dag in fitted_task_dags @@ -544,7 +557,7 @@ def add_dimensionality_reduction_operators(self, builder, current_node_id, rep_d dimensionality_reduction_operators = ( self._get_dimensionality_reduction_operators(modality_type) ) - + if self._dag_output_has_time_axis(rep_dag): for dimensionality_reduction_op in dimensionality_reduction_operators: self.pruned.append( @@ -592,7 +605,7 @@ def _build_modality_dag( if dimensionality_reduction_dags is not None: dags.extend(dimensionality_reduction_dags) - if operator.needs_context: + if operator.needs_context_for(modality.modality_type): context_operators = self._get_context_operators(modality.modality_type) for context_op in context_operators: if operator.initial_context_length is not None: @@ -709,29 +722,51 @@ def _dag_output_has_time_axis(self, dag: RepresentationDag) -> bool: def _aggregation_needed(self, dag: RepresentationDag) -> bool: return len(self._node_output_stats(dag)[dag.root_node_id].output_shape) > 1 + def _is_vector_sequence(self, dag: RepresentationDag) -> bool: + return len(self._node_output_stats(dag)[dag.root_node_id].output_shape) == 2 + def add_aggregation_operator(self, builder, dags): + aggregate = self._tasks_require_same_dims and self.expected_dimensions == 1 new_dags = [] - if self._tasks_require_same_dims and self.expected_dimensions == 1: - aggregated_dags = [] - for dag in dags: - if self._aggregation_needed(dag): - agg_op = AggregatedRepresentation( - target_dimensions=self.expected_dimensions, - aggregate_leading=True, - ) - agg_node_id = builder.create_operation_node( - agg_op.__class__, - [dag.root_node_id], - agg_op.get_current_parameters(), - ) - aggregated_dags.append(builder.build(agg_node_id, dag.dag_id)) - else: - aggregated_dags.append(dag) - new_dags = aggregated_dags - else: - new_dags = dags + for dag in dags: + if aggregate and self._aggregation_needed(dag): + agg_op = AggregatedRepresentation( + target_dimensions=self.expected_dimensions, + aggregate_leading=True, + ) + agg_node_id = builder.create_operation_node( + agg_op.__class__, + [dag.root_node_id], + agg_op.get_current_parameters(), + ) + new_dags.append(builder.build(agg_node_id, dag.dag_id)) + else: + new_dags.append(dag) + if self._is_vector_sequence(dag): + new_dags.append(self._lstm_aggregation(builder, dag.root_node_id)) return new_dags + @staticmethod + def _lstm_aggregation(builder, node_id): + lstm_node_id = builder.create_operation_node( + LSTM, [node_id], LSTM().get_current_parameters() + ) + return builder.build(lstm_node_id) + + @staticmethod + def _encode_below_lstm(dag, modality): + root = dag.get_node_by_id(dag.root_node_id) + below = RepresentationDag(dag.nodes, root.inputs[0]) + if root.operation is not LSTM or dag_requires_fitting(below): + return dag, modality, 0.0 + start = time.perf_counter() + encoded = below.execute([modality]) + leaf = RepresentationNode("lstm_input", None, [], modality.modality_id) + lstm = copy.copy(root) + lstm.inputs = [leaf.node_id] + dag = RepresentationDag([leaf, lstm], lstm.node_id, dag.dag_id) + return dag, encoded, time.perf_counter() - start + def default_context_operators( self, modality, builder, leaf_id, rep_dag, apply_context_to_leaf=False ): @@ -743,6 +778,19 @@ def default_context_operators( and modality.modality_type != ModalityType.IMAGE ): context_operators = self._get_context_operators(modality.modality_type) + if self._leaf_is_vector(modality): + for context_op in context_operators: + self.pruned.append( + { + "operation": context_op.__name__, + "input": "leaf", + "reason": ( + "leaf is a rank-1 feature vector " + "with no time axis to window over" + ), + } + ) + context_operators = [] for context_op in context_operators: context_node_id = builder.create_operation_node( context_op, @@ -785,6 +833,16 @@ def default_context_operators( return dags + @staticmethod + def _leaf_is_vector(modality) -> bool: + stats = getattr(modality, "stats", None) + shape = getattr(stats, "output_shape", None) + return ( + shape is not None + and len(shape) == 1 + and getattr(stats, "max_length", None) is None + ) + def _size_context_operator(self, context_operator_instance, output_length) -> bool: parameter = getattr(context_operator_instance, "granularity_parameter", None) kind = getattr(context_operator_instance, "granularity_kind", None) @@ -923,6 +981,9 @@ def combine_windowed_representations(self, builder, window_node_ids): dags.append( builder.build(chain_tip if summary_id is None else summary_id) ) + + if start == 0 and self._is_vector_sequence(builder.build(chain_tip)): + dags.append(self._lstm_aggregation(builder, chain_tip)) return dags def _summarize_windows(self, builder, node_id): @@ -1147,19 +1208,24 @@ def get_k_best_results( task_cache = None if not task_cache: cache = [] - for result in results: - if dag_requires_fitting(result.dag): - SplitAwareDAGEvaluator.materialize(modality) - transformed = SplitAwareDAGEvaluator( - result.dag, modality, task - ).fit_transform( - task.train_indices, - range(len(modality.data)), - ) - transformed.modality.fold_source = (result.dag, modality) - cache.append(transformed.modality) - else: - cache.append(result.dag.execute([modality])) + lazy_data = self._decode_lazy_leaf_once(modality, results) + try: + for result in results: + if dag_requires_fitting(result.dag): + SplitAwareDAGEvaluator.materialize(modality) + transformed = SplitAwareDAGEvaluator( + result.dag, modality, task + ).fit_transform( + task.train_indices, + range(len(modality.data)), + ) + transformed.modality.fold_source = (result.dag, modality) + cache.append(transformed.modality) + else: + cache.append(result.dag.execute([modality])) + finally: + if lazy_data is not None: + modality._data = lazy_data elif isinstance(task_cache, list): cache = task_cache @@ -1186,6 +1252,28 @@ def get_k_best_results( return results, cache + @staticmethod + def _decode_lazy_leaf_once(modality, results): + data = modality.data if modality.has_data() else None + if not isinstance(data, LazyFileSequence): + return None + if sum(not dag_requires_fitting(result.dag) for result in results) < 2: + return None + estimate = getattr(modality.data_loader, "estimate_decoded_bytes", None) + metadata = ( + modality.metadata if len(modality.metadata or []) == len(data) else None + ) + decoded_bytes = estimate(metadata) if callable(estimate) else None + budget = cpu_memory_budget_bytes(modality.modality_type.name) + if not decoded_bytes or decoded_bytes > budget * _lazy_leaf_budget_fraction(): + return None + arrays = data.materialize() + for array in arrays: + if hasattr(array, "flags"): + array.flags.writeable = False + modality._data = arrays + return data + def get_k_most_robust_results( self, modality, diff --git a/src/main/python/systemds/scuro/drsearch/worker_pool.py b/src/main/python/systemds/scuro/drsearch/worker_pool.py index f54578aca1f..b5afc9731fc 100644 --- a/src/main/python/systemds/scuro/drsearch/worker_pool.py +++ b/src/main/python/systemds/scuro/drsearch/worker_pool.py @@ -31,6 +31,7 @@ import torch from systemds.scuro.utils.memory_utility import cuda_oom_peak_bytes, is_cuda_oom +from systemds.scuro.utils.static_variables import scheduled_device _THREAD_ENV_VARS = ( "OMP_NUM_THREADS", @@ -127,7 +128,9 @@ def _worker_main( 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) + + with scheduled_device(local_gpu_id): + 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: cuda_oom = is_cuda_oom(e) @@ -388,14 +391,17 @@ def _enqueue(self, pid: int, kind: str, payload: Any, gpu_id: Optional[int]) -> self._workers[pid]["job_q"].put(job) return job_id - def wait(self) -> _JobResult: + def wait(self, timeout: Optional[float] = None) -> Optional[_JobResult]: + """The next job result; None if `timeout` seconds pass without one.""" while True: sentinel_to_pid = { w["process"].sentinel: pid for pid, w in self._workers.items() } ready = mp_connection.wait( - [self._result_q._reader, *sentinel_to_pid.keys()] + [self._result_q._reader, *sentinel_to_pid.keys()], timeout ) + if not ready: + return None if self._result_q._reader in ready: jr = self._result_q.get() entry = self._running.pop(jr.job_id, None) diff --git a/src/main/python/systemds/scuro/modality/transformed.py b/src/main/python/systemds/scuro/modality/transformed.py index 5806e3d4098..eb5aa7c4edd 100644 --- a/src/main/python/systemds/scuro/modality/transformed.py +++ b/src/main/python/systemds/scuro/modality/transformed.py @@ -66,6 +66,7 @@ def __init__( else True ) self.aggregate_dim = aggregate_dim + self.is_fused = getattr(modality, "is_fused", False) if modality.__class__.__name__ == "UnimodalModality": for m in self.metadata: @@ -190,6 +191,7 @@ def combine(self, other: Union[Modality, List[Modality]], fusion_method): fused_modality = TransformedModality( self, fusion_method, ModalityType.EMBEDDING ) + fused_modality.is_fused = True modalities = self.create_modality_list(other) start_time = time.time() fused_modality.data = fusion_method.transform(modalities) diff --git a/src/main/python/systemds/scuro/representations/clip.py b/src/main/python/systemds/scuro/representations/clip.py index 5be30b42f82..6301e817bba 100644 --- a/src/main/python/systemds/scuro/representations/clip.py +++ b/src/main/python/systemds/scuro/representations/clip.py @@ -384,6 +384,7 @@ def _pool_visual_output(self, output: torch.Tensor) -> torch.Tensor: @register_expensive_representation(ModalityType.TEXT) class CLIPText(UnimodalRepresentation): supports_aggregation_pushdown = True + max_context_words = 60 cache_in_worker = True layer_parameter = "layer_name" MODEL_BYTES = CLIP_MODEL_BYTES diff --git a/src/main/python/systemds/scuro/representations/fusion.py b/src/main/python/systemds/scuro/representations/fusion.py index 95bdec2c1a9..14c407aaa04 100644 --- a/src/main/python/systemds/scuro/representations/fusion.py +++ b/src/main/python/systemds/scuro/representations/fusion.py @@ -44,6 +44,7 @@ from systemds.scuro.utils.schema_helpers import get_shape DEFAULT_MAX_ALIGNMENT_PAD_RATIO = None +L2_NORM_EPS = 1e-12 class FusionInputTooLargeError(ValueError): @@ -160,7 +161,7 @@ def _prepare_modalities(self, modalities: List[Modality]): target_dimensions=1, aggregate_leading=True ) agg_modality = agg_operator.transform(modality) - mods.append(agg_modality if agg_modality else modality) + mods.append(self._l2_normalize(agg_modality if agg_modality else modality)) if self.needs_alignment: mods = [self._detached_copy(modality) for modality in mods] @@ -173,6 +174,31 @@ def _prepare_modalities(self, modalities: List[Modality]): return mods + @classmethod + def _l2_normalize(cls, modality: Modality) -> Modality: + if getattr(modality, "is_fused", False) or not modality.has_data(): + return modality + + data = modality.data + try: + arr = np.asarray(data) + except ValueError: + arr = np.empty(0, dtype=object) + normalized = cls._detached_copy(modality) + if arr.dtype != object and arr.ndim == 1 and len(modality.metadata or []) <= 1: + arr = arr.reshape(1, -1) + if arr.dtype != object and arr.size: + flat = arr.reshape(len(arr), -1) + dtype = np.result_type(arr.dtype, np.float32) + norms = np.linalg.norm(flat, axis=1, keepdims=True).astype(dtype) + normalized.data = (flat / np.maximum(norms, L2_NORM_EPS)).reshape(arr.shape) + else: + normalized.data = [ + np.asarray(x) / max(np.linalg.norm(np.asarray(x)), L2_NORM_EPS) + for x in data + ] + return normalized + @staticmethod def _detached_copy(modality: Modality) -> Modality: detached = copy.copy(modality) @@ -214,6 +240,8 @@ def _validate_indexed_modalities(inputs): def fit(self, inputs, labels=None) -> RepresentationFitState: if self.fit_requirement is not FitRequirement.SUPERVISED: return super().fit(inputs, labels) + if isinstance(inputs, IndexedModality): + inputs = [inputs] instance_ids = self._validate_indexed_modalities(inputs) if labels is None: raise ValueError(f"{self.name} requires labels for fitting") @@ -231,6 +259,8 @@ def fit(self, inputs, labels=None) -> RepresentationFitState: ) def transform_fitted(self, inputs, fit_state, **kwargs): + if isinstance(inputs, IndexedModality): + inputs = [inputs] instance_ids = self._validate_indexed_modalities(inputs) if ( not isinstance(fit_state, RepresentationFitState) @@ -246,6 +276,7 @@ def transform_fitted(self, inputs, fit_state, **kwargs): modalities[0], fit_state.payload, ModalityType.EMBEDDING ) transformed.data = fit_state.payload.apply_representation(modalities) + transformed.is_fused = len(inputs) > 1 return IndexedModality(transformed, instance_ids) def transform_data(self, modalities: List[Modality], indices=None): diff --git a/src/main/python/systemds/scuro/representations/image_bind.py b/src/main/python/systemds/scuro/representations/image_bind.py index 9d33fd352e7..2f1b2fe5d57 100644 --- a/src/main/python/systemds/scuro/representations/image_bind.py +++ b/src/main/python/systemds/scuro/representations/image_bind.py @@ -25,9 +25,12 @@ import numpy as np import torch +from PIL import Image from pytorchvideo import transforms as pv_transforms import torchaudio +from torchvision import transforms as tv_transforms +from systemds.scuro.dataloader.image_loader import ImageStats from systemds.scuro.drsearch.operator_registry import ( register_representation, register_expensive_representation, @@ -48,12 +51,15 @@ from systemds.scuro.utils.torch_dataset import TextDataset, TextSpanDataset -@register_representation([ModalityType.VIDEO, ModalityType.AUDIO, ModalityType.TEXT]) +@register_representation( + [ModalityType.IMAGE, ModalityType.VIDEO, ModalityType.AUDIO, ModalityType.TEXT] +) @register_expensive_representation( - [ModalityType.VIDEO, ModalityType.AUDIO, ModalityType.TEXT] + [ModalityType.IMAGE, ModalityType.VIDEO, ModalityType.AUDIO, ModalityType.TEXT] ) class ImageBind(UnimodalRepresentation): - """ImageBind (Girdhar et al., CVPR 2023) embeddings for video, audio and text. + """ImageBind (Girdhar et al., CVPR 2023) embeddings for images, video, audio + and text. Input preparation is implemented here; only the pretrained network comes from the optional ``imagebind`` package. Its code and weights are licensed @@ -77,15 +83,18 @@ class ImageBind(UnimodalRepresentation): _MEL_MEAN = -4.268 _MEL_STD = 9.138 _TEXT_CONTEXT_LENGTH = 77 - # ImageBind's text encoder uses the OpenAI CLIP BPE vocabulary. + _TEXT_WIDTH = 1024 + _TEXT_MLP_WIDTH = 4096 + _TEXT_HEADS = 16 + _TEXT_LAYERS = 24 _TOKENIZER_NAME = "openai/clip-vit-base-patch32" _CHECKPOINT_PATH = os.path.join(".checkpoints", "imagebind_huge.pth") _tokenizer = None supports_aggregation_pushdown = True + max_context_words = 60 cache_in_worker = True def __init__(self, output_file=None, batch_size=8, params=None): - # parameters = {"batch_size": [1, 2, 4, 8, 16, 32]} parameters = {} super().__init__("ImageBind", ModalityType.EMBEDDING, parameters) self.params = params @@ -98,6 +107,11 @@ def __init__(self, output_file=None, batch_size=8, params=None): self.model = None self.device = get_device() self._gpu_id = self.device.index + self.needs_context = True + self.initial_context_length = 55 + + def needs_context_for(self, modality_type): + return modality_type == ModalityType.TEXT @property def gpu_id(self): @@ -111,25 +125,92 @@ def gpu_id(self, gpu_id): self.model = self.model.to(self.device) def get_output_stats(self, input_stats) -> RepresentationStats: + chunked = isinstance(input_stats, RepresentationStats) and not ( + self.params and "_pushdown_aggregation" in self.params + ) return RepresentationStats( input_stats.num_instances, - (self._EMBEDDING_DIM,), - aggregate_dim=None, + ( + (input_stats.output_shape[0], self._EMBEDDING_DIM) + if chunked + else (self._EMBEDDING_DIM,) + ), + aggregate_dim=(0,) if chunked else 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 + shape = self.get_output_stats(input_stats).output_shape + return int(input_stats.num_instances * np.prod(shape) * self.data_type.itemsize) + + @staticmethod + def _is_text_stats(input_stats): + return hasattr(input_stats, "max_words") or isinstance( + input_stats, RepresentationStats + ) + + def _estimate_text_peak_memory_bytes(self, input_stats, model_bytes) -> dict: + num_instances = max(getattr(input_stats, "num_instances", 1), 1) + shape = tuple(getattr(input_stats, "output_shape", ()) or ()) + chunks_per_instance = ( + int(shape[0]) + if isinstance(input_stats, RepresentationStats) and len(shape) == 2 + else 1 + ) + num_sequences = num_instances * max(chunks_per_instance, 1) + itemsize = self.data_type.itemsize + batch_tokens = self.batch_size * self._TEXT_CONTEXT_LENGTH + + token_bytes = batch_tokens * np.dtype(np.int64).itemsize + activation_bytes = 4 * ( + batch_tokens * (3 * self._TEXT_WIDTH + self._TEXT_MLP_WIDTH) * itemsize + + self.batch_size + * self._TEXT_HEADS + * self._TEXT_CONTEXT_LENGTH**2 + * itemsize + ) + output_bytes = ( + num_sequences * self._EMBEDDING_DIM * itemsize + + self.estimate_output_memory_bytes(input_stats) + ) + cuda_overhead_bytes = 512 * 1024 * 1024 + gpu_peak = ( + model_bytes + + token_bytes + + activation_bytes + + output_bytes + + cuda_overhead_bytes + ) + + raw_text_bytes = num_instances * max( + int(getattr(input_stats, "max_length", 0) or 0), 1 + ) + tokenizer_bytes = 64 * 1024 * 1024 + safety_margin_bytes = 512 * 1024 * 1024 + cpu_peak = ( + model_bytes + + tokenizer_bytes + + raw_text_bytes + + num_sequences * self._TEXT_CONTEXT_LENGTH * 8 + + output_bytes + + safety_margin_bytes + ) + return {"cpu_peak_bytes": int(cpu_peak), "gpu_peak_bytes": int(gpu_peak)} 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 + if self._is_text_stats(input_stats): + return self._estimate_text_peak_memory_bytes(input_stats, model_bytes) + is_image = isinstance(input_stats, ImageStats) clip_bytes = ( - self._CLIPS_PER_VIDEO - * self._SPATIAL_CROPS - * self._FRAMES_PER_CLIP + ( + 1 + if is_image + else self._CLIPS_PER_VIDEO * self._SPATIAL_CROPS * self._FRAMES_PER_CLIP + ) * 3 * self._CROP_SIZE * self._CROP_SIZE @@ -142,7 +223,7 @@ def estimate_peak_memory_bytes(self, input_stats) -> dict: output_bytes = self.estimate_output_memory_bytes(input_stats) decoded_bytes = ( - getattr(input_stats, "max_length", 0) + (1 if is_image else 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) @@ -324,6 +405,34 @@ def _transform_video_data(self, samples, metadata): return torch.stack(video_outputs).to(self.device) + def _transform_image_data(self, samples): + transform = tv_transforms.Compose( + [ + tv_transforms.Resize( + self._CROP_SIZE, + interpolation=tv_transforms.InterpolationMode.BICUBIC, + ), + tv_transforms.CenterCrop(self._CROP_SIZE), + tv_transforms.ToTensor(), + tv_transforms.Normalize(mean=self._PIXEL_MEAN, std=self._PIXEL_STD), + ] + ) + image_outputs = [] + for sample in samples: + image = np.asarray(sample) + if image.ndim == 2: + image = np.repeat(image[..., None], 3, axis=-1) + if image.ndim != 3 or image.shape[-1] != 3: + raise ValueError( + "ImageBind image samples must have shape (height, width, 3) " + "or (height, width)" + ) + if not np.issubdtype(image.dtype, np.integer): + image = np.clip(np.rint(image * 255.0), 0, 255) + image_outputs.append(transform(Image.fromarray(image.astype(np.uint8)))) + + return torch.stack(image_outputs).to(self.device) + @classmethod def _get_tokenizer(cls): if cls._tokenizer is None: @@ -366,6 +475,8 @@ def _prepare_inputs(self, samples, metadata, modality_type): return IBModalityType.AUDIO, self._transform_audio_data(samples, metadata) if modality_type == ModalityType.VIDEO: return IBModalityType.VISION, self._transform_video_data(samples, metadata) + if modality_type == ModalityType.IMAGE: + return IBModalityType.VISION, self._transform_image_data(samples) raise ValueError(f"ImageBind does not support {modality_type}") def transform(self, modality, aggregation=None): diff --git a/src/main/python/systemds/scuro/representations/lstm.py b/src/main/python/systemds/scuro/representations/lstm.py index a5c1476860e..d182b496032 100644 --- a/src/main/python/systemds/scuro/representations/lstm.py +++ b/src/main/python/systemds/scuro/representations/lstm.py @@ -35,10 +35,7 @@ stats_bytes, ) -from systemds.scuro.drsearch.operator_registry import register_fusion_operator - -@register_fusion_operator() class LSTM(Fusion): fit_requirement = FitRequirement.SUPERVISED preserves_leading_axis = True @@ -98,7 +95,6 @@ def __init__( set_random_seeds() - @staticmethod def _shape_description(modality: Modality) -> str: data = modality.data @@ -193,7 +189,9 @@ def _prepare_data(self, modalities: List[Modality]) -> np.ndarray: aligned_modalities.append(data) concatenated_data = np.concatenate(aligned_modalities, axis=2) - return concatenated_data.astype(np.float32) + return np.nan_to_num( + concatenated_data.astype(np.float32), nan=0.0, posinf=0.0, neginf=0.0 + ) def _parameter_bytes(self, input_dim: int, num_classes: int) -> int: directions = 2 diff --git a/src/main/python/systemds/scuro/representations/multimodal_attention_fusion.py b/src/main/python/systemds/scuro/representations/multimodal_attention_fusion.py index 80d9fa1b946..9505386fb61 100644 --- a/src/main/python/systemds/scuro/representations/multimodal_attention_fusion.py +++ b/src/main/python/systemds/scuro/representations/multimodal_attention_fusion.py @@ -96,7 +96,6 @@ def __init__( set_random_seeds() - def _prepare_data(self, modalities: List[Modality]): inputs = {} input_dimensions = {} @@ -120,7 +119,10 @@ def _prepare_data(self, modalities: List[Modality]): input_dimensions[modality_name] = data.shape[2] # Feature dimension max_sequence_length = max(max_sequence_length, data.shape[1]) - inputs[modality_name] = torch.from_numpy(data.astype(np.float32)) + data = np.nan_to_num( + data.astype(np.float32), nan=0.0, posinf=0.0, neginf=0.0 + ) + inputs[modality_name] = torch.from_numpy(data) for modality_name, tensor in inputs.items(): if tensor.shape[1] < max_sequence_length: diff --git a/src/main/python/systemds/scuro/representations/openface.py b/src/main/python/systemds/scuro/representations/openface.py index bbb1e6240a6..26bd8b23b08 100644 --- a/src/main/python/systemds/scuro/representations/openface.py +++ b/src/main/python/systemds/scuro/representations/openface.py @@ -18,6 +18,7 @@ # under the License. # # ------------------------------------------------------------- +import logging import os import tempfile from contextlib import contextmanager @@ -95,6 +96,21 @@ def patched_init(self, args): self.log_dir = os.path.join(self.work_dir, "log") star_alignment.Alignment.__init__ = patched_init + original_init_instance = star_alignment.Alignment.init_instance + + def patched_init_instance(self): + root = logging.getLogger() + level, handlers = root.level, list(root.handlers) + try: + original_init_instance(self) + finally: + for handler in list(root.handlers): + if handler not in handlers: + root.removeHandler(handler) + handler.close() + root.setLevel(level) + + star_alignment.Alignment.init_instance = patched_init_instance _star_dirs_patched = True diff --git a/src/main/python/systemds/scuro/representations/timeseries_representations.py b/src/main/python/systemds/scuro/representations/timeseries_representations.py index c67fde8692e..f4e66bc3e6b 100644 --- a/src/main/python/systemds/scuro/representations/timeseries_representations.py +++ b/src/main/python/systemds/scuro/representations/timeseries_representations.py @@ -212,7 +212,7 @@ def __init__(self, params=None): def compute_feature(self, signal, axis=-1): result = np.asarray(stats.skew(signal, axis=axis)) zero_variance = np.std(signal, axis=axis) <= np.finfo(float).eps - return np.where(zero_variance, 0.0, result) + return np.where(zero_variance | ~np.isfinite(result), 0.0, result) @register_representation([ModalityType.TIMESERIES, ModalityType.PHYSIOLOGICAL]) @@ -257,7 +257,7 @@ def __init__(self, params=None): def compute_feature(self, signal, axis=-1): result = np.asarray(stats.kurtosis(signal, fisher=True, bias=True, axis=axis)) zero_variance = np.std(signal, axis=axis) <= np.finfo(float).eps - return np.where(zero_variance, 0.0, result) + return np.where(zero_variance | ~np.isfinite(result), 0.0, result) @register_representation([ModalityType.TIMESERIES, ModalityType.PHYSIOLOGICAL]) diff --git a/src/main/python/systemds/scuro/representations/unimodal.py b/src/main/python/systemds/scuro/representations/unimodal.py index 5c61629099b..a24e6c6af81 100644 --- a/src/main/python/systemds/scuro/representations/unimodal.py +++ b/src/main/python/systemds/scuro/representations/unimodal.py @@ -24,14 +24,9 @@ class UnimodalRepresentation(Representation): - # Parameter that selects which layer's activations an encoder returns. An - # encoder that sets it implements transform_layers, so a search over that - # parameter can obtain every layer from one pass over the data. layer_parameter = None - # Set by the executor (e.g. HPOExecutionCache.intermediates) for - # representations that share an intermediate result with other - # configurations; None computes it every time. intermediate_cache = None + max_context_words = None def __init__( self, name: str, output_modality_type, parameters=None, self_contained=True @@ -50,6 +45,9 @@ def __init__( self.needs_context = False self.initial_context_length = None + def needs_context_for(self, modality_type): + return self.needs_context + @abc.abstractmethod def transform(self, data, aggregation=None): # TODO: check if there is a way to time the transformation in here (needed for chunked execution) diff --git a/src/main/python/systemds/scuro/representations/utils.py b/src/main/python/systemds/scuro/representations/utils.py index 9ae5129eb15..5b27da3a1dd 100644 --- a/src/main/python/systemds/scuro/representations/utils.py +++ b/src/main/python/systemds/scuro/representations/utils.py @@ -428,7 +428,16 @@ def update(self, vectors, owner_ids, chunk_ids): 0, owner_ids, torch.ones_like(owner_ids, dtype=torch.long) ) if name in ("mean", "sum"): - self._values.index_add_(0, owner_ids, vectors) + owners, inverse = owner_ids.unique(return_inverse=True) + one_hot = torch.zeros( + (len(owners), len(owner_ids)), + device=vectors.device, + dtype=vectors.dtype, + ) + one_hot[inverse, torch.arange(len(owner_ids), device=vectors.device)] = 1 + with torch.autocast(device_type=vectors.device.type, enabled=False): + sums = one_hot @ vectors + self._values.index_add_(0, owners, sums) elif name in ("max", "min"): indices = owner_ids[:, None].expand_as(vectors) self._values.scatter_reduce_( diff --git a/src/main/python/systemds/scuro/representations/window_aggregation.py b/src/main/python/systemds/scuro/representations/window_aggregation.py index 115fb433d79..59aed6a0e26 100644 --- a/src/main/python/systemds/scuro/representations/window_aggregation.py +++ b/src/main/python/systemds/scuro/representations/window_aggregation.py @@ -693,7 +693,9 @@ def _max_window_size(self, max_length: int) -> int: length = max(int(max_length), num_windows) weights = np.geomspace(4, 256, num=num_windows) last_share = weights[-1] / np.sum(weights) - geometric_bound = int(np.ceil(num_windows + last_share * (length - num_windows))) + geometric_bound = int( + np.ceil(num_windows + last_share * (length - num_windows)) + ) return max(1, min(length - num_windows + 1, geometric_bound)) def _feature_shape(self, in_shape): diff --git a/src/main/python/systemds/scuro/utils/memory_utility.py b/src/main/python/systemds/scuro/utils/memory_utility.py index 5b80431add0..521c8d40237 100644 --- a/src/main/python/systemds/scuro/utils/memory_utility.py +++ b/src/main/python/systemds/scuro/utils/memory_utility.py @@ -32,6 +32,7 @@ from typing import List, Optional, Tuple import psutil import gc +from systemds.scuro.utils.static_variables import cuda_allowed, deterministic_enabled _CPU_MEMORY_BUDGET_FRACTION = float( os.environ.get("SCURO_CPU_MEMORY_BUDGET_FRACTION", "0.5") @@ -88,8 +89,10 @@ def get_best_gpu() -> Tuple[int, float]: def get_device_for_model( model: torch.nn.Module, memory_factor: float = 1.5 ) -> torch.device: - if not torch.cuda.is_available(): + if not cuda_allowed(): return torch.device("cpu") + if deterministic_enabled(): + return torch.device(f"cuda:{torch.cuda.current_device()}") model_mb = get_model_size_mb(model) required_mb = model_mb * memory_factor @@ -138,7 +141,7 @@ def compute_batch_size( memory_fraction: float = 0.8, ) -> int: - if device.type == "cpu": + if device.type == "cpu" or deterministic_enabled(): return max_batch_size gpu_idx = device.index if device.index is not None else 0 @@ -162,7 +165,7 @@ def compute_batch_size( def get_device(gpu_id: int = None): if gpu_id is not None: return torch.device(f"cuda:{gpu_id}") - return torch.device("cuda" if torch.cuda.is_available() else "cpu") + return torch.device("cuda" if cuda_allowed() else "cpu") def get_cpu_memory_mb(): diff --git a/src/main/python/systemds/scuro/utils/static_variables.py b/src/main/python/systemds/scuro/utils/static_variables.py index c163496f6f3..754f4518f5c 100644 --- a/src/main/python/systemds/scuro/utils/static_variables.py +++ b/src/main/python/systemds/scuro/utils/static_variables.py @@ -19,9 +19,10 @@ # # ------------------------------------------------------------- import os +from contextlib import contextmanager import numpy as np import torch -from typing import Tuple +from typing import Optional, Tuple PY_LIST_HEADER_BYTES = 56 PY_LIST_SLOT_BYTES = 8 @@ -29,7 +30,6 @@ DEBUG = False -# Read from the environment so spawned worker processes inherit the parent's seed. SEED = int(os.environ.get("SCURO_SEED", 42)) global_rng = np.random.default_rng(SEED) @@ -46,6 +46,40 @@ def get_global_seed() -> int: return SEED +def get_train_seed(default: int) -> int: + value = os.environ.get("SCURO_TRAIN_SEED", "").strip() + return int(value) if value else default + + +def deterministic_enabled() -> bool: + return os.environ.get("SCURO_DETERMINISTIC", "0").strip().lower() in ( + "1", + "true", + "yes", + "on", + ) + + +_CPU_JOB = False + + +@contextmanager +def scheduled_device(gpu_id: Optional[int]): + """Runs a worker job on the device it was scheduled for: gpu_id None is a + CPU job.""" + global _CPU_JOB + previous = _CPU_JOB + _CPU_JOB = gpu_id is None and deterministic_enabled() + try: + yield + finally: + _CPU_JOB = previous + + +def cuda_allowed() -> bool: + return not _CPU_JOB and torch.cuda.is_available() + + def get_seed(): return global_rng.integers(0, 1024) @@ -80,8 +114,10 @@ def get_best_gpu() -> Tuple[int, float]: def get_device_for_model( model: torch.nn.Module, memory_factor: float = 1.5 ) -> torch.device: - if not torch.cuda.is_available(): + if not cuda_allowed(): return torch.device("cpu") + if deterministic_enabled(): + return torch.device(f"cuda:{torch.cuda.current_device()}") model_mb = get_model_size_mb(model) required_mb = model_mb * memory_factor @@ -130,7 +166,7 @@ def compute_batch_size( memory_fraction: float = 0.8, ) -> int: - if device.type == "cpu": + if device.type == "cpu" or deterministic_enabled(): return max_batch_size gpu_idx = device.index if device.index is not None else 0 @@ -154,4 +190,4 @@ def compute_batch_size( def get_device(gpu_id: int = None): if gpu_id is not None: return torch.device(f"cuda:{gpu_id}") - return torch.device("cuda" if torch.cuda.is_available() else "cpu") + return torch.device("cuda" if cuda_allowed() else "cpu") diff --git a/src/main/python/systemds/scuro/utils/torch_dataset.py b/src/main/python/systemds/scuro/utils/torch_dataset.py index 7dcdea6f504..5f6d806868a 100644 --- a/src/main/python/systemds/scuro/utils/torch_dataset.py +++ b/src/main/python/systemds/scuro/utils/torch_dataset.py @@ -122,7 +122,6 @@ def __call__(self, frames, device) -> torch.Tensor: return torch.stack([self(f, device) for f in frames]) device = torch.device(device) - # Antialiased interpolation is not implemented on MPS. work_device = torch.device("cpu") if device.type == "mps" else device frames = frames.to(work_device, non_blocking=frames.is_pinned()) leading = frames.shape[:-3] @@ -140,14 +139,24 @@ def _process(self, frames: torch.Tensor) -> torch.Tensor: x = x[:, :3] if x.is_floating_point(): x = x.float().mul(255.0).round_().clamp_(0, 255).to(torch.uint8) - x = TF.resize( - x, [self.resize], interpolation=self.interpolation, antialias=True - ) + x = self._resize(x) x = TF.center_crop(x, [self.crop, self.crop]) x = x.to(self.dtype).div_(255.0) - # Frames arrive channel-last; hand the model the NCHW layout it got before. return TF.normalize(x, self.mean, self.std).contiguous() + def _resize(self, x: torch.Tensor) -> torch.Tensor: + try: + return TF.resize( + x, [self.resize], interpolation=self.interpolation, antialias=True + ) + except RuntimeError as e: + if x.device.type != "cuda" or "shared memory" not in str(e): + raise + resized = TF.resize( + x.cpu(), [self.resize], interpolation=self.interpolation, antialias=True + ) + return resized.to(x.device) + class CustomDataset(torch.utils.data.Dataset): def __init__(self, data, data_type, device, size=None, tf=None): diff --git a/src/main/python/tests/scuro/test_algebra_canonical.py b/src/main/python/tests/scuro/test_algebra_canonical.py index 9801d4817a3..e354c30154c 100644 --- a/src/main/python/tests/scuro/test_algebra_canonical.py +++ b/src/main/python/tests/scuro/test_algebra_canonical.py @@ -106,6 +106,23 @@ def pushed(c, params): ("concat_order_kept", C("a", "b"), C("b", "a"), ALL_LAWS, False, False), ("concat_perm", C("a", "b"), C("b", "a"), ALL_LAWS, True, True), ("concat_assoc", C(C("a", "b"), "c"), C("a", C("b", "c")), ALL_LAWS, False, True), + ( + "concat_perm_flattened", + C(C("a", "b"), "c"), + C("c", C("b", "a")), + ALL_LAWS, + True, + True, + ), + ( + "no_perm_below_sum", + S(C("a", "b"), "c"), + S(C("b", "a"), "c"), + ALL_LAWS, + True, + False, + ), + ("no_perm_below_agg", G(C("a", "b")), G(C("b", "a")), ALL_LAWS, True, False), ("window_concat_never_sorted", W("a", "b"), W("b", "a"), ALL_LAWS, True, False), ( "window_concat_assoc", diff --git a/src/main/python/tests/scuro/test_algebra_cse.py b/src/main/python/tests/scuro/test_algebra_cse.py index df86a55bd21..24f853ab8e2 100644 --- a/src/main/python/tests/scuro/test_algebra_cse.py +++ b/src/main/python/tests/scuro/test_algebra_cse.py @@ -49,22 +49,13 @@ def _nodes(builder, op): class TestLawAwareCSE(unittest.TestCase): - def test_window_concatenation_rotations(self): - for algebra, expected in ((LEGACY, 1), (LAW_AWARE, 2), (FULL, 2)): - with self.subTest(algebra=algebra): - optimizer, builder, windows = _builder(algebra) - optimizer.combine_windowed_representations(builder, windows) - self.assertEqual(len(_nodes(builder, Concatenation)), expected) - prevented = builder.algebra_stats["cse_unsafe_merges_prevented"] - self.assertEqual(prevented, int(algebra.cse_law_aware)) - - def test_commutative_fusion_merges_in_both_modes(self): + def test_only_commutative_fusion_merges_in_both_modes(self): cases = [ (Sum, {}, LEGACY, 1), (Sum, {}, LAW_AWARE, 1), - (Concatenation, {"preserve_leading_axis": False}, LEGACY, 1), + (Concatenation, {"preserve_leading_axis": False}, LEGACY, 2), (Concatenation, {"preserve_leading_axis": False}, LAW_AWARE, 2), - (Concatenation, {"preserve_leading_axis": False}, FULL, 1), + (Concatenation, {"preserve_leading_axis": False}, FULL, 2), ] for op, params, algebra, expected in cases: with self.subTest(op=op.__name__, algebra=algebra): diff --git a/src/main/python/tests/scuro/test_algebra_laws.py b/src/main/python/tests/scuro/test_algebra_laws.py index c0d55ddc6b6..5c0cccac7e2 100644 --- a/src/main/python/tests/scuro/test_algebra_laws.py +++ b/src/main/python/tests/scuro/test_algebra_laws.py @@ -58,16 +58,20 @@ def _operands(widths, dtypes, rank): def _fuse(op, params, *operands): modalities = [] for data in operands: + if isinstance(data, Modality): + modalities.append(data) + continue metadata = [ModalityType.EMBEDDING.create_metadata(d) for d in data] modality = TransformedModality( Modality(ModalityType.EMBEDDING, 0, metadata), "operand" ) modality.data = data modalities.append(modality) - return np.asarray(op(params).transform(modalities)) + return modalities[0].combine(modalities[1:], op(params)) def _assert_same(x, y, exact): + x, y = (np.asarray(getattr(v, "data", v)) for v in (x, y)) assert x.shape == y.shape if exact: np.testing.assert_array_equal(x, y) @@ -76,8 +80,8 @@ def _assert_same(x, y, exact): def _prepared(x, width): - # What _prepare_modalities hands to execute: hidden mean aggregation, then zero-padding. x = x.mean(axis=1) if x.ndim == 3 else x + x = x / np.linalg.norm(x, axis=-1, keepdims=True) return np.pad(x, ((0, 0), (0, width - x.shape[-1]))) @@ -98,7 +102,7 @@ def _avg_flatten(f, a, b, c, exact): def _concat_perm(f, a, b, c, exact): - _assert_same(f(a, b), np.roll(f(b, a), -b.shape[-1], axis=-1), exact=True) + _assert_same(f(a, b), np.roll(f(b, a).data, -b.shape[-1], axis=-1), exact=True) LAW_CHECKS = { @@ -134,7 +138,8 @@ def test_window_axis_concatenation_is_order_sensitive(self): a, b, _ = _operands(WIDTHS["equal"], DTYPES["float32"], rank=2) self.assertFalse( np.array_equal( - _fuse(Concatenation, params, a, b), _fuse(Concatenation, params, b, a) + _fuse(Concatenation, params, a, b).data, + _fuse(Concatenation, params, b, a).data, ) ) self.assertFalse({"comm", "concat_perm"} & Concatenation.algebraic_laws(params)) diff --git a/src/main/python/tests/scuro/test_algebra_plan_counts.py b/src/main/python/tests/scuro/test_algebra_plan_counts.py index 8861d89b766..5fa87adf437 100644 --- a/src/main/python/tests/scuro/test_algebra_plan_counts.py +++ b/src/main/python/tests/scuro/test_algebra_plan_counts.py @@ -56,11 +56,13 @@ def test_closed_forms_match_brute_force(self): def test_known_values(self): counts = count_fusion_plans(3, [Concatenation, Sum, Average], ALL_LAWS, True) self.assertEqual((counts["legacy_exhaustive"], counts["ordered"]), (648, 108)) - # comm alone divides the ordered trees by 2^(n-1) self.assertEqual( count_fusion_plans(4, COMM_ASSOC, COMM, True)["commutative"], 3240 // 8 ) self.assertIsNone(count_fusion_plans(5, MIXED, ALL_LAWS, True)["canonical"]) + self.assertEqual( + count_fusion_plans(3, [Concatenation], COMM, True)["commutative"], 6 + ) def test_search_space_sums_over_subsets(self): per_size = {n: count_fusion_plans(n, MIXED, ALL_LAWS, False) for n in (2, 3)} diff --git a/src/main/python/tests/scuro/test_algebra_stats.py b/src/main/python/tests/scuro/test_algebra_stats.py index 619c686e37b..9b62a9fbd2a 100644 --- a/src/main/python/tests/scuro/test_algebra_stats.py +++ b/src/main/python/tests/scuro/test_algebra_stats.py @@ -28,7 +28,6 @@ from systemds.scuro.representations.representation import RepresentationStats from tests.scuro.test_algebra_laws import _fuse -# Other test modules register fusion operators of their own. OPERATORS = [ op for op in Registry().get_fusion_operators() @@ -51,7 +50,7 @@ def test_predicted_stats_match_execution(self): for shape in OPERANDS[case] ] stats = [RepresentationStats(len(x), x.shape[1:]) for x in operands] - fused = _fuse(op, {}, *operands) + fused = np.asarray(_fuse(op, {}, *operands).data) self.assertEqual( tuple(op().get_output_stats(stats).output_shape), fused.shape[1:] ) diff --git a/src/main/python/tests/scuro/test_algebra_typecheck.py b/src/main/python/tests/scuro/test_algebra_typecheck.py index 4e34a3be095..4819cee368e 100644 --- a/src/main/python/tests/scuro/test_algebra_typecheck.py +++ b/src/main/python/tests/scuro/test_algebra_typecheck.py @@ -34,15 +34,21 @@ RepresentationDAGBuilder, ) from systemds.scuro.modality.type import ModalityType +from systemds.scuro.representations.aggregated_representation import ( + AggregatedRepresentation, +) from systemds.scuro.representations.bow import BoW from systemds.scuro.representations.concatenation import Concatenation from systemds.scuro.representations.mfcc import MFCC +from systemds.scuro.representations.mlp_averaging import MLPAveraging from systemds.scuro.representations.representation import RepresentationStats from systemds.scuro.representations.sum import Sum +from systemds.scuro.representations.window_aggregation import WindowAggregation from tests.scuro.test_multimodal_ga_optimizer import _FakeUnimodalResults, _make_task -TEXT, EMB = ModalityType.TEXT, ModalityType.EMBEDDING +TEXT, EMB, TS = ModalityType.TEXT, ModalityType.EMBEDDING, ModalityType.TIMESERIES SEQ, VEC = RepresentationStats(5, (4, 8)), RepresentationStats(5, (3,)) +RAW = SimpleNamespace(output_shape=(100,), max_length=100) # a loader's raw signal def _check(op, params, leaves, with_stats=True): @@ -74,6 +80,15 @@ def _check(op, params, leaves, with_stats=True): [("implicit", "hidden_aggregation")], ), (Concatenation, {"preserve_leading_axis": True}, [(EMB, SEQ), (EMB, SEQ)], []), + (WindowAggregation, {}, [(TS, None)], []), + (WindowAggregation, {}, [(TS, RAW)], []), + (AggregatedRepresentation, {}, [(TS, RAW)], []), + (WindowAggregation, {}, [(EMB, SEQ)], []), + (WindowAggregation, {}, [(EMB, VEC)], [("error", "rank")]), + (AggregatedRepresentation, {}, [(EMB, SEQ)], []), + (AggregatedRepresentation, {}, [(EMB, VEC)], [("error", "rank")]), + (MLPAveraging, {}, [(EMB, VEC)], []), + (MLPAveraging, {}, [(EMB, SEQ)], [("error", "rank")]), ] @@ -87,6 +102,17 @@ def test_shapes_unknown_without_leaf_stats(self): violations = _check(Sum, {}, [(EMB, SEQ), (EMB, VEC)], with_stats=False) self.assertEqual(violations, []) + def test_rank_follows_the_declared_output_rank(self): + builder = RepresentationDAGBuilder() + agg = builder.create_operation_node( + AggregatedRepresentation, [builder.create_leaf_node("hr", 0)], {} + ) + dag = builder.build(builder.create_operation_node(WindowAggregation, [agg], {})) + violations = typecheck(dag, {("hr", 0): TS}) + self.assertEqual( + [(v.kind, v.operator) for v in violations], [("rank", "WindowAggregation")] + ) + def test_builder_modes(self): for mode, raises, counted in [ ("off", False, 0), diff --git a/src/main/python/tests/scuro/test_data_loaders.py b/src/main/python/tests/scuro/test_data_loaders.py index 4affe2d3b7f..6cd35e930a3 100644 --- a/src/main/python/tests/scuro/test_data_loaders.py +++ b/src/main/python/tests/scuro/test_data_loaders.py @@ -75,9 +75,50 @@ def test_audio_loader_stats(self): self.assertIsInstance(stats, AudioStats) self.assertEqual(stats.num_instances, 2) - self.assertEqual(stats.sampling_rate, 22050) - self.assertEqual(stats.max_length, 44100) - self.assertAlmostEqual(stats.avg_length, (44100 * 2) / 2.0) + self.assertEqual(stats.sampling_rate, 16000) + self.assertEqual(stats.max_length, 32000) + self.assertAlmostEqual(stats.avg_length, 32000) + + def test_audio_loader_keeps_native_rate(self): + loader = AudioLoader( + self.data_generator.get_modality_path(ModalityType.AUDIO), + self.data_generator.indices, + ) + data, metadata = loader.load() + + for arr, md in zip(data, metadata): + self.assertEqual(md["frequency"], 16000) + self.assertEqual(md["length"], arr.shape[0]) + self.assertEqual(arr.shape[0], loader.stats.max_length) + + def test_audio_loader_resamples_on_request(self): + loader = AudioLoader( + self.data_generator.get_modality_path(ModalityType.AUDIO), + self.data_generator.indices, + sampling_rate=22050, + ) + data, metadata = loader.load() + + self.assertEqual(loader.stats.sampling_rate, 22050) + self.assertEqual(loader.stats.max_length, 44100) + for arr, md in zip(data, metadata): + self.assertEqual(md["frequency"], 22050) + self.assertEqual(arr.shape[0], 44100) + + def test_audio_loader_chunked_matches_full_load(self): + path = self.data_generator.get_modality_path(ModalityType.AUDIO) + full_data, full_metadata = AudioLoader(path, self.data_generator.indices).load() + chunked = AudioLoader(path, self.data_generator.indices, chunk_size=1) + + chunks = list(chunked.iter_loaded_chunks()) + + self.assertEqual(len(chunks), self.num_instances) + for (data, metadata, _), expected, expected_md in zip( + chunks, full_data, full_metadata + ): + self.assertEqual(len(data), 1) + np.testing.assert_array_equal(data[0], expected) + self.assertEqual(metadata[0]["frequency"], expected_md["frequency"]) def test_video_loader_loads_all_instances(self): loader = VideoLoader( diff --git a/src/main/python/tests/scuro/test_hp_tuner.py b/src/main/python/tests/scuro/test_hp_tuner.py index 8f7aa0b1284..1e0bad3991e 100644 --- a/src/main/python/tests/scuro/test_hp_tuner.py +++ b/src/main/python/tests/scuro/test_hp_tuner.py @@ -51,6 +51,14 @@ from systemds.scuro.modality.type import ModalityType from systemds.scuro.drsearch.hyperparameter_tuner import HyperparameterTuner +from systemds.scuro.drsearch.representation_dag import CSEAwareDAGBuilder +from systemds.scuro.representations.bert import Bert +from systemds.scuro.representations.clip import CLIPText +from systemds.scuro.representations.image_bind import ImageBind +from systemds.scuro.representations.text_context_with_indices import ( + OverlappingSplitIndices, + SentenceBoundarySplitIndices, +) from unittest.mock import patch @@ -200,6 +208,37 @@ def run_hp_for_modality( == 2 ) + def test_context_max_words_is_limited_by_the_encoder_text_window(self): + hp = HyperparameterTuner( + modalities=[], + tasks=[SimpleNamespace(model=SimpleNamespace(name="dummy_task"))], + optimization_results=None, + n_jobs=1, + ) + for context_op in (SentenceBoundarySplitIndices, OverlappingSplitIndices): + domain = context_op().parameters["max_words"] + for encoder, expected in ( + (ImageBind, [v for v in domain if v <= 60]), + (CLIPText, [v for v in domain if v <= 60]), + (Bert, domain), + ): + with self.subTest( + context=context_op.__name__, encoder=encoder.__name__ + ): + builder = CSEAwareDAGBuilder() + leaf_id = builder.create_leaf_node("text") + context_id = builder.create_operation_node( + context_op, [leaf_id], context_op(55).get_current_parameters() + ) + encoder_id = builder.create_operation_node( + encoder, [context_id], {} + ) + dag = builder.build(encoder_id) + params = hp._HyperparameterTuner__get_params_for_node( + dag.get_node_by_id(context_id), dag + ) + self.assertEqual(params["max_words"], expected) + def test_evaluate_configs_deduplicates_candidates(self): class DummyOptimizationResults: def get_k_best_results(self, modality, task, performance_metric_name): diff --git a/src/main/python/tests/scuro/test_modality_pad.py b/src/main/python/tests/scuro/test_modality_pad.py index b6c0001250c..1200e928084 100644 --- a/src/main/python/tests/scuro/test_modality_pad.py +++ b/src/main/python/tests/scuro/test_modality_pad.py @@ -96,8 +96,10 @@ def test_fusion_sum_aligns_mismatched_embedding_sizes(self): 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) + a = np.array([1.0, 2.0, 3.0, 4.0, 0.0, 0.0]) + b = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0]) + np.testing.assert_allclose( + fused[0], a / np.linalg.norm(a) + b / np.linalg.norm(b), rtol=1e-6 ) diff --git a/src/main/python/tests/scuro/test_multimodal_ga_optimizer.py b/src/main/python/tests/scuro/test_multimodal_ga_optimizer.py index e5eac1c5cfb..0affa8acaf2 100644 --- a/src/main/python/tests/scuro/test_multimodal_ga_optimizer.py +++ b/src/main/python/tests/scuro/test_multimodal_ga_optimizer.py @@ -23,6 +23,8 @@ import os import pickle import random +import subprocess +import sys import tempfile import unittest from types import SimpleNamespace @@ -34,6 +36,7 @@ DagGenome, MultimodalDeapOptimizer, _collect_leaf_indices, + _dispatch_genome_evaluation, _failure_fitness, _objective_value, _task_uses_gpu, @@ -595,6 +598,26 @@ def test_same_leaves_produces_structurally_valid_children(self): set(child.fusion_ops.keys()), set(optimizer._internal_paths(child.tree)) ) + def test_subtree_swap_keeps_every_leaf_exactly_once(self): + """Swapping subtrees with different leaf sets used to duplicate one + leaf and drop another from the child's tree.""" + optimizer, _, task = _make_optimizer( + n_modalities=4, reps_per_modality=2, min_modalities=4, max_modalities=4 + ) + optimizer.fusion_operators = [Concatenation, Average] + optimizer.rng = random.Random(0) + for _ in range(100): + g1 = optimizer._random_genome(task.model.name) + g2 = copy.deepcopy(g1) + g2.tree = optimizer._random_binary_tree(len(g2.leaves)) + g2.fusion_ops = {} + optimizer._assign_fusion_ops(g2.tree, g2.fusion_ops, "") + for child in optimizer._crossover_genomes(g1, g2): + self.assertEqual( + sorted(_collect_leaf_indices(child.tree)), + list(range(len(child.leaves))), + ) + def test_mismatched_leaves_still_recombines(self): """Regression test: the original crossover returned both parents completely unchanged whenever they didn't select the exact same @@ -628,6 +651,49 @@ def test_mismatched_leaves_still_recombines(self): self.assertEqual(c1.leaves, g1.leaves) self.assertEqual(c2.leaves, g2.leaves) + def test_mismatched_leaves_crossover_is_independent_of_hash_seed(self): + """Regression test: the op-only fallback iterated a set of path + strings while drawing from the rng, so every process (PYTHONHASHSEED) + paired the paths with different draws and two runs of the same seed + bred different children.""" + script = ( + "import json, random\n" + "from tests.scuro.test_multimodal_ga_optimizer import _make_optimizer\n" + "from systemds.scuro.representations.average import Average\n" + "from systemds.scuro.representations.concatenation import Concatenation\n" + "from systemds.scuro.representations.sum import Sum\n" + "opt, _, task = _make_optimizer(n_modalities=6, reps_per_modality=2,\n" + " min_modalities=4, max_modalities=6)\n" + "opt.fusion_operators = [Concatenation, Average, Sum]\n" + "opt.rng = random.Random(0)\n" + "children = []\n" + "for _ in range(100):\n" + " g1 = opt._random_genome(task.model.name)\n" + " g2 = opt._random_genome(task.model.name)\n" + " for c in opt._crossover_genomes(g1, g2):\n" + " children.append(sorted((p, op.__name__) for p, op in c.fusion_ops.items()))\n" + "print(json.dumps(children))\n" + ) + root = os.path.dirname( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + ) + outputs = set() + for seed in ("0", "1", "2", "3"): + env = dict(os.environ, PYTHONHASHSEED=seed) + env["PYTHONPATH"] = os.pathsep.join( + filter(None, [root, env.get("PYTHONPATH")]) + ) + result = subprocess.run( + [sys.executable, "-c", script], + cwd=root, + env=env, + capture_output=True, + text=True, + check=True, + ) + outputs.add(result.stdout.strip().splitlines()[-1]) + self.assertEqual(len(outputs), 1) + class TestGenomeSignature(unittest.TestCase): def test_signature_ignores_fusion_ops_dict_insertion_order(self): @@ -936,6 +1002,58 @@ def test_write_is_atomic_no_partial_file_on_failure(self): self.assertEqual(os.listdir(d), []) +class TestSearchReproducibility(unittest.TestCase): + def _run(self, first_task_offset, global_seed): + """Searches two tasks and returns the genomes evaluated per task. The + first task's scores are shifted by `first_task_offset`; the global + random module is seeded with `global_seed`.""" + tasks = [_make_task("first"), _make_task("second")] + modalities = [_FakeModality(f"m{i}") for i in range(4)] + reps = {m.modality_id: [object() for _ in range(3)] for m in modalities} + optimizer = MultimodalDeapOptimizer( + modalities, + _FakeUnimodalResults(reps), + tasks, + debug=False, + population_size=8, + generations=4, + early_stopping_patience=None, + random_seed=7, + ) + optimizer.fusion_operators = [Concatenation, Average, Sum] + evaluated = {"first": [], "second": []} + + def evaluate(population, task): + for ind in population: + key = ( + tuple(ind[0].leaves), + tuple( + sorted((p, c.__name__) for p, c in ind[0].fusion_ops.items()) + ), + ) + value = (hash(repr(key)) % 1000) / 1000.0 + if task.model.name == "first": + value = (value + first_task_offset) % 1.0 + ind.fitness.values = (value,) + evaluated[task.model.name].append(key) + + random.seed(global_seed) + with patch.object(optimizer, "_evaluate_population", side_effect=evaluate): + optimizer._optimize() + return evaluated + + def test_a_task_search_ignores_earlier_tasks_and_global_random(self): + a = self._run(first_task_offset=0.0, global_seed=0) + b = self._run(first_task_offset=0.37, global_seed=1) + self.assertNotEqual( + a["first"], b["first"], "test setup needs diverging first task" + ) + self.assertEqual(a["second"], b["second"]) + + def test_same_scores_give_the_same_search(self): + self.assertEqual(self._run(0.0, 0), self._run(0.0, 1)) + + class TestRealFusionIntegration(unittest.TestCase): """Exercises the actual RepresentationDag.execute() path end-to-end, including across process boundaries for the parallel case (a spawned diff --git a/src/main/python/tests/scuro/test_operator_registry.py b/src/main/python/tests/scuro/test_operator_registry.py index d60788ea3f1..87271d52a97 100644 --- a/src/main/python/tests/scuro/test_operator_registry.py +++ b/src/main/python/tests/scuro/test_operator_registry.py @@ -120,6 +120,20 @@ def test_video_representations_in_registry(self): ], ) + def test_image_representations_in_registry(self): + registry = Registry() + self.assertCountEqual( + registry.get_representations(ModalityType.IMAGE), + [ + ResNet, + OpenFace, + ImageBind, + ColorHistogram, + VGG19, + CLIPVisual, + ], + ) + def test_timeseries_representations_in_registry(self): registry = Registry() assert registry.get_representations(ModalityType.TIMESERIES) == [ diff --git a/src/main/python/tests/scuro/test_scheduler.py b/src/main/python/tests/scuro/test_scheduler.py index e46d59c15f3..cec30c47bb8 100644 --- a/src/main/python/tests/scuro/test_scheduler.py +++ b/src/main/python/tests/scuro/test_scheduler.py @@ -77,7 +77,7 @@ def __init__(self, dimensionality=1000, params=None): def get_output_stats(self, input_stats): return RepresentationStats( - input_stats.num_instances, (self.dimensionality) + input_stats.num_instances, (self.dimensionality,) ) def estimate_memory_bytes(self, input_stats): diff --git a/src/main/python/tests/scuro/test_transformer_text_aggregation.py b/src/main/python/tests/scuro/test_transformer_text_aggregation.py index ded26859917..f24b6cbd136 100644 --- a/src/main/python/tests/scuro/test_transformer_text_aggregation.py +++ b/src/main/python/tests/scuro/test_transformer_text_aggregation.py @@ -49,7 +49,25 @@ def to(self, device): return self -class _Tokenizer: +class _PadMixin: + def pad(self, features, **kwargs): + rows = [ + { + key: torch.as_tensor(feature[key]) + for key in ("input_ids", "attention_mask") + } + for feature in features + ] + padded_length = max(len(row["input_ids"]) for row in rows) + input_ids = torch.zeros((len(rows), padded_length), dtype=torch.long) + attention_mask = torch.zeros_like(input_ids) + for index, row in enumerate(rows): + input_ids[index, : len(row["input_ids"])] = row["input_ids"] + attention_mask[index, : len(row["attention_mask"])] = row["attention_mask"] + return _BatchEncoding(input_ids=input_ids, attention_mask=attention_mask) + + +class _Tokenizer(_PadMixin): def __call__(self, batch, **kwargs): ids = torch.tensor([int(text) for text in batch]) input_ids = ids.unsqueeze(1).repeat(1, 3) @@ -61,8 +79,8 @@ def __call__(self, batch, **kwargs): ) -class _BertModel: - def __call__(self, input_ids, attention_mask): +class _BertModel(torch.nn.Module): + def forward(self, input_ids, attention_mask): ids = input_ids[:, 0].float() hidden = torch.stack( ( @@ -75,7 +93,7 @@ def __call__(self, input_ids, attention_mask): return SimpleNamespace(last_hidden_state=hidden) -class _DynamicTokenizer: +class _DynamicTokenizer(_PadMixin): def __call__(self, batch, **kwargs): tokens = [[int(token) for token in text.split()] for text in batch] max_length = kwargs.get("max_length") @@ -103,11 +121,12 @@ def __call__(self, batch, **kwargs): ) -class _IntermediateBertModel: +class _IntermediateBertModel(torch.nn.Module): def __init__(self, representation): - self.representation = representation + super().__init__() + self.intermediate = torch.nn.Identity() - def __call__(self, input_ids, attention_mask): + def forward(self, input_ids, attention_mask): values = input_ids.float() hidden = torch.stack((values, values + 10), dim=-1) hidden = torch.where( @@ -115,8 +134,7 @@ def __call__(self, input_ids, attention_mask): hidden, torch.full_like(hidden, 1000), ) - self.representation.bert_output = hidden - return SimpleNamespace(last_hidden_state=hidden) + return SimpleNamespace(last_hidden_state=self.intermediate(hidden)) class _CLIPProcessor: @@ -128,25 +146,28 @@ def __call__(self, text, **kwargs): ) -class _CLIPTextModel: - def __init__(self, representation): - self.representation = representation +class _CLIPTextModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.intermediate = torch.nn.Identity() - def __call__(self, input_ids, attention_mask): + def forward(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, + return self.intermediate( + 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) + self.text_model = _CLIPTextModel() class TestTransformerTextAggregation(unittest.TestCase): @@ -288,7 +309,7 @@ def test_clip_intermediate_layer_uses_masked_mean_before_chunk_aggregation(self) 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): + def test_pushdown_skipped_when_plain_output_is_also_needed(self): builder = CSEAwareDAGBuilder() leaf_id = builder.create_leaf_node("transformer_pushdown") bert = Bert() @@ -312,17 +333,17 @@ def test_pushdown_keeps_shared_plain_and_aggregated_paths_distinct(self): pushdown_aggregation([plain_dag, aggregated_dag]) + # Bert runs once; its plain output feeds both DAGs and the aggregation + # is applied to it instead of being fused into a second Bert run. 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)) + aggregation_node = aggregated_dag.get_node_by_id(aggregation_id) + self.assertIs(aggregation_node.operation, AggregatedRepresentation) + self.assertEqual(aggregation_node.inputs, [bert_id]) + self.assertEqual(aggregation_node.parameters, aggregation_params) + self.assertIs(aggregated_dag.get_node_by_id(bert_id), plain_node) def test_clip_text_supports_aggregation_pushdown(self): builder = CSEAwareDAGBuilder() diff --git a/src/main/python/tests/scuro/test_unimodal_optimizer.py b/src/main/python/tests/scuro/test_unimodal_optimizer.py index 6e002c1688e..215b4712466 100644 --- a/src/main/python/tests/scuro/test_unimodal_optimizer.py +++ b/src/main/python/tests/scuro/test_unimodal_optimizer.py @@ -46,13 +46,20 @@ 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.image_bind import ImageBind +from systemds.scuro.representations.text_context_with_indices import ( + OverlappingSplitIndices, + SentenceBoundarySplitIndices, +) from systemds.scuro.representations.word2vec import W2V from systemds.scuro.representations.window_aggregation import Window +from systemds.scuro.representations.concatenation import Concatenation from tests.scuro.test_unimodal_representations import ( PHYSIOLOGICAL_REPRESENTATIONS, TIMESERIES_REPRESENTATIONS, ) from systemds.scuro.modality.unimodal_modality import UnimodalModality +from systemds.scuro.dataloader.base_loader import LazyFileSequence from tests.scuro.data_generator import ( ModalityRandomDataGenerator, TestDataLoader, @@ -63,8 +70,13 @@ from systemds.scuro.drsearch.representation_dag import ( CSEAwareDAGBuilder, RepresentationDag, + RepresentationDAGBuilder, pushdown_aggregation, ) +from systemds.scuro.drsearch.split_aware_evaluator import ( + SplitAwareDAGEvaluator, + dag_requires_fitting, +) from systemds.scuro.representations.aggregated_representation import ( AggregatedRepresentation, ) @@ -139,6 +151,111 @@ def test_unimodal_optimizer_for_text_modality(self): ) self.optimize_unimodal_representation_for_modality([text]) + def test_fitted_text_leaves_are_refit_per_fold_in_fusion(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 + ) + ) + task = self.tasks[0] + results = UnimodalResults([text], [task], k=2) + for accuracy, representation in ((0.9, BoW), (0.8, TfIdf)): + builder = RepresentationDAGBuilder() + leaf = builder.create_leaf_node(text.modality_id) + dag = builder.build( + builder.create_operation_node(representation, [leaf], {}) + ) + results.results[text.modality_id][task.model.name].append( + SimpleNamespace( + dag=dag, + val_score={"accuracy": accuracy}, + representation_time=0.0, + task_time=0.0, + ) + ) + _, leaves = results.get_k_best_results(text, task, "accuracy") + self.assertEqual(len(leaves), 2) + self.assertTrue(all(leaf.fold_source for leaf in leaves)) + + builder = RepresentationDAGBuilder() + fused = builder.build( + builder.create_operation_node( + Concatenation, + [ + builder.create_leaf_node(text.modality_id, 0), + builder.create_leaf_node(text.modality_id, 1), + ], + {}, + ) + ) + self.assertTrue(dag_requires_fitting(fused, leaves)) + + fit_sets = [] + fits = {cls: cls.fit for cls in (BoW, TfIdf)} + + def recording_fit(cls): + def fit(op, inputs, labels=None): + fit_sets.append(set(inputs.instance_ids)) + return fits[cls](op, inputs, labels=labels) + + return fit + + with patch.object(BoW, "fit", recording_fit(BoW)), patch.object( + TfIdf, "fit", recording_fit(TfIdf) + ): + SplitAwareDAGEvaluator(fused, leaves, task).evaluate() + SplitAwareDAGEvaluator(fused, leaves, task).evaluate() + + allowed = [set(split.fit_indices) for split in task.evaluation_splits] + self.assertEqual(len(fit_sets), 2 * len(allowed)) + self.assertTrue(all(fit_set in allowed for fit_set in fit_sets)) + + def test_k_best_dags_decode_a_lazy_leaf_once(self): + images, image_md = ModalityRandomDataGenerator().create_visual_modality( + self.num_instances + ) + task = self.tasks[0] + + def k_best(decoded_bytes): + loader = TestDataLoader( + self.indices, None, ModalityType.IMAGE, images, np.uint8, image_md + ) + loader.estimate_decoded_bytes = lambda metadata=None: decoded_bytes + image = UnimodalModality(loader) + decoded = [] + + def decode(name): + decoded.append(name) + return images[int(name)].copy() + + lazy = LazyFileSequence([str(i) for i in self.indices], decode) + image._data, image.metadata = lazy, image_md + results = UnimodalResults([image], [task], k=2) + for accuracy, bins in ((0.9, 16), (0.8, 32)): + builder = RepresentationDAGBuilder() + leaf = builder.create_leaf_node(image.modality_id) + dag = builder.build( + builder.create_operation_node( + ColorHistogram, [leaf], {"bins": bins} + ) + ) + results.results[image.modality_id][task.model.name].append( + SimpleNamespace(dag=dag, val_score={"accuracy": accuracy}) + ) + _, cache = results.get_k_best_results(image, task, "accuracy") + self.assertIs(image.data, lazy) + return len(decoded), [np.asarray(c.data) for c in cache] + + lazy_decodes, lazy_outputs = k_best(None) + once_decodes, once_outputs = k_best(1) + self.assertEqual(lazy_decodes, 2 * self.num_instances) + self.assertEqual(once_decodes, self.num_instances) + for lazy_output, once_output in zip(lazy_outputs, once_outputs): + np.testing.assert_array_equal(lazy_output, once_output) + def test_robust_results_ignore_non_finite_scores(self): modality = SimpleNamespace(modality_id="modality") task = SimpleNamespace(model=SimpleNamespace(name="task")) @@ -232,6 +349,55 @@ def test_bow_and_tfidf_require_dimensionality_reduction_before_task(self): task_input = dag.get_node_by_id(task_node.inputs[0]) self.assertIs(task_input.operation, MLPAveraging) + def test_dimensionality_reduction_is_not_applied_to_sequence_outputs(self): + video_data, video_md = ModalityRandomDataGenerator().create_visual_modality( + self.num_instances, 10, 10, 10 + ) + video = UnimodalModality( + TestDataLoader( + self.indices, None, ModalityType.VIDEO, video_data, np.float32, video_md + ) + ) + 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 + ) + ) + dimensionality_reduction_operators = {ModalityType.EMBEDDING: [MLPAveraging]} + for modality, representations in ((video, [ResNet]), (audio, [Spectral])): + with self.subTest(modality=modality.modality_type.name), patch.object( + Registry, + "_representations", + registry_for(modality.modality_type, representations), + ), patch.object( + Registry, + "_dimensionality_reduction_operators", + dimensionality_reduction_operators, + ): + optimizer = UnimodalOptimizer( + [modality], self.tasks, False, enable_checkpointing=False + ) + dags, _, _ = optimizer._build_execution_dags_for_modality(modality) + stats = {} + for dag in dags: + stats.update(optimizer._node_output_stats(dag)) + reduce_input_ranks = [ + len(stats[node.inputs[0]].output_shape) + for dag in dags + for node in dag.nodes + if node.operation is MLPAveraging + ] + self.assertEqual([rank for rank in reduce_input_ranks if rank > 1], []) + self.assertTrue( + any( + p["operation"] == "MLPAveraging" and "sequence" in p["reason"] + for p in optimizer.pruned + ) + ) + def _window_inputs(self, modality, representations): """(input operation, input rank) of every window node in the search.""" with patch.object( diff --git a/src/main/python/tests/scuro/test_window_operations.py b/src/main/python/tests/scuro/test_window_operations.py index c6a258fb465..ddfecc16862 100644 --- a/src/main/python/tests/scuro/test_window_operations.py +++ b/src/main/python/tests/scuro/test_window_operations.py @@ -30,6 +30,9 @@ 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.aggregated_representation import ( + AggregatedRepresentation, +) from systemds.scuro.representations.timeseries_representations import ( FrequencyMagnitude, Mean, @@ -265,6 +268,30 @@ def test_window_aggregation_without_padding_returns_one_array_per_instance(self) for instance in windowed: self.assertEqual(np.asarray(instance).shape, (10,)) + def test_aggregation_skips_zero_padded_windows(self): + """Ragged instances are padded to the longest window count; Agg over + the windows must see only an instance's own windows.""" + modality = self._timeseries_modality(10) + modality.data = [np.arange(n, dtype=np.float32) + 1 for n in (40, 25, 10)] + modality.metadata = [ + ModalityType.TIMESERIES.create_metadata(["signal"], x) + for x in modality.data + ] + unpadded = WindowAggregation("mean", window_size=10, pad=False).execute( + modality + ) + windowed = modality.context(WindowAggregation("mean", window_size=10)) + self.assertEqual(np.asarray(windowed.data).shape, (3, 4)) + + for aggregation in ["mean", "min", "sum"]: + with self.subTest(aggregation=aggregation): + aggregated = AggregatedRepresentation(aggregation).transform(windowed) + expected = [getattr(np, aggregation)(w) for w in unpadded] + np.testing.assert_allclose( + np.asarray(aggregated.data).ravel(), expected, rtol=1e-6 + ) + self.assertNotIn("valid_windows", aggregated.metadata[0]) + # ------------------------------------------------------------------ # Preconditions # ------------------------------------------------------------------ @@ -345,6 +372,30 @@ def test_dynamic_window_clamps_num_windows_to_the_signal_length(self): (signal_length,), ) + def test_dynamic_window_stats_bound_the_largest_window(self): + """A shorter instance can have a longer last window (it takes the + rounding remainder), so the declared per-window feature shape must + bound the largest window over every length up to L, and by L - q + 1.""" + for num_windows in [5, 16, 100]: + operator = DynamicWindow(FrequencyMagnitude(), num_windows=num_windows) + largest = 0 + for length in range(num_windows, 3000): + largest = max(largest, int(operator._window_sizes(length).max())) + bound = operator._max_window_size(length) + self.assertGreaterEqual(bound, largest) + self.assertLessEqual(bound, length - num_windows + 1) + self.assertLessEqual(bound - largest, num_windows) + + num_windows = 16 + instances = [ + np.random.rand(length).astype(np.float32) for length in (1005, 1007) + ] + operator = DynamicWindow(FrequencyMagnitude(), num_windows=num_windows) + stats = operator.get_output_stats(RepresentationStats(len(instances), (1007,))) + windowed = np.asarray(operator.execute(_FakeModality(instances))) + self.assertEqual(windowed.shape[:2], (len(instances), num_windows)) + self.assertLessEqual(windowed.shape[2], stats.output_shape[1]) + def test_window_operators_on_variable_length_instances(self): """A window *count* is length-independent, so ragged instances still stack into one rectangular block.""" diff --git a/src/main/python/tests/scuro/test_window_representation_batching.py b/src/main/python/tests/scuro/test_window_representation_batching.py index e561830d081..650ab3299cf 100644 --- a/src/main/python/tests/scuro/test_window_representation_batching.py +++ b/src/main/python/tests/scuro/test_window_representation_batching.py @@ -27,6 +27,7 @@ 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.spectrogram import Spectrogram from systemds.scuro.representations.unimodal import UnimodalRepresentation from systemds.scuro.representations.window_aggregation import ( DynamicWindow, @@ -127,6 +128,7 @@ def test_audio_representations_preserve_results_when_batched(self): representations = ( MFCC(n_mfcc=4, n_mels=8, hop_length=8, n_fft=32), MelSpectrogram(n_mels=8, hop_length=8, n_fft=32), + Spectrogram(hop_length=8, n_fft=32), ) for representation in representations: with self.subTest(representation=representation.name): From 8c47a9bbd7316145c995b414e69c634b4c1763ca Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Wed, 28 Oct 2026 14:37:19 +0100 Subject: [PATCH 26/26] missing file --- .../python/systemds/scuro/models/fused_mlp.py | 255 ++++++++++++++++++ 1 file changed, 255 insertions(+) create mode 100644 src/main/python/systemds/scuro/models/fused_mlp.py diff --git a/src/main/python/systemds/scuro/models/fused_mlp.py b/src/main/python/systemds/scuro/models/fused_mlp.py new file mode 100644 index 00000000000..7688f98b8a4 --- /dev/null +++ b/src/main/python/systemds/scuro/models/fused_mlp.py @@ -0,0 +1,255 @@ +# ------------------------------------------------------------- +# +# 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. +# +# ------------------------------------------------------------- +"""Fused training of many surrogate MLPs in one pass (SIGMOD28_PLAN B1). + +The surrogates are tiny, so one fit is bound by per-step Python and kernel +launch overhead, not by FLOPs. Here G = candidates x folds copies of a model's +own `nn.Sequential` train together: every Linear becomes one batched matmul +over a leading group axis, so a step costs about the same for 5 groups as for +250. + +A model opts in with: + build(input_dim) -> nn.Sequential the architecture fit() uses + make_optimizer(params) the optimizer fit() uses (elementwise, + so stacking the parameters changes nothing) + fused_loss "bce", "bce_logits" or "bce_logits_balanced" + scores_from_output(output, y) the metrics test() computes + batch_size, num_steps, drop_last_batch(n, batch_size) + +Each fold starts from the init fit() uses (torch.manual_seed(seed), then +build). Shuffles and dropout masks are drawn once per fold from one generator +seeded with `seed`, and every candidate reads its fold's draws. So a +candidate's numbers do not depend on which other candidates share its batch, +which depends on timing. For the same reason every batch of one input width +runs with the same number of candidate slots (candidate_slots), the unused +ones filled with copies: the GEMM kernels, and with them the last bit of every +result, depend on the number of groups. +""" + +import math +import os + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F + +_ELEMENTWISE = {nn.ReLU: F.relu, nn.GELU: F.gelu, nn.Sigmoid: torch.sigmoid} +_SUPPORTED = (nn.Linear, nn.BatchNorm1d, nn.LayerNorm, nn.Dropout, *_ELEMENTWISE) +_LOSSES = ("bce", "bce_logits", "bce_logits_balanced") + + +def fused_training_enabled() -> bool: + """SCURO_FUSED_TRAINING=0 keeps every task node on its own worker.""" + return ( + os.environ.get("SCURO_FUSED_TRAINING", "1") != "0" and torch.cuda.is_available() + ) + + +def supports_fused_training(model) -> bool: + return ( + getattr(model, "fused_loss", None) in _LOSSES + and callable(getattr(model, "build", None)) + and callable(getattr(model, "make_optimizer", None)) + and callable(getattr(model, "scores_from_output", None)) + and getattr(model, "num_steps", None) + ) + + +class GroupedSequential: + """G stacked copies of one nn.Sequential; activations are G x rows x features.""" + + def __init__(self, net: nn.Sequential, groups: int, device): + unsupported = [type(m).__name__ for m in net if not isinstance(m, _SUPPORTED)] + if unsupported: + raise TypeError(f"fused training does not support {unsupported}") + + def stack(t, shape): + return t.detach().reshape(shape).repeat(groups, 1, 1).to(device) + + self.layers, self.params = [], [] + for m in net: + if isinstance(m, nn.Linear): + w = stack(m.weight.t(), (1, *m.weight.t().shape)).requires_grad_() + b = stack(m.bias, (1, 1, -1)).requires_grad_() + self.layers.append(("linear", w, b)) + self.params += [w, b] + elif isinstance(m, nn.BatchNorm1d): + w = stack(m.weight, (1, 1, -1)).requires_grad_() + b = stack(m.bias, (1, 1, -1)).requires_grad_() + mean = stack(m.running_mean, (1, 1, -1)) + var = stack(m.running_var, (1, 1, -1)) + self.layers.append(("batchnorm", w, b, mean, var, m.momentum, m.eps)) + self.params += [w, b] + elif isinstance(m, nn.LayerNorm): + w = stack(m.weight, (1, 1, -1)).requires_grad_() + b = stack(m.bias, (1, 1, -1)).requires_grad_() + self.layers.append(("layernorm", w, b, m.normalized_shape, m.eps)) + self.params += [w, b] + elif isinstance(m, nn.Dropout): + self.layers.append(("dropout", m.p)) + else: + self.layers.append(("elementwise", _ELEMENTWISE[type(m)])) + + def forward(self, h, train=False, mask=None, noise=None): + """mask: G x rows x 1 marks the real rows of a short batch (train only); + noise(features) returns G x rows x features uniforms for dropout.""" + for kind, *args in self.layers: + if kind == "linear": + h = torch.baddbmm(args[1], h, args[0]) + elif kind == "batchnorm": + w, b, running_mean, running_var, momentum, eps = args + if train: + count = mask.sum(1, keepdim=True) + mean = (h * mask).sum(1, keepdim=True) / count + var = ((h - mean) ** 2 * mask).sum(1, keepdim=True) / count + with torch.no_grad(): + unbiased = var * count / (count - 1).clamp(min=1) + running_mean.lerp_(mean, momentum) + running_var.lerp_(unbiased, momentum) + else: + mean, var = running_mean, running_var + h = (h - mean) * torch.rsqrt(var + eps) * w + b + elif kind == "layernorm": + w, b, shape, eps = args + h = F.layer_norm(h, shape, eps=eps) * w + b + elif kind == "dropout": + if train and args[0] > 0: + h = h * (noise(h.shape[-1]) >= args[0]) / (1 - args[0]) + else: + h = args[0](h) + return h + + +def candidate_bytes(model, dim, n_rows, n_folds): + with torch.random.fork_rng(devices=[]), torch.device("meta"): + net = model.build(dim) + params = sum(p.numel() for p in net.parameters()) + widths = sum(m.out_features for m in net if isinstance(m, nn.Linear)) + per_fold = 4 * (4 * params + 3 * int(model.batch_size) * (dim + widths)) + return 4 * n_rows * dim + n_folds * per_fold + + +def fused_gpu_bytes() -> int: + return int(os.environ.get("SCURO_FUSED_GPU_MB", "4096")) * 2**20 + + +def candidate_slots(model, dim, n_rows, n_folds, per_candidate=None): + if per_candidate is None: + per_candidate = candidate_bytes(model, dim, n_rows, n_folds) + cap = int(os.environ.get("SCURO_FUSED_CANDIDATES", "32")) + return max(1, min(cap, fused_gpu_bytes() // per_candidate)) + + +def fit_folds(model, features, labels, splits, seed, device, eval_elements=2**27): + """Train `model` on every fold of every candidate at once. + + features: C arrays of shape rows x d (one d for all), labels: rows (x O), + splits: K (fit_positions, val_positions) into the rows. Returns the + eval-mode outputs, C x K x rows x O, as a numpy array. + """ + n_candidates, n_folds = len(features), len(splits) + n_rows, dim = features[0].shape + slots = candidate_slots(model, dim, n_rows, n_folds) + if n_candidates > slots: + raise ValueError(f"{n_candidates} candidates for {slots} fused slots") + groups = slots * n_folds + cand = torch.zeros(slots, dtype=torch.long, device=device) + cand[:n_candidates] = torch.arange(n_candidates, device=device) + cand = cand.repeat_interleave(n_folds) + fold = torch.arange(n_folds, device=device).repeat(slots) + + X = torch.from_numpy(np.stack(features).astype(np.float32, copy=False)).to(device) + y = torch.as_tensor(np.asarray(labels, dtype=np.float32), device=device) + y = y.reshape(len(y), -1) + + fit_idx = [torch.as_tensor(np.asarray(f), device=device) for f, _ in splits] + mu = torch.empty(groups, 1, dim, device=device) + sd = torch.empty(groups, 1, dim, device=device) + for g in range(groups): + x = X[cand[g], fit_idx[fold[g]]].double() + scale = x.std(0, unbiased=False) + mu[g, 0], sd[g, 0] = x.mean(0), torch.where(scale == 0, 1.0, scale) + + torch.manual_seed(seed) + net = GroupedSequential(model.build(dim), groups, device) + optimizer = model.make_optimizer(net.params) + + loss_kind = model.fused_loss + pos_weight = None + if loss_kind == "bce_logits_balanced": + pos = torch.stack([y[i].sum() for i in fit_idx]) + neg = torch.tensor([float(y[i].numel()) for i in fit_idx], device=device) - pos + pos_weight = (neg / pos.clamp(min=1.0))[fold].view(groups, 1, 1) + + batch = int(model.batch_size) + lengths = torch.tensor([len(i) for i in fit_idx], device=device) + padded = torch.nn.utils.rnn.pad_sequence(fit_idx, batch_first=True) + per_pass = [ + n // batch if model.drop_last_batch(n, batch) else math.ceil(n / batch) + for n in lengths.tolist() + ] + per_pass_t = torch.tensor(per_pass, device=device) + gen = torch.Generator(device=device) + gen.manual_seed(seed) + order = torch.zeros_like(padded) + offsets = torch.arange(batch, device=device) + invalid = torch.arange(padded.shape[1], device=device) >= lengths[:, None] + + for step in range(int(model.num_steps)): + in_pass = step % per_pass_t + if any(step % p == 0 for p in per_pass): + keys = torch.rand(padded.shape, generator=gen, device=device) + keys = torch.argsort(keys.masked_fill(invalid, 2.0), dim=1) + order = torch.where(in_pass[:, None] == 0, padded.gather(1, keys), order) + positions = in_pass[:, None] * batch + offsets + valid = positions < lengths[:, None] + rows = order.gather(1, positions.clamp(max=padded.shape[1] - 1)) + rows, mask = rows[fold], valid[fold].unsqueeze(-1).float() + + def noise(width): + return torch.rand(n_folds, batch, width, generator=gen, device=device)[fold] + + xb = (X[cand[:, None], rows] - mu) / sd + out = net.forward(xb, train=True, mask=mask, noise=noise) + yb = y[rows] + if loss_kind == "bce": + loss = F.binary_cross_entropy(out.nan_to_num(0.5), yb, reduction="none") + else: + loss = F.binary_cross_entropy_with_logits( + out, yb, pos_weight=pos_weight, reduction="none" + ) + loss = (loss * mask).sum((1, 2)) / (mask.sum((1, 2)) * yb.shape[-1]) + optimizer.zero_grad(set_to_none=True) + loss.sum().backward() + optimizer.step() + + outputs = [] + widest = max( + [dim] + [layer[1].shape[-1] for layer in net.layers if layer[0] == "linear"] + ) + chunk = max(1, eval_elements // (groups * widest)) + with torch.no_grad(): + for start in range(0, n_rows, chunk): + xb = (X[cand, start : start + chunk] - mu) / sd + outputs.append(net.forward(xb).cpu()) + out = torch.cat(outputs, dim=1)[: n_candidates * n_folds].numpy() + return out.reshape(n_candidates, n_folds, n_rows, -1)