From 2651d489998e35bc4d0a5df031a39fc712a9de39 Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Tue, 1 Sep 2026 11:13:49 +0200 Subject: [PATCH 01/15] add fit requirement to representations --- src/main/python/systemds/scuro/representations/bow.py | 7 ++++++- .../python/systemds/scuro/representations/lstm.py | 3 +++ .../representations/mlp_learned_dim_reduction.py | 3 +++ .../representations/multimodal_attention_fusion.py | 3 +++ .../systemds/scuro/representations/representation.py | 11 +++++++++++ .../python/systemds/scuro/representations/tfidf.py | 7 ++++++- 6 files changed, 32 insertions(+), 2 deletions(-) diff --git a/src/main/python/systemds/scuro/representations/bow.py b/src/main/python/systemds/scuro/representations/bow.py index 12fc2e4eb69..7437e8ac18c 100644 --- a/src/main/python/systemds/scuro/representations/bow.py +++ b/src/main/python/systemds/scuro/representations/bow.py @@ -24,7 +24,10 @@ from sklearn.feature_extraction.text import CountVectorizer from systemds.scuro.modality.transformed import TransformedModality -from systemds.scuro.representations.representation import RepresentationStats +from systemds.scuro.representations.representation import ( + FitRequirement, + RepresentationStats, +) from systemds.scuro.representations.unimodal import UnimodalRepresentation from systemds.scuro.representations.utils import save_embeddings @@ -37,6 +40,8 @@ @register_representation(ModalityType.TEXT) class BoW(UnimodalRepresentation): + fit_requirement = FitRequirement.UNSUPERVISED + def __init__(self, ngram_range=2, min_df=2, output_file=None, params=None): parameters = {"ngram_range": [2, 3, 5, 10], "min_df": [1, 2, 4, 8]} super().__init__("BoW", ModalityType.EMBEDDING, parameters) diff --git a/src/main/python/systemds/scuro/representations/lstm.py b/src/main/python/systemds/scuro/representations/lstm.py index 08d5bc8af78..346c8e992db 100644 --- a/src/main/python/systemds/scuro/representations/lstm.py +++ b/src/main/python/systemds/scuro/representations/lstm.py @@ -30,12 +30,15 @@ from systemds.scuro.modality.modality import Modality from systemds.scuro.representations.fusion import Fusion +from systemds.scuro.representations.representation import FitRequirement from systemds.scuro.drsearch.operator_registry import register_fusion_operator @register_fusion_operator() class LSTM(Fusion): + fit_requirement = FitRequirement.SUPERVISED + def __init__( self, width=128, diff --git a/src/main/python/systemds/scuro/representations/mlp_learned_dim_reduction.py b/src/main/python/systemds/scuro/representations/mlp_learned_dim_reduction.py index 0cdc2b8fad8..9c24dfdb4db 100644 --- a/src/main/python/systemds/scuro/representations/mlp_learned_dim_reduction.py +++ b/src/main/python/systemds/scuro/representations/mlp_learned_dim_reduction.py @@ -34,6 +34,7 @@ from systemds.scuro.representations.dimensionality_reduction import ( DimensionalityReduction, ) +from systemds.scuro.representations.representation import FitRequirement from systemds.scuro.modality.type import ModalityType from systemds.scuro.utils.utils import set_random_seeds @@ -50,6 +51,8 @@ class MLPLearnedDimReduction(DimensionalityReduction): :param epochs: The number of epochs to train for """ + fit_requirement = FitRequirement.SUPERVISED + def __init__(self, output_dim=256, batch_size=32, learning_rate=0.001, epochs=5): parameters = { "output_dim": [64, 128, 256, 512, 1024], 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 af5a83c7e50..b972d0a12d5 100644 --- a/src/main/python/systemds/scuro/representations/multimodal_attention_fusion.py +++ b/src/main/python/systemds/scuro/representations/multimodal_attention_fusion.py @@ -29,11 +29,14 @@ from systemds.scuro.drsearch.operator_registry import register_fusion_operator from systemds.scuro.modality.modality import Modality from systemds.scuro.representations.fusion import Fusion +from systemds.scuro.representations.representation import FitRequirement from systemds.scuro.utils.static_variables import get_device @register_fusion_operator() class AttentionFusion(Fusion): + fit_requirement = FitRequirement.SUPERVISED + def __init__( self, hidden_dim=256, diff --git a/src/main/python/systemds/scuro/representations/representation.py b/src/main/python/systemds/scuro/representations/representation.py index ae7df5e4477..52f38dbbc39 100644 --- a/src/main/python/systemds/scuro/representations/representation.py +++ b/src/main/python/systemds/scuro/representations/representation.py @@ -20,6 +20,7 @@ # ------------------------------------------------------------- import abc from dataclasses import dataclass +from enum import Enum from typing import Any, Optional import numpy as np @@ -33,6 +34,14 @@ DEFAULT_DTYPE = np.dtype(np.float32) +class FitRequirement(str, Enum): + """Data required to fit a representation before it can transform inputs.""" + + NONE = "none" + UNSUPERVISED = "unsupervised" + SUPERVISED = "supervised" + + @dataclass class RepresentationStats: num_instances: int @@ -154,6 +163,8 @@ def derive_stats(stats: RepresentationStats, **overrides) -> RepresentationStats class Representation: + fit_requirement = FitRequirement.NONE + def __init__(self, name, parameters): self.name = name self._parameters = parameters diff --git a/src/main/python/systemds/scuro/representations/tfidf.py b/src/main/python/systemds/scuro/representations/tfidf.py index a68fe50e407..ba4caf41862 100644 --- a/src/main/python/systemds/scuro/representations/tfidf.py +++ b/src/main/python/systemds/scuro/representations/tfidf.py @@ -23,7 +23,10 @@ from sklearn.feature_extraction.text import TfidfVectorizer from systemds.scuro.dataloader.text_loader import TextStats from systemds.scuro.modality.transformed import TransformedModality -from systemds.scuro.representations.representation import RepresentationStats +from systemds.scuro.representations.representation import ( + FitRequirement, + RepresentationStats, +) from systemds.scuro.representations.unimodal import UnimodalRepresentation from systemds.scuro.representations.utils import save_embeddings @@ -33,6 +36,8 @@ @register_representation(ModalityType.TEXT) class TfIdf(UnimodalRepresentation): + fit_requirement = FitRequirement.UNSUPERVISED + def __init__(self, min_df=2, output_file=None, params=None): parameters = {"min_df": [min_df, 4, 8]} super().__init__("TF-IDF", ModalityType.EMBEDDING, parameters) From 4e7e23959757e3b553794ee701595009200ad87e Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Tue, 1 Sep 2026 11:28:15 +0200 Subject: [PATCH 02/15] add evaluation split class and indexed modality --- src/main/python/systemds/scuro/__init__.py | 8 ++- .../scuro/drsearch/evaluation_split.py | 48 +++++++++++++ .../python/systemds/scuro/drsearch/task.py | 7 ++ .../python/systemds/scuro/modality/indexed.py | 68 +++++++++++++++++++ .../scuro/representations/representation.py | 23 +++++++ 5 files changed, 153 insertions(+), 1 deletion(-) create mode 100644 src/main/python/systemds/scuro/drsearch/evaluation_split.py create mode 100644 src/main/python/systemds/scuro/modality/indexed.py diff --git a/src/main/python/systemds/scuro/__init__.py b/src/main/python/systemds/scuro/__init__.py index c8ba06e72d9..bea9fab2d1f 100644 --- a/src/main/python/systemds/scuro/__init__.py +++ b/src/main/python/systemds/scuro/__init__.py @@ -25,7 +25,11 @@ from systemds.scuro.dataloader.text_loader import TextLoader from systemds.scuro.dataloader.json_loader import JSONLoader from systemds.scuro.dataloader.tabular_loader import TabularLoader -from systemds.scuro.representations.representation import Representation +from systemds.scuro.representations.representation import ( + FitRequirement, + Representation, + RepresentationFitState, +) from systemds.scuro.representations.aggregate import Aggregation from systemds.scuro.representations.aggregated_representation import ( AggregatedRepresentation, @@ -109,11 +113,13 @@ from systemds.scuro.modality.joined import JoinedModality from systemds.scuro.modality.joined_transformed import JoinedTransformedModality from systemds.scuro.modality.modality import Modality +from systemds.scuro.modality.indexed import IndexedModality from systemds.scuro.utils.identifier import Identifier from systemds.scuro.modality.transformed import TransformedModality from systemds.scuro.modality.type import ModalityType from systemds.scuro.modality.unimodal_modality import UnimodalModality from systemds.scuro.drsearch.task import Task +from systemds.scuro.drsearch.evaluation_split import EvaluationSplit from systemds.scuro.drsearch.operator_registry import Registry from systemds.scuro.drsearch.optimization_data import OptimizationData from systemds.scuro.representations.covarep_audio_features import ( diff --git a/src/main/python/systemds/scuro/drsearch/evaluation_split.py b/src/main/python/systemds/scuro/drsearch/evaluation_split.py new file mode 100644 index 00000000000..2564a15fac8 --- /dev/null +++ b/src/main/python/systemds/scuro/drsearch/evaluation_split.py @@ -0,0 +1,48 @@ +# ------------------------------------------------------------- +# +# 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. +# +# ------------------------------------------------------------- +from dataclasses import dataclass + + +@dataclass(frozen=True) +class EvaluationSplit: + """Global instance IDs used to fit and evaluate one pipeline.""" + + fit_indices: tuple + evaluation_indices: tuple + name: str = "" + + def __post_init__(self): + fit_indices = tuple(int(i) for i in self.fit_indices) + evaluation_indices = tuple(int(i) for i in self.evaluation_indices) + + if len(set(fit_indices)) != len(fit_indices): + raise ValueError("fit_indices contains duplicate instance IDs") + if len(set(evaluation_indices)) != len(evaluation_indices): + raise ValueError("evaluation_indices contains duplicate instance IDs") + overlap = set(fit_indices).intersection(evaluation_indices) + if overlap: + raise ValueError( + "Fit and evaluation instance IDs must be disjoint: " + f"{sorted(overlap)}" + ) + + object.__setattr__(self, "fit_indices", fit_indices) + object.__setattr__(self, "evaluation_indices", evaluation_indices) diff --git a/src/main/python/systemds/scuro/drsearch/task.py b/src/main/python/systemds/scuro/drsearch/task.py index 7ddeaede76c..b98b152ca04 100644 --- a/src/main/python/systemds/scuro/drsearch/task.py +++ b/src/main/python/systemds/scuro/drsearch/task.py @@ -25,6 +25,7 @@ from systemds.scuro.models.model import Model import numpy as np from sklearn.model_selection import train_test_split +from systemds.scuro.drsearch.evaluation_split import EvaluationSplit from systemds.scuro.representations.representation import RepresentationStats _GPU_CONTEXT_FLOOR_BYTES = ( @@ -251,6 +252,12 @@ def _create_cv_splits(self): self.fusion_train_indices = [ idx for idx in self.train_indices if idx not in all_val_indices ] + self.evaluation_splits = tuple( + EvaluationSplit(train_indices, val_indices, name=f"fold_{fold_idx}") + for fold_idx, (train_indices, val_indices) in enumerate( + zip(self.cv_train_indices, self.cv_val_indices) + ) + ) def create_model(self): """ diff --git a/src/main/python/systemds/scuro/modality/indexed.py b/src/main/python/systemds/scuro/modality/indexed.py new file mode 100644 index 00000000000..9841b8b222d --- /dev/null +++ b/src/main/python/systemds/scuro/modality/indexed.py @@ -0,0 +1,68 @@ +# ------------------------------------------------------------- +# +# 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. +# +# ------------------------------------------------------------- +from dataclasses import dataclass + +from systemds.scuro.modality.modality import Modality + + +@dataclass(frozen=True) +class IndexedModality: + """A modality partition paired with its global instance IDs.""" + + modality: Modality + instance_ids: tuple + + def __post_init__(self): + instance_ids = tuple(int(i) for i in self.instance_ids) + if len(set(instance_ids)) != len(instance_ids): + raise ValueError("instance_ids contains duplicate instance IDs") + + if self.modality.has_data(): + num_instances = len(self.modality.data) + elif self.modality.has_metadata(): + num_instances = len(self.modality.metadata) + else: + raise ValueError( + "IndexedModality requires materialized data or instance metadata" + ) + + if len(instance_ids) != num_instances: + raise ValueError( + f"Expected {num_instances} instance IDs, got {len(instance_ids)}" + ) + + object.__setattr__(self, "instance_ids", instance_ids) + + def __len__(self): + return len(self.instance_ids) + + def subset(self, instance_ids) -> "IndexedModality": + requested_ids = tuple(int(i) for i in instance_ids) + positions = {instance_id: i for i, instance_id in enumerate(self.instance_ids)} + unknown_ids = [i for i in requested_ids if i not in positions] + if unknown_ids: + raise KeyError(f"Unknown instance IDs: {unknown_ids}") + + local_indices = [positions[i] for i in requested_ids] + return IndexedModality( + self.modality.subset(local_indices), + requested_ids, + ) diff --git a/src/main/python/systemds/scuro/representations/representation.py b/src/main/python/systemds/scuro/representations/representation.py index 52f38dbbc39..94e66adacb0 100644 --- a/src/main/python/systemds/scuro/representations/representation.py +++ b/src/main/python/systemds/scuro/representations/representation.py @@ -42,6 +42,24 @@ class FitRequirement(str, Enum): SUPERVISED = "supervised" +@dataclass(frozen=True) +class RepresentationFitState: + """State learned by a representation from a specific set of instances.""" + + fit_requirement: FitRequirement + fit_instance_ids: tuple = () + payload: Any = None + + def __post_init__(self): + fit_requirement = FitRequirement(self.fit_requirement) + fit_instance_ids = tuple(int(i) for i in self.fit_instance_ids) + if len(set(fit_instance_ids)) != len(fit_instance_ids): + raise ValueError("fit_instance_ids contains duplicate instance IDs") + + object.__setattr__(self, "fit_requirement", fit_requirement) + object.__setattr__(self, "fit_instance_ids", fit_instance_ids) + + @dataclass class RepresentationStats: num_instances: int @@ -192,6 +210,11 @@ def set_parameters(self, parameters): for parameter in parameters: setattr(self, parameter, parameters[parameter]) + def fit(self, inputs, labels=None) -> RepresentationFitState: + if self.fit_requirement is not FitRequirement.NONE: + raise NotImplementedError(f"fit is not implemented for {self.name}") + return RepresentationFitState(FitRequirement.NONE) + def check_preconditions(self, input_stats) -> Optional[str]: return None From ac59718712776bd74914cbb7de355989de7600d2 Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Tue, 1 Sep 2026 11:35:33 +0200 Subject: [PATCH 03/15] add split awareness to tfidf and bow --- .../systemds/scuro/representations/bow.py | 49 ++++++++++++++++--- .../scuro/representations/representation.py | 5 ++ .../systemds/scuro/representations/tfidf.py | 45 +++++++++++++++-- 3 files changed, 88 insertions(+), 11 deletions(-) diff --git a/src/main/python/systemds/scuro/representations/bow.py b/src/main/python/systemds/scuro/representations/bow.py index 7437e8ac18c..93e72ddf8fb 100644 --- a/src/main/python/systemds/scuro/representations/bow.py +++ b/src/main/python/systemds/scuro/representations/bow.py @@ -23,9 +23,11 @@ import numpy as np from sklearn.feature_extraction.text import CountVectorizer +from systemds.scuro.modality.indexed import IndexedModality from systemds.scuro.modality.transformed import TransformedModality from systemds.scuro.representations.representation import ( FitRequirement, + RepresentationFitState, RepresentationStats, ) from systemds.scuro.representations.unimodal import UnimodalRepresentation @@ -83,21 +85,29 @@ def estimate_peak_memory_bytes(self, input_stats: TextStats) -> dict: "gpu_peak_bytes": 0, } - def transform(self, modality, aggregation=None): - transformed_modality = TransformedModality(modality, self) - vectorizer = CountVectorizer( + def _create_vectorizer(self): + return CountVectorizer( ngram_range=(1, self.ngram_range), min_df=self.min_df, max_features=_MAX_VOCAB_FEATURES, dtype=np.float32, ) - X = ( - vectorizer.fit_transform(modality.data) - .toarray() - .astype(np.float32, copy=False) + def fit(self, inputs, labels=None) -> RepresentationFitState: + if not isinstance(inputs, IndexedModality): + raise TypeError("BoW.fit expects an IndexedModality") + + vectorizer = self._create_vectorizer() + vectorizer.fit(inputs.modality.data) + return RepresentationFitState( + self.fit_requirement, + fit_instance_ids=inputs.instance_ids, + payload=vectorizer, ) + def _to_modality(self, modality, matrix, aggregation=None): + transformed_modality = TransformedModality(modality, self) + X = matrix.toarray().astype(np.float32, copy=False) if self.output_file is not None: save_embeddings(X, self.output_file) @@ -106,3 +116,28 @@ def transform(self, modality, aggregation=None): if aggregation is not None: transformed_modality.data = aggregation.execute(transformed_modality.data) return transformed_modality + + def transform_fitted(self, inputs, fit_state, aggregation=None): + if not isinstance(inputs, IndexedModality): + raise TypeError("BoW.transform_fitted expects an IndexedModality") + if ( + not isinstance(fit_state, RepresentationFitState) + or fit_state.fit_requirement is not self.fit_requirement + or type(fit_state.payload) is not CountVectorizer + ): + raise ValueError("BoW requires a compatible fitted state") + + transformed = self._to_modality( + inputs.modality, + fit_state.payload.transform(inputs.modality.data), + aggregation, + ) + return IndexedModality(transformed, inputs.instance_ids) + + def transform(self, modality, aggregation=None): + vectorizer = self._create_vectorizer() + return self._to_modality( + modality, + vectorizer.fit_transform(modality.data), + aggregation, + ) diff --git a/src/main/python/systemds/scuro/representations/representation.py b/src/main/python/systemds/scuro/representations/representation.py index 94e66adacb0..c5f198f0f82 100644 --- a/src/main/python/systemds/scuro/representations/representation.py +++ b/src/main/python/systemds/scuro/representations/representation.py @@ -215,6 +215,11 @@ def fit(self, inputs, labels=None) -> RepresentationFitState: raise NotImplementedError(f"fit is not implemented for {self.name}") return RepresentationFitState(FitRequirement.NONE) + def transform_fitted(self, inputs, fit_state, **kwargs): + raise NotImplementedError( + f"transform_fitted is not implemented for {self.name}" + ) + def check_preconditions(self, input_stats) -> Optional[str]: return None diff --git a/src/main/python/systemds/scuro/representations/tfidf.py b/src/main/python/systemds/scuro/representations/tfidf.py index ba4caf41862..8ff2757e7c3 100644 --- a/src/main/python/systemds/scuro/representations/tfidf.py +++ b/src/main/python/systemds/scuro/representations/tfidf.py @@ -22,9 +22,11 @@ from sklearn.feature_extraction.text import TfidfVectorizer from systemds.scuro.dataloader.text_loader import TextStats +from systemds.scuro.modality.indexed import IndexedModality from systemds.scuro.modality.transformed import TransformedModality from systemds.scuro.representations.representation import ( FitRequirement, + RepresentationFitState, RepresentationStats, ) from systemds.scuro.representations.unimodal import UnimodalRepresentation @@ -71,15 +73,50 @@ def estimate_peak_memory_bytes(self, input_stats: TextStats) -> dict: cpu_peak = int(dense_bytes * 2.2 + 32 * 1024 * 1024) return {"cpu_peak_bytes": cpu_peak, "gpu_peak_bytes": 0} - def transform(self, modality, aggregation=None): - transformed_modality = TransformedModality(modality, self) + def _create_vectorizer(self): + return TfidfVectorizer(min_df=self.min_df) + + def fit(self, inputs, labels=None) -> RepresentationFitState: + if not isinstance(inputs, IndexedModality): + raise TypeError("TfIdf.fit expects an IndexedModality") - vectorizer = TfidfVectorizer(min_df=self.min_df) + vectorizer = self._create_vectorizer() + vectorizer.fit(inputs.modality.data) + return RepresentationFitState( + self.fit_requirement, + fit_instance_ids=inputs.instance_ids, + payload=vectorizer, + ) - X = vectorizer.fit_transform(modality.data).astype(np.float32, copy=False) + def _to_modality(self, modality, matrix): + transformed_modality = TransformedModality(modality, self) + X = matrix.astype(np.float32, copy=False) if self.output_file is not None: save_embeddings(X, self.output_file) transformed_modality.data_type = self.data_type transformed_modality.data = X.toarray() return transformed_modality + + def transform_fitted(self, inputs, fit_state, aggregation=None): + if not isinstance(inputs, IndexedModality): + raise TypeError("TfIdf.transform_fitted expects an IndexedModality") + if ( + not isinstance(fit_state, RepresentationFitState) + or fit_state.fit_requirement is not self.fit_requirement + or type(fit_state.payload) is not TfidfVectorizer + ): + raise ValueError("TfIdf requires a compatible fitted state") + + transformed = self._to_modality( + inputs.modality, + fit_state.payload.transform(inputs.modality.data), + ) + return IndexedModality(transformed, inputs.instance_ids) + + def transform(self, modality, aggregation=None): + vectorizer = self._create_vectorizer() + return self._to_modality( + modality, + vectorizer.fit_transform(modality.data), + ) From d56777d1db40318db76013380fced241f7563b0c Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Tue, 1 Sep 2026 12:06:16 +0200 Subject: [PATCH 04/15] add split aware evaluator of dags --- .../scuro/drsearch/split_aware_evaluator.py | 265 ++++++++++++++++++ .../scuro/drsearch/unimodal_optimizer.py | 85 +++++- 2 files changed, 336 insertions(+), 14 deletions(-) create mode 100644 src/main/python/systemds/scuro/drsearch/split_aware_evaluator.py diff --git a/src/main/python/systemds/scuro/drsearch/split_aware_evaluator.py b/src/main/python/systemds/scuro/drsearch/split_aware_evaluator.py new file mode 100644 index 00000000000..31a1566dea1 --- /dev/null +++ b/src/main/python/systemds/scuro/drsearch/split_aware_evaluator.py @@ -0,0 +1,265 @@ +# ------------------------------------------------------------- +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# +# ------------------------------------------------------------- +import copy +import time +from dataclasses import dataclass + +import numpy as np + +from systemds.scuro.drsearch.representation_dag import ( + RepresentationDag, + pushdown_aggregation_for_node, +) +from systemds.scuro.drsearch.task import PerformanceMeasure +from systemds.scuro.modality.indexed import IndexedModality +from systemds.scuro.representations.aggregated_representation import ( + AggregatedRepresentation, +) +from systemds.scuro.representations.context import Context +from systemds.scuro.representations.dimensionality_reduction import ( + DimensionalityReduction, +) +from systemds.scuro.representations.representation import ( + FitRequirement, + Representation, +) +from systemds.scuro.representations.unimodal import UnimodalRepresentation + + +@dataclass +class SplitEvaluationResult: + scores: list + representation_time: float + task_time: float + + +def dag_requires_fitting(dag: RepresentationDag) -> bool: + return any( + node.operation is not None + and getattr(node.operation, "fit_requirement", FitRequirement.NONE) + is not FitRequirement.NONE + for node in dag.nodes + ) + + +class SplitAwareDAGEvaluator: + def __init__(self, dag: RepresentationDag, modality, task): + self.dag = dag + self.modality = modality + self.task = task + + @staticmethod + def materialize(modality) -> None: + if modality.has_data(): + return + + loader = getattr(modality, "data_loader", None) + if loader is None: + raise ValueError("Split-aware evaluation requires materialized data") + + if getattr(loader, "chunk_size", None) is None: + modality.extract_raw_data() + return + + full_loader = copy.copy(loader) + full_loader.chunk_size = None + full_loader.reset() + data, metadata = full_loader.load() + modality._data = data + modality.metadata = metadata + + @staticmethod + def _wrap(transformed, instance_ids) -> IndexedModality: + return IndexedModality(transformed, instance_ids) + + def _apply_stateless(self, operation, inputs, node): + aggregation = pushdown_aggregation_for_node(node.parameters) + + if isinstance(operation, Context): + transformed = inputs.modality.context(operation) + elif isinstance(operation, DimensionalityReduction): + if getattr(operation, "needs_training", False): + raise NotImplementedError( + "Trainable dimensionality reduction is not yet implemented for split-aware evaluation" + ) + transformed = inputs.modality.dimensionality_reduction(operation) + elif isinstance(operation, AggregatedRepresentation): + transformed = operation.transform(inputs.modality) + elif isinstance(operation, UnimodalRepresentation): + transformed = inputs.modality.apply_representation( + operation, aggregation=aggregation + ) + elif isinstance(operation, Representation): + transformed = inputs.modality.apply_representation( + operation, aggregation=aggregation + ) + else: + raise TypeError(f"Unsupported unimodal operation: {type(operation).__name__}") + + return self._wrap(transformed, inputs.instance_ids) + + def _execute_partitions(self, *partitions): + cache = {} + + def execute_node(node_id): + if node_id in cache: + return cache[node_id] + + node = self.dag.get_node_by_id(node_id) + if not node.inputs: + result = partitions + cache[node_id] = result + return result + if len(node.inputs) != 1: + raise NotImplementedError( + "Multimodal split-aware execution is handled in step 4" + ) + + input_partitions = execute_node(node.inputs[0]) + operation = node.operation(params=node.parameters) + requirement = getattr( + operation, "fit_requirement", FitRequirement.NONE + ) + + if requirement is not FitRequirement.NONE: + labels = None + if requirement is FitRequirement.SUPERVISED: + labels = np.asarray(self.task.labels)[ + list(input_partitions[0].instance_ids) + ] + fit_state = operation.fit(input_partitions[0], labels=labels) + aggregation = pushdown_aggregation_for_node(node.parameters) + result = tuple( + operation.transform_fitted( + partition, + fit_state, + aggregation=aggregation, + ) + for partition in input_partitions + ) + else: + result = tuple( + self._apply_stateless(operation, partition, node) + for partition in input_partitions + ) + + cache[node_id] = result + return result + + return execute_node(self.dag.root_node_id) + + def fit_transform(self, fit_indices, transform_indices) -> IndexedModality: + self.materialize(self.modality) + full = IndexedModality(self.modality, range(len(self.modality.data))) + fit_partition = full.subset(fit_indices) + transform_partition = full.subset(transform_indices) + _, transformed = self._execute_partitions( + fit_partition, + transform_partition, + ) + return transformed + + @staticmethod + def _partition_data(partition): + return list(partition.modality.data) + + def evaluate(self) -> SplitEvaluationResult: + self.materialize(self.modality) + full = IndexedModality(self.modality, range(len(self.modality.data))) + + train_scores = PerformanceMeasure( + "train", self.task.performance_measures + ) + val_scores = PerformanceMeasure("val", self.task.performance_measures) + test_scores = PerformanceMeasure("test", self.task.performance_measures) + self.task.training_time = [] + self.task.inference_time = [] + + representation_time = 0.0 + task_time = 0.0 + for split in self.task.evaluation_splits: + train = full.subset(split.fit_indices) + evaluation = full.subset(split.evaluation_indices) + test = full.subset(self.task.test_indices) + + start = time.perf_counter() + transformed_train, transformed_evaluation, transformed_test = ( + self._execute_partitions(train, evaluation, test) + ) + representation_time += time.perf_counter() - start + + model = self.task.create_model() + train_X = self._partition_data(transformed_train) + evaluation_X = self._partition_data(transformed_evaluation) + test_X = self._partition_data(transformed_test) + train_y = np.asarray(self.task.labels)[list(split.fit_indices)] + evaluation_y = np.asarray(self.task.labels)[ + list(split.evaluation_indices) + ] + test_y = np.asarray(self.task.labels)[list(self.task.test_indices)] + + start = time.perf_counter() + train_start = time.perf_counter() + train_score = model.fit( + train_X, + train_y, + evaluation_X, + evaluation_y, + ) + self.task.training_time.append(time.perf_counter() - train_start) + val_score = model.test(evaluation_X, evaluation_y) + test_start = time.perf_counter() + test_score = model.test(np.asarray(test_X), test_y) + self.task.inference_time.append(time.perf_counter() - test_start) + task_time += time.perf_counter() - start + + train_scores.add_scores(train_score[0]) + val_scores.add_scores(val_score[0]) + test_scores.add_scores(test_score[0]) + + if hasattr(model, "clean_up"): + model.clean_up() + + self.task.last_run_timing = { + "train_time_per_fold_s": list(self.task.training_time), + "test_inference_time_per_fold_s": list(self.task.inference_time), + "train_time_mean_s": ( + float(np.mean(self.task.training_time)) + if self.task.training_time + else 0.0 + ), + "test_inference_time_mean_s": ( + float(np.mean(self.task.inference_time)) + if self.task.inference_time + else 0.0 + ), + "n_test_instances": len(self.task.test_indices), + } + + return SplitEvaluationResult( + scores=[ + train_scores.compute_averages(), + val_scores.compute_averages(), + test_scores.compute_averages(), + ], + representation_time=representation_time, + task_time=task_time, + ) diff --git a/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py b/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py index 5f130c369a3..1eaf9303fdd 100644 --- a/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py +++ b/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py @@ -29,6 +29,10 @@ from systemds.scuro.modality.type import ModalityType from systemds.scuro.drsearch.node_executor import NodeExecutor, ResultEntry +from systemds.scuro.drsearch.split_aware_evaluator import ( + SplitAwareDAGEvaluator, + dag_requires_fitting, +) from systemds.scuro.representations.representation import RepresentationStats from systemds.scuro.drsearch.ranking import rank_by_tradeoff, rank_by_robustness from systemds.scuro.drsearch.task import PerformanceMeasure @@ -109,7 +113,8 @@ def __init__( if tasks[i - 1].expected_dim != tasks[i].expected_dim: self._tasks_require_same_dims = False - self._combination_operators = [Concatenation(), Hadamard(), Sum()] + # self._combination_operators = [Concatenation(), Hadamard(), Sum(), Average()] + self._combination_operators = [Concatenation()] @lru_cache(maxsize=128) def _get_modality_operators(self, modality_type): @@ -350,23 +355,66 @@ def _process_modality(self, modality, skip_remaining: int = 0, scheduler=None): self._build_execution_dags_for_modality(modality, skip_remaining) ) - node_executor = NodeExecutor( - expanded_dags_with_task_roots, - [modality], - self.tasks, - max_num_workers=self.max_num_workers, - result_path=self.result_path, - enable_checkpointing=self.enable_checkpointing, - search_start=self._search_start, - ) start_time = time.perf_counter() - exec_out = node_executor.run() + fitted_dags = [ + dag for dag in dags_with_pushdown if dag_requires_fitting(dag) + ] + stateless_dag_ids = { + dag.dag_id + for dag in dags_with_pushdown + if not dag_requires_fitting(dag) + } + stateless_task_dags = [ + dag + for dag in expanded_dags_with_task_roots + if dag.dag_id in stateless_dag_ids + ] + + task_results = [] + statistics = {"worker_stats": {}, "node_stats": {}, "reuse": {}} + if fitted_dags: + SplitAwareDAGEvaluator.materialize(modality) + + if stateless_task_dags: + node_executor = NodeExecutor( + stateless_task_dags, + [modality], + self.tasks, + max_num_workers=self.max_num_workers, + result_path=self.result_path, + enable_checkpointing=self.enable_checkpointing, + search_start=self._search_start, + ) + exec_out = node_executor.run() + task_results.extend(exec_out["task_results"]) + statistics = exec_out["statistics"] + + if fitted_dags: + fitted_task_dags = self._expand_dags_with_task_roots(fitted_dags) + for dag in fitted_dags: + for task_idx, task in enumerate(self.tasks): + evaluation = SplitAwareDAGEvaluator(dag, modality, task).evaluate() + result_dag = next( + task_dag + for task_dag in fitted_task_dags + if task_dag.dag_id == dag.dag_id + and task_dag.nodes[-1].parameters.get("_task_idx") == task_idx + ) + task_results.append( + ResultEntry( + train_score=evaluation.scores[0].average_scores, + val_score=evaluation.scores[1].average_scores, + test_score=evaluation.scores[2].average_scores, + representation_time=evaluation.representation_time, + task_time=evaluation.task_time, + dag=result_dag, + ) + ) + end_time = time.perf_counter() - task_results = exec_out["task_results"] for task_result in task_results: local_results.add_task_result(task_result, dags) - statistics = exec_out["statistics"] local_results.add_worker_stat(statistics["worker_stats"], modality.modality_id) local_results.add_node_stat(statistics["node_stats"], modality.modality_id) @@ -967,7 +1015,16 @@ def get_k_best_results( if not task_cache: cache = [] for result in results: - cache.append(result.dag.execute([modality])) + if dag_requires_fitting(result.dag): + transformed = SplitAwareDAGEvaluator( + result.dag, modality, task + ).fit_transform( + task.train_indices, + range(len(modality.data)), + ) + cache.append(transformed.modality) + else: + cache.append(result.dag.execute([modality])) elif isinstance(task_cache, list): cache = task_cache From 2991b33682411935ea6f04a10a83f883d8a1a629 Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Wed, 2 Sep 2026 13:41:23 +0200 Subject: [PATCH 05/15] add split awareness to fusion --- .../scuro/dataloader/timeseries_loader.py | 9 +- .../scuro/drsearch/hyperparameter_tuner.py | 88 +++++++++----- .../scuro/drsearch/multimodal_ga_optimizer.py | 93 +++++++++++++-- .../scuro/drsearch/multimodal_optimizer.py | 41 +++++++ .../systemds/scuro/drsearch/node_executor.py | 43 ++++++- .../scuro/drsearch/representation_dag.py | 77 +++++++++--- .../scuro/drsearch/split_aware_evaluator.py | 112 ++++++++++++------ .../python/systemds/scuro/drsearch/task.py | 14 --- .../scuro/drsearch/test_set_evaluation.py | 30 ++++- .../scuro/drsearch/unimodal_optimizer.py | 10 ++ .../systemds/scuro/modality/transformed.py | 16 --- .../dimensionality_reduction.py | 75 ++++++++---- .../systemds/scuro/representations/fusion.py | 86 ++++++++++---- .../mlp_learned_dim_reduction.py | 12 +- 14 files changed, 516 insertions(+), 190 deletions(-) diff --git a/src/main/python/systemds/scuro/dataloader/timeseries_loader.py b/src/main/python/systemds/scuro/dataloader/timeseries_loader.py index a2f515f7d26..685105e31e5 100644 --- a/src/main/python/systemds/scuro/dataloader/timeseries_loader.py +++ b/src/main/python/systemds/scuro/dataloader/timeseries_loader.py @@ -60,6 +60,7 @@ def __init__( self.normalize = normalize self.channel_index = channel_index self.file_format = file_format.lower() + self.channel_index = channel_index self.stats = self.get_stats(source_path, sampling_rate) if self.file_format not in ["npy", "mat", "hdf5", "txt"]: raise ValueError(f"Unsupported file format: {self.file_format}") @@ -69,7 +70,9 @@ def extract(self, file: str, index: Optional[Union[str, List[str]]] = None): data = self._load_data(file) - if data.ndim > 1 and len(self.signal_names) == 1: + if self.channel_index is not None and data.ndim > 1: + data = data[:, self.channel_index] + elif data.ndim > 1 and len(self.signal_names) == 1: data = data.flatten() if self.normalize: @@ -160,10 +163,12 @@ def get_stats(self, source_path: str, sampling_rate: int): for file_name in self.indices: file = source_path + file_name + "." + self.file_format data = self._load_data(file) + if self.channel_index is not None and data.ndim > 1: + data = data[:, self.channel_index] max_length = max(max_length, data.shape[0]) avg_length += data.shape[0] num_instances += 1 - num_signals = max(num_signals, data.shape[1]) + num_signals = max(num_signals, data.shape[1] if data.ndim > 1 else 1) return TimeseriesStats( max_length, num_instances, diff --git a/src/main/python/systemds/scuro/drsearch/hyperparameter_tuner.py b/src/main/python/systemds/scuro/drsearch/hyperparameter_tuner.py index 094a4b9585e..529e38032de 100644 --- a/src/main/python/systemds/scuro/drsearch/hyperparameter_tuner.py +++ b/src/main/python/systemds/scuro/drsearch/hyperparameter_tuner.py @@ -31,10 +31,13 @@ import math import random from systemds.scuro.drsearch.representation_dag import ( - RepresentationDAGBuilder, RepresentationNode, ) from systemds.scuro.modality.modality import Modality +from systemds.scuro.drsearch.split_aware_evaluator import ( + SplitAwareDAGEvaluator, + dag_requires_fitting, +) from systemds.scuro.drsearch.task import PerformanceMeasure import pickle from systemds.scuro.utils.checkpointing import CheckpointManager @@ -229,12 +232,40 @@ def _apply_trial_params_to_node( return {**base_params, **trial_params} +class HyperparamEvaluationScores(list): + """Backward-compatible score triple with complete CV details.""" + + def __init__(self, scores, metric): + split_names = ("train", "val", "test") + for split_name, score in zip(split_names, scores): + if isinstance(score, PerformanceMeasure): + averages = dict(score.average_scores or {}) + fold_scores = score.fold_scores() + score_stds = score.score_stds() + else: + value = float(score) + averages = {metric: value} + fold_scores = {metric: [value]} + score_stds = {metric: 0.0} + setattr(self, f"all_{split_name}_scores", averages) + setattr(self, f"{split_name}_fold_scores", fold_scores) + setattr(self, f"{split_name}_score_stds", score_stds) + + super().__init__( + [ + self.all_train_scores.get(metric, np.nan), + self.val_fold_scores.get(metric, []), + self.all_test_scores.get(metric, np.nan), + ] + ) + + @dataclass class HyperparamResult: representation_name: str best_params: Dict[str, Any] best_score: float - all_results: List[Tuple[Dict[str, Any], float]] + all_results: List[Tuple[Dict[str, Any], Any]] tuning_time: float modality_id: int task_name: str @@ -272,25 +303,23 @@ def get_k_best_dags(self, modality, task): results = self.results[task.model.name][modality.modality_id] dags = [] for result in results: - dag_with_best_params = RepresentationDAGBuilder() - prev_node_id = None - for node in result.dag.nodes: - if node.operation is not None and node.parameters: - params = _apply_trial_params_to_node(node, result.best_params) - prev_node_id = dag_with_best_params.create_operation_node( - node.operation, [prev_node_id], params - ) - else: # it is a leaf node - prev_node_id = dag_with_best_params.create_leaf_node( - node.modality_id + dag_with_best_params = copy.deepcopy(result.dag) + for node in dag_with_best_params.nodes: + if node.operation is not None: + node.parameters = _apply_trial_params_to_node( + node, result.best_params ) - - dags.append(dag_with_best_params.build(prev_node_id)) + dags.append(dag_with_best_params) return results, dags def get_k_best_results(self, modality, task, performance_metric_name): results, dags = self.get_k_best_dags(modality, task) - representations = [list(dag.execute([modality]).values())[-1] for dag in dags] + representations = [] + for dag in dags: + representation = dag.execute([modality]) + if isinstance(representation, dict): + representation = representation[list(representation.keys())[-1]] + representations.append(representation) return results, representations @@ -499,11 +528,9 @@ def visit_node(node_id): ) baseline = ( baseline_params, - [ - self._score_value(baseline_raw_scores[0]), - self._score_value_list(baseline_raw_scores[1]), - self._score_value(baseline_raw_scores[2]), - ], + HyperparamEvaluationScores( + baseline_raw_scores, self.scoring_metric + ), ) if not hyperparams: @@ -927,21 +954,19 @@ def objective(trial: optuna.Trial) -> float: task, modalities_override=modalities_override, ) - train_score = self._score_value(scores[0]) val_score = self._score_value(scores[1]) - test_score = self._score_value(scores[2]) if np.isnan(val_score): raise optuna.TrialPruned() - seen[self._config_key(params)] = ( - params, - [train_score, self._score_value_list(scores[1]), test_score], + evaluation_scores = HyperparamEvaluationScores( + scores, self.scoring_metric ) + seen[self._config_key(params)] = (params, evaluation_scores) - trial_results.append( - (params, [train_score, self._score_value_list(scores[1]), test_score]) + trial_results.append((params, evaluation_scores)) + val_folds = evaluation_scores.val_fold_scores.get( + self.scoring_metric, [] ) - val_folds = self._score_value_list(scores[1]) if val_folds is not None and len(val_folds) > 1: lam = 0.5 robust_val = np.mean(val_folds) - lam * np.std(val_folds, ddof=1) @@ -1004,6 +1029,11 @@ def evaluate_dag_config( if modalities_override is not None else self.get_modalities_by_id(modality_ids) ) + if dag_requires_fitting(dag_copy): + evaluation = SplitAwareDAGEvaluator( + dag_copy, modalities, task + ).evaluate() + return params, evaluation.scores modified_modality = dag_copy.execute(modalities, task) score = task.run(modified_modality.data) 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 44858d96a0e..62bf203f4f8 100644 --- a/src/main/python/systemds/scuro/drsearch/multimodal_ga_optimizer.py +++ b/src/main/python/systemds/scuro/drsearch/multimodal_ga_optimizer.py @@ -44,8 +44,13 @@ from systemds.scuro.drsearch.representation_dag import ( RepresentationDAGBuilder, RepresentationDag, + compose_dag_leaves, ) from systemds.scuro.drsearch.task import Task +from systemds.scuro.drsearch.split_aware_evaluator import ( + SplitAwareDAGEvaluator, + dag_requires_fitting, +) from systemds.scuro.drsearch.worker_pool import PersistentWorkerPool, create_mp_context from systemds.scuro.representations.aggregated_representation import ( AggregatedRepresentation, @@ -140,6 +145,9 @@ class FusionSearchResult: val_fold_scores: dict = field(default_factory=dict) train_fold_scores: dict = field(default_factory=dict) test_fold_scores: dict = field(default_factory=dict) + train_score_stds: dict = field(default_factory=dict) + val_score_stds: dict = field(default_factory=dict) + test_score_stds: dict = field(default_factory=dict) task_timing: dict = field(default_factory=dict) @@ -176,12 +184,6 @@ def _failure_fitness(objective_specs: List[ObjectiveSpec]) -> Tuple[float, ...]: ) -def _fold_scores(performance_measure) -> Dict[str, List[float]]: - return { - metric: [float(value) for value in values] - for metric, values in performance_measure.scores.items() - } - def _evaluate_genome_body( dag: RepresentationDag, @@ -190,6 +192,34 @@ def _evaluate_genome_body( objective_specs: List[ObjectiveSpec], ) -> Tuple[Optional[Tuple[float, ...]], Optional[Dict[str, Any]]]: start = time.time() + if dag_requires_fitting(dag): + evaluation = SplitAwareDAGEvaluator(dag, modalities, task).evaluate() + scores = evaluation.scores + total = evaluation.representation_time + evaluation.task_time + val_score = scores[1].average_scores + timing = { + "runtime": total, + "task_time": evaluation.task_time, + "representation_time": evaluation.representation_time, + } + fitness = tuple( + _objective_value(name, val_score, timing) for name, _ in objective_specs + ) + payload = { + "train_score": scores[0].average_scores, + "val_score": val_score, + "test_score": scores[2].average_scores, + "train_fold_scores": scores[0].fold_scores(), + "val_fold_scores": scores[1].fold_scores(), + "test_fold_scores": scores[2].fold_scores(), + "train_score_stds": scores[0].score_stds(), + "val_score_stds": scores[1].score_stds(), + "test_score_stds": scores[2].score_stds(), + "task_timing": getattr(task, "last_run_timing", {}), + **timing, + } + return fitness, payload + fused = dag.execute(modalities, task, enable_cache=False) if fused is None: return None, None @@ -218,9 +248,12 @@ def _evaluate_genome_body( "train_score": scores[0].average_scores, "val_score": val_score, "test_score": scores[2].average_scores, - "train_fold_scores": _fold_scores(scores[0]), - "val_fold_scores": _fold_scores(scores[1]), - "test_fold_scores": _fold_scores(scores[2]), + "train_fold_scores": scores[0].fold_scores(), + "val_fold_scores": scores[1].fold_scores(), + "test_fold_scores": scores[2].fold_scores(), + "train_score_stds": scores[0].score_stds(), + "val_score_stds": scores[1].score_stds(), + "test_score_stds": scores[2].score_stds(), "task_timing": getattr(task, "last_run_timing", {}), **timing, } @@ -324,6 +357,7 @@ def __init__( "MultimodalDeapOptimizer requires at least one registered " "fusion operator." ) + self.k_best_dags = {} self.k_best_representations = self._extract_k_best( unimodal_optimization_results ) @@ -685,9 +719,8 @@ def _record_evaluation( self.evaluation_errors.get(task_name, 0) + 1 ) if self.debug and error is not None: - last_line = error.strip().splitlines()[-1] if error.strip() else error print( - f"[GA] genome evaluation failed for task={task_name}: {last_line}" + f"[GA] genome evaluation failed for task={task_name}:\n{error}" ) return @@ -699,6 +732,9 @@ def _record_evaluation( train_fold_scores=payload.get("train_fold_scores", {}), val_fold_scores=payload.get("val_fold_scores", {}), test_fold_scores=payload.get("test_fold_scores", {}), + train_score_stds=payload.get("train_score_stds", {}), + val_score_stds=payload.get("val_score_stds", {}), + test_score_stds=payload.get("test_score_stds", {}), task_timing=payload.get("task_timing", {}), runtime=payload["runtime"], task_time=payload["task_time"], @@ -764,13 +800,46 @@ def _extract_k_best(self, unimodal_results) -> Dict[str, Dict[str, List[Any]]]: for task in self.tasks: name = task.model.name k_best[name] = {} + self.k_best_dags[name] = {} for modality in self.modalities: - _, cached_data = unimodal_results.get_k_best_results( + results, cached_data = unimodal_results.get_k_best_results( modality, task, self.metric_name ) + if hasattr(unimodal_results, "get_k_best_dags"): + _, source_dags = unimodal_results.get_k_best_dags(modality, task) + else: + source_dags = [getattr(result, "dag", None) for result in results] + if len(source_dags) != len(cached_data): + raise ValueError( + "Unimodal result DAGs and cached representations are misaligned" + ) k_best[name][modality.modality_id] = cached_data + self.k_best_dags[name][modality.modality_id] = source_dags return k_best + def compose_end_to_end_dag(self, dag, task): + """Expand fusion leaves into the selected unimodal representation DAGs.""" + task_name = task.model.name if hasattr(task, "model") else str(task) + leaf_dags = {} + for leaf_id in dag.get_leaf_nodes(): + leaf = dag.get_node_by_id(leaf_id) + representation_index = int(leaf.representation_index) + try: + source_dag = self.k_best_dags[task_name][leaf.modality_id][ + representation_index + ] + except (KeyError, IndexError, TypeError) as exc: + raise KeyError( + "No selected unimodal DAG for multimodal leaf " + f"({leaf.modality_id}, {leaf.representation_index})" + ) from exc + if source_dag is None: + raise ValueError( + "The unimodal result source does not expose representation DAGs" + ) + leaf_dags[(leaf.modality_id, leaf.representation_index)] = source_dag + return compose_dag_leaves(dag, leaf_dags) + def _available_modality_ids(self, task_name: str) -> List[Any]: reps = self.k_best_representations[task_name] return [ diff --git a/src/main/python/systemds/scuro/drsearch/multimodal_optimizer.py b/src/main/python/systemds/scuro/drsearch/multimodal_optimizer.py index 42369e423be..c146a6485a2 100644 --- a/src/main/python/systemds/scuro/drsearch/multimodal_optimizer.py +++ b/src/main/python/systemds/scuro/drsearch/multimodal_optimizer.py @@ -30,6 +30,10 @@ RepresentationDag, RepresentationDAGBuilder, ) +from systemds.scuro.drsearch.split_aware_evaluator import ( + SplitAwareDAGEvaluator, + dag_requires_fitting, +) from systemds.scuro.representations.aggregated_representation import ( AggregatedRepresentation, ) @@ -57,6 +61,22 @@ 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): + evaluation = SplitAwareDAGEvaluator( + dag, modalities_for_dag, task + ).evaluate() + total_time = evaluation.representation_time + evaluation.task_time + return OptimizationResult( + dag=dag, + train_score=evaluation.scores[0].average_scores, + val_score=evaluation.scores[1].average_scores, + test_score=evaluation.scores[2].average_scores, + runtime=total_time, + task_name=task.model.name, + task_time=evaluation.task_time, + representation_time=evaluation.representation_time, + ) + fused_representation = dag.execute(modalities_for_dag, task) if fused_representation is None: return None @@ -406,6 +426,27 @@ def _evaluate_dag(self, dag: RepresentationDag, task: Task) -> "OptimizationResu start_time = time.time() try: + modalities_for_dag = list( + chain.from_iterable( + self.k_best_representations[task.model.name].values() + ) + ) + if dag_requires_fitting(dag): + evaluation = SplitAwareDAGEvaluator( + dag, modalities_for_dag, task + ).evaluate() + total_time = evaluation.representation_time + evaluation.task_time + return OptimizationResult( + dag=dag, + train_score=evaluation.scores[0].average_scores, + val_score=evaluation.scores[1].average_scores, + test_score=evaluation.scores[2].average_scores, + runtime=total_time, + representation_time=evaluation.representation_time, + task_name=task.model.name, + task_time=evaluation.task_time, + ) + fused_representation = dag.execute( list( chain.from_iterable( diff --git a/src/main/python/systemds/scuro/drsearch/node_executor.py b/src/main/python/systemds/scuro/drsearch/node_executor.py index c56a6f4ca1f..bc5cf3f6b43 100644 --- a/src/main/python/systemds/scuro/drsearch/node_executor.py +++ b/src/main/python/systemds/scuro/drsearch/node_executor.py @@ -21,7 +21,7 @@ import multiprocessing as mp import os import time -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Any, Dict, List, Optional, Union import torch @@ -188,8 +188,8 @@ def _run_node_op(): else: fusion_op = node_operation if getattr(fusion_op, "needs_training", False): - return input_mods[0].combine_with_training( - input_mods[1:], fusion_op, None + raise RuntimeError( + "Fitted representations require SplitAwareDAGEvaluator" ) return input_mods[0].combine(input_mods[1:], fusion_op) @@ -276,6 +276,7 @@ def _run_task(): return { "scores": result[0], "task_time": result[1], + "task_timing": task.last_run_timing, "memory": measurement, "peak_bytes": measurement.increment_bytes if measurement else -1, "gpu_peak_bytes": gpu_peak_bytes, @@ -406,6 +407,18 @@ class ResultEntry: task_time: float = 0.0 dag: RepresentationDag = None tradeoff_score: float = 0.0 + train_fold_scores: dict = field(default_factory=dict) + val_fold_scores: dict = field(default_factory=dict) + test_fold_scores: dict = field(default_factory=dict) + train_score_stds: dict = field(default_factory=dict) + val_score_stds: dict = field(default_factory=dict) + test_score_stds: dict = field(default_factory=dict) + task_timing: dict = field(default_factory=dict) + peak_cpu_bytes: int = -1 + peak_gpu_bytes: int = -1 + eval_index: int = -1 + t_since_search_start_s: float = 0.0 + t_eval_end_unix: float = 0.0 class NodeExecutor: @@ -460,6 +473,7 @@ def __init__( self._job_units: Dict[int, Union[_NodeUnit, _BatchUnit]] = {} self._job_retained_shm: Dict[int, List[str]] = {} self._leaf_shm_names: List[str] = [] + self._leaf_shm_owners: Dict[str, Any] = {} self._task_results: Dict[str, ResultEntry] = {} self._owns_pool = worker_pool is None @@ -543,6 +557,7 @@ def _load_leaf_modalities(self) -> None: 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 @staticmethod def _loads_in_chunks(modality: Any) -> bool: @@ -551,8 +566,17 @@ def _loads_in_chunks(modality: Any) -> bool: def _cleanup_leaf_shared_memory(self) -> None: for shm_name in self._leaf_shm_names: + modality = self._leaf_shm_owners.pop(shm_name, None) + if modality is not None and shm_name in collect_shm_names_from_payload( + getattr(modality, "_data", None) + ): + close = getattr(modality._data, "close", None) + if callable(close): + close() + modality._data = None unlink_shm(shm_name) self._leaf_shm_names = [] + self._leaf_shm_owners = {} def _submit_node( self, node_id: str, allow_gpu_worker_for_cpu: bool = False @@ -702,6 +726,19 @@ def _handle_node_success(self, node_id: str, value: Dict[str, Any]) -> None: entry.train_score = value["scores"][0].average_scores entry.val_score = value["scores"][1].average_scores entry.test_score = value["scores"][2].average_scores + entry.train_fold_scores = value["scores"][0].fold_scores() + entry.val_fold_scores = value["scores"][1].fold_scores() + entry.test_fold_scores = value["scores"][2].fold_scores() + entry.train_score_stds = value["scores"][0].score_stds() + entry.val_score_stds = value["scores"][1].score_stds() + entry.test_score_stds = value["scores"][2].score_stds() + entry.task_timing = value.get("task_timing", {}) + entry.peak_cpu_bytes = value.get("peak_bytes", -1) + entry.peak_gpu_bytes = value.get("gpu_peak_bytes", -1) + entry.eval_index = self._eval_counter + entry.t_since_search_start_s = time.perf_counter() - self._search_start + entry.t_eval_end_unix = time.time() + self._eval_counter += 1 if self.enable_checkpointing: self._checkpoint_manager.increment(node_id) self._checkpoint_manager.checkpoint_if_due( diff --git a/src/main/python/systemds/scuro/drsearch/representation_dag.py b/src/main/python/systemds/scuro/drsearch/representation_dag.py index 6315e853dda..de14b462735 100644 --- a/src/main/python/systemds/scuro/drsearch/representation_dag.py +++ b/src/main/python/systemds/scuro/drsearch/representation_dag.py @@ -379,15 +379,11 @@ def execute_node(node_id: str, task) -> TransformedModality: else: # It's a fusion operation fusion_op = node_operation - if ( - hasattr(fusion_op, "needs_training") - and fusion_op.needs_training - ): - result = input_mods[0].combine_with_training( - input_mods[1:], fusion_op, task + if getattr(fusion_op, "needs_training", False): + raise RuntimeError( + "Fitted representations require SplitAwareDAGEvaluator" ) - else: - result = input_mods[0].combine(input_mods[1:], fusion_op) + result = input_mods[0].combine(input_mods[1:], fusion_op) else: node_operation = node.operation(params=node.parameters) if gpu_id is not None and hasattr(node_operation, "gpu_id"): @@ -411,15 +407,11 @@ def execute_node(node_id: str, task) -> TransformedModality: else: # It's a fusion operation fusion_op = node_operation - if ( - hasattr(fusion_op, "needs_training") - and fusion_op.needs_training - ): - result = input_mods[0].combine_with_training( - input_mods[1:], fusion_op, task + if getattr(fusion_op, "needs_training", False): + raise RuntimeError( + "Fitted representations require SplitAwareDAGEvaluator" ) - else: - result = input_mods[0].combine(input_mods[1:], fusion_op) + result = input_mods[0].combine(input_mods[1:], fusion_op) if ( enable_cache @@ -475,6 +467,59 @@ def get_leaf_node_id(self) -> str: return None +def compose_dag_leaves( + dag: RepresentationDag, leaf_dags: Dict[Tuple[Any, Any], RepresentationDag] +) -> RepresentationDag: + """Replace indexed leaves with their complete source representation DAGs.""" + composed_nodes = [] + replacement_roots = {} + + for leaf_id in dag.get_leaf_nodes(): + leaf = dag.get_node_by_id(leaf_id) + key = (leaf.modality_id, leaf.representation_index) + source_dag = leaf_dags.get(key) + if source_dag is None: + raise KeyError( + f"No source representation DAG found for multimodal leaf {key}" + ) + + id_map = { + node.node_id: f"{leaf.node_id}__{node.node_id}" + for node in source_dag.nodes + } + for source_node in source_dag.nodes: + cloned = copy.deepcopy(source_node) + cloned.node_id = id_map[source_node.node_id] + cloned.inputs = [id_map[input_id] for input_id in source_node.inputs] + composed_nodes.append(cloned) + replacement_roots[leaf.node_id] = id_map[source_dag.root_node_id] + + internal_nodes = [node for node in dag.nodes if node.inputs] + internal_ids = { + node.node_id: f"fusion__{node.node_id}" for node in internal_nodes + } + for node in internal_nodes: + cloned = copy.deepcopy(node) + cloned.node_id = internal_ids[node.node_id] + cloned.inputs = [ + replacement_roots[input_id] + if input_id in replacement_roots + else internal_ids[input_id] + for input_id in node.inputs + ] + composed_nodes.append(cloned) + + root_node_id = ( + replacement_roots[dag.root_node_id] + if dag.root_node_id in replacement_roots + else internal_ids[dag.root_node_id] + ) + composed = RepresentationDag(composed_nodes, root_node_id, dag_id=dag.dag_id) + if not composed.validate(): + raise ValueError("Composed representation DAG is invalid") + return composed + + def _spec_safe(value): if isinstance(value, (str, int, float, bool)) or value is None: return value 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 31a1566dea1..12c11cee44d 100644 --- a/src/main/python/systemds/scuro/drsearch/split_aware_evaluator.py +++ b/src/main/python/systemds/scuro/drsearch/split_aware_evaluator.py @@ -37,6 +37,7 @@ from systemds.scuro.representations.dimensionality_reduction import ( DimensionalityReduction, ) +from systemds.scuro.representations.fusion import Fusion from systemds.scuro.representations.representation import ( FitRequirement, Representation, @@ -61,17 +62,23 @@ def dag_requires_fitting(dag: RepresentationDag) -> bool: class SplitAwareDAGEvaluator: - def __init__(self, dag: RepresentationDag, modality, task): + def __init__(self, dag: RepresentationDag, modalities, task): self.dag = dag - self.modality = modality + self.modalities = ( + list(modalities) if isinstance(modalities, (list, tuple)) else [modalities] + ) self.task = task @staticmethod def materialize(modality) -> None: + loader = getattr(modality, "data_loader", None) if modality.has_data(): - return + if loader is None: + return + expected_instances = len(getattr(loader, "indices", ())) + if not expected_instances or len(modality.data) == expected_instances: + return - loader = getattr(modality, "data_loader", None) if loader is None: raise ValueError("Split-aware evaluation requires materialized data") @@ -93,7 +100,13 @@ def _wrap(transformed, instance_ids) -> IndexedModality: def _apply_stateless(self, operation, inputs, node): aggregation = pushdown_aggregation_for_node(node.parameters) - if isinstance(operation, Context): + if isinstance(operation, Fusion): + instance_ids = operation._validate_indexed_modalities(inputs) + transformed = inputs[0].modality.combine( + [item.modality for item in inputs[1:]], operation + ) + return self._wrap(transformed, instance_ids) + elif isinstance(operation, Context): transformed = inputs.modality.context(operation) elif isinstance(operation, DimensionalityReduction): if getattr(operation, "needs_training", False): @@ -112,10 +125,33 @@ def _apply_stateless(self, operation, inputs, node): operation, aggregation=aggregation ) else: - raise TypeError(f"Unsupported unimodal operation: {type(operation).__name__}") + raise TypeError( + f"Unsupported unimodal operation: {type(operation).__name__}" + ) return self._wrap(transformed, inputs.instance_ids) + @staticmethod + def _leaf_input(partition, node): + if len(partition) == 1: + return partition[0] + modality_id = node.modality_id + try: + modality_id = int(modality_id) + except (TypeError, ValueError): + pass + matches = [ + item for item in partition if item.modality.modality_id == modality_id + ] + if not matches: + raise KeyError(f"No modality found for leaf {node.node_id}") + if node.representation_index in (None, -1) or len(matches) == 1: + return matches[0] + index = int(node.representation_index) + if index >= len(matches): + raise IndexError(f"No representation {index} for leaf {node.node_id}") + return matches[index] + def _execute_partitions(self, *partitions): cache = {} @@ -125,27 +161,34 @@ def execute_node(node_id): node = self.dag.get_node_by_id(node_id) if not node.inputs: - result = partitions + result = tuple( + self._leaf_input(partition, node) for partition in partitions + ) cache[node_id] = result return result - if len(node.inputs) != 1: - raise NotImplementedError( - "Multimodal split-aware execution is handled in step 4" - ) - input_partitions = execute_node(node.inputs[0]) - operation = node.operation(params=node.parameters) - requirement = getattr( - operation, "fit_requirement", FitRequirement.NONE + input_results = [execute_node(input_id) for input_id in node.inputs] + input_partitions = tuple( + tuple(result[i] for result in input_results) + for i in range(len(partitions)) + ) + input_partitions = tuple( + items[0] if len(items) == 1 else items for items in input_partitions ) + operation = node.operation(params=node.parameters) + requirement = getattr(operation, "fit_requirement", FitRequirement.NONE) if requirement is not FitRequirement.NONE: labels = None + fit_input = input_partitions[0] + fit_ids = ( + fit_input[0].instance_ids + if isinstance(fit_input, tuple) + else fit_input.instance_ids + ) if requirement is FitRequirement.SUPERVISED: - labels = np.asarray(self.task.labels)[ - list(input_partitions[0].instance_ids) - ] - fit_state = operation.fit(input_partitions[0], labels=labels) + labels = np.asarray(self.task.labels)[list(fit_ids)] + fit_state = operation.fit(fit_input, labels=labels) aggregation = pushdown_aggregation_for_node(node.parameters) result = tuple( operation.transform_fitted( @@ -166,11 +209,17 @@ def execute_node(node_id): return execute_node(self.dag.root_node_id) + def _indexed_modalities(self): + full = [] + for modality in self.modalities: + self.materialize(modality) + full.append(IndexedModality(modality, range(len(modality.data)))) + return tuple(full) + def fit_transform(self, fit_indices, transform_indices) -> IndexedModality: - self.materialize(self.modality) - full = IndexedModality(self.modality, range(len(self.modality.data))) - fit_partition = full.subset(fit_indices) - transform_partition = full.subset(transform_indices) + full = self._indexed_modalities() + fit_partition = tuple(item.subset(fit_indices) for item in full) + transform_partition = tuple(item.subset(transform_indices) for item in full) _, transformed = self._execute_partitions( fit_partition, transform_partition, @@ -182,12 +231,9 @@ def _partition_data(partition): return list(partition.modality.data) def evaluate(self) -> SplitEvaluationResult: - self.materialize(self.modality) - full = IndexedModality(self.modality, range(len(self.modality.data))) + full = self._indexed_modalities() - train_scores = PerformanceMeasure( - "train", self.task.performance_measures - ) + train_scores = PerformanceMeasure("train", self.task.performance_measures) val_scores = PerformanceMeasure("val", self.task.performance_measures) test_scores = PerformanceMeasure("test", self.task.performance_measures) self.task.training_time = [] @@ -196,9 +242,9 @@ def evaluate(self) -> SplitEvaluationResult: representation_time = 0.0 task_time = 0.0 for split in self.task.evaluation_splits: - train = full.subset(split.fit_indices) - evaluation = full.subset(split.evaluation_indices) - test = full.subset(self.task.test_indices) + train = tuple(item.subset(split.fit_indices) for item in full) + evaluation = tuple(item.subset(split.evaluation_indices) for item in full) + test = tuple(item.subset(self.task.test_indices) for item in full) start = time.perf_counter() transformed_train, transformed_evaluation, transformed_test = ( @@ -211,9 +257,7 @@ def evaluate(self) -> SplitEvaluationResult: evaluation_X = self._partition_data(transformed_evaluation) test_X = self._partition_data(transformed_test) train_y = np.asarray(self.task.labels)[list(split.fit_indices)] - evaluation_y = np.asarray(self.task.labels)[ - list(split.evaluation_indices) - ] + evaluation_y = np.asarray(self.task.labels)[list(split.evaluation_indices)] test_y = np.asarray(self.task.labels)[list(self.task.test_indices)] start = time.perf_counter() diff --git a/src/main/python/systemds/scuro/drsearch/task.py b/src/main/python/systemds/scuro/drsearch/task.py index b98b152ca04..164830c4042 100644 --- a/src/main/python/systemds/scuro/drsearch/task.py +++ b/src/main/python/systemds/scuro/drsearch/task.py @@ -87,7 +87,6 @@ def __init__( kfold=5, measure_performance=True, performance_measures=["accuracy"], - fusion_train_split=0.8, ): """ Parent class for the prediction task that is performed on top of the aligned representation @@ -114,7 +113,6 @@ def __init__( self.train_scores = PerformanceMeasure("train", performance_measures) self.val_scores = PerformanceMeasure("val", performance_measures) self.test_scores = PerformanceMeasure("test", performance_measures) - self.fusion_train_indices = None self._create_cv_splits() def get_output_stats(self, input_stats): @@ -151,11 +149,6 @@ def estimate_peak_memory_bytes(self, input_stats): + n_test + sum(len(fold) for fold in self.cv_train_indices) + sum(len(fold) for fold in self.cv_val_indices) - + ( - len(self.fusion_train_indices) - if self.fusion_train_indices is not None - else 0 - ) ) index_bytes = total_index_entries * np.dtype(np.int64).itemsize @@ -245,13 +238,6 @@ def _create_cv_splits(self): f"Fold {fold_idx}: Overlap detected between train and val indices: {overlap}" ) - all_val_indices = set() - for val_indices in self.cv_val_indices: - all_val_indices.update(val_indices) - - self.fusion_train_indices = [ - idx for idx in self.train_indices if idx not in all_val_indices - ] self.evaluation_splits = tuple( EvaluationSplit(train_indices, val_indices, name=f"fold_{fold_idx}") for fold_idx, (train_indices, val_indices) in enumerate( diff --git a/src/main/python/systemds/scuro/drsearch/test_set_evaluation.py b/src/main/python/systemds/scuro/drsearch/test_set_evaluation.py index f160cd4ccd0..4479727e9d1 100644 --- a/src/main/python/systemds/scuro/drsearch/test_set_evaluation.py +++ b/src/main/python/systemds/scuro/drsearch/test_set_evaluation.py @@ -25,6 +25,10 @@ from systemds.scuro.modality.modality import Modality from systemds.scuro.drsearch.representation_dag import RepresentationDag +from systemds.scuro.drsearch.split_aware_evaluator import ( + SplitAwareDAGEvaluator, + dag_requires_fitting, +) from systemds.scuro.representations.aggregated_representation import ( AggregatedRepresentation, ) @@ -60,9 +64,17 @@ def measure_representation_time_on_test_set( output = None for _ in range(max(1, repeats)): t0 = time.perf_counter() - output = _match_expected_dim( - _unwrap(dag.execute(subsets, task, enable_cache=False)), task - ) + if dag_requires_fitting(dag): + if task is None: + raise ValueError("A task is required to fit the representation DAG") + transformed = SplitAwareDAGEvaluator(dag, modalities, task).fit_transform( + task.train_indices, test_indices + ) + output = _match_expected_dim(transformed.modality, task) + else: + output = _match_expected_dim( + _unwrap(dag.execute(subsets, task, enable_cache=False)), task + ) timings.append(time.perf_counter() - t0) timings.sort() @@ -103,9 +115,15 @@ def measure_test_set_application( if full_data is None: t0 = time.perf_counter() - output = _match_expected_dim( - _unwrap(dag.execute(modalities, task, enable_cache=False)), task - ) + if dag_requires_fitting(dag): + transformed = SplitAwareDAGEvaluator(dag, modalities, task).fit_transform( + task.train_indices, range(len(modalities[0].data)) + ) + output = _match_expected_dim(transformed.modality, task) + else: + output = _match_expected_dim( + _unwrap(dag.execute(modalities, task, enable_cache=False)), task + ) record["full_representation_time_s"] = time.perf_counter() - t0 full_data = None if output is None else output.data diff --git a/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py b/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py index 1eaf9303fdd..fabe6f1e6fc 100644 --- a/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py +++ b/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py @@ -405,6 +405,13 @@ def _process_modality(self, modality, skip_remaining: int = 0, scheduler=None): train_score=evaluation.scores[0].average_scores, val_score=evaluation.scores[1].average_scores, test_score=evaluation.scores[2].average_scores, + train_fold_scores=evaluation.scores[0].fold_scores(), + val_fold_scores=evaluation.scores[1].fold_scores(), + test_fold_scores=evaluation.scores[2].fold_scores(), + train_score_stds=evaluation.scores[0].score_stds(), + val_score_stds=evaluation.scores[1].score_stds(), + test_score_stds=evaluation.scores[2].score_stds(), + task_timing=getattr(task, "last_run_timing", {}), representation_time=evaluation.representation_time, task_time=evaluation.task_time, dag=result_dag, @@ -948,6 +955,9 @@ def add_result( train_fold_scores=scores[0].fold_scores(), val_fold_scores=scores[1].fold_scores(), test_fold_scores=scores[2].fold_scores(), + train_score_stds=scores[0].score_stds(), + val_score_stds=scores[1].score_stds(), + test_score_stds=scores[2].score_stds(), representation_time=transform_time, task_time=task_time, dag=dag, diff --git a/src/main/python/systemds/scuro/modality/transformed.py b/src/main/python/systemds/scuro/modality/transformed.py index e02695166db..62cc03eec04 100644 --- a/src/main/python/systemds/scuro/modality/transformed.py +++ b/src/main/python/systemds/scuro/modality/transformed.py @@ -191,22 +191,6 @@ def combine(self, other: Union[Modality, List[Modality]], fusion_method): fused_modality.transform_time = end_time - start_time return fused_modality - def combine_with_training( - self, other: Union[Modality, List[Modality]], fusion_method, task - ): - fused_modality = TransformedModality( - self, fusion_method, ModalityType.EMBEDDING - ) - modalities = self.create_modality_list(other) - start_time = time.time() - fused_modality.data = fusion_method.transform_with_training(modalities, task) - end_time = time.time() - fused_modality.transform_time = ( - end_time - start_time - ) # Note: this incldues the training time - - return fused_modality - def create_modality_list(self, other: Union[Modality, List[Modality]]): modalities = [self] if isinstance(other, list): diff --git a/src/main/python/systemds/scuro/representations/dimensionality_reduction.py b/src/main/python/systemds/scuro/representations/dimensionality_reduction.py index 731b3b9523a..3c8351c9451 100644 --- a/src/main/python/systemds/scuro/representations/dimensionality_reduction.py +++ b/src/main/python/systemds/scuro/representations/dimensionality_reduction.py @@ -22,8 +22,14 @@ import numpy as np +from systemds.scuro.modality.indexed import IndexedModality from systemds.scuro.modality.modality import Modality -from systemds.scuro.representations.representation import Representation +from systemds.scuro.modality.transformed import TransformedModality +from systemds.scuro.representations.representation import ( + FitRequirement, + Representation, + RepresentationFitState, +) class DimensionalityReduction(Representation): @@ -43,10 +49,7 @@ def execute(self, data, labels=None): :param labels: labels for learned dimensionality reduction :return: dimensionality reduced data """ - if labels is not None: - self.execute_with_training(data, labels) - else: - self.execute(data) + raise NotImplementedError(f"execute is not implemented for {self.name}") def apply_representation(self, data): """ @@ -56,26 +59,48 @@ def apply_representation(self, data): """ raise f"Not implemented for Dimensionality Reduction Operator: {self.name}" - def execute_with_training(self, modality, task): - fusion_train_indices = task.fusion_train_indices - # Handle 3d data - data = modality.data - if ( - len(np.array(modality.data).shape) == 3 - and np.array(modality.data).shape[1] == 1 - ): - data = np.array([x.reshape(-1) for x in modality.data]) - transformed_train = self.execute( - np.array(data)[fusion_train_indices], task.labels[fusion_train_indices] - ) + @staticmethod + def _prepare_data(modality): + data = np.asarray(modality.data) + if data.ndim == 3 and data.shape[1] == 1: + return data.reshape(data.shape[0], -1) + return data + + def fit(self, inputs, labels=None) -> RepresentationFitState: + if self.fit_requirement is not FitRequirement.SUPERVISED: + return super().fit(inputs, labels) + if not isinstance(inputs, IndexedModality): + raise TypeError( + "Trainable dimensionality reduction expects an IndexedModality" + ) + if labels is None: + raise ValueError(f"{self.name} requires labels for fitting") + if len(labels) != len(inputs): + raise ValueError( + f"Expected {len(inputs)} labels for {self.name}, got {len(labels)}" + ) - all_other_indices = [ - i for i in range(len(modality.data)) if i not in fusion_train_indices - ] - transformed_other = self.apply_representation(np.array(data)[all_other_indices]) + self.execute(self._prepare_data(inputs.modality), np.asarray(labels)) + return RepresentationFitState( + self.fit_requirement, + fit_instance_ids=inputs.instance_ids, + payload=self, + ) - transformed_data = np.zeros((len(data), transformed_train.shape[1])) - transformed_data[fusion_train_indices] = transformed_train - transformed_data[all_other_indices] = transformed_other + def transform_fitted(self, inputs, fit_state, **kwargs): + if not isinstance(inputs, IndexedModality): + raise TypeError( + "Trainable dimensionality reduction expects an IndexedModality" + ) + if ( + not isinstance(fit_state, RepresentationFitState) + or fit_state.fit_requirement is not self.fit_requirement + or not isinstance(fit_state.payload, type(self)) + ): + raise ValueError(f"{self.name} requires a compatible fitted state") - return transformed_data + transformed = TransformedModality(inputs.modality, fit_state.payload) + transformed.data = fit_state.payload.apply_representation( + self._prepare_data(inputs.modality) + ) + return IndexedModality(transformed, inputs.instance_ids) diff --git a/src/main/python/systemds/scuro/representations/fusion.py b/src/main/python/systemds/scuro/representations/fusion.py index e56eb465bc6..496fc2fba1f 100644 --- a/src/main/python/systemds/scuro/representations/fusion.py +++ b/src/main/python/systemds/scuro/representations/fusion.py @@ -23,6 +23,7 @@ import numpy as np +from systemds.scuro.modality.indexed import IndexedModality from systemds.scuro.modality.type import ModalityType from systemds.scuro.representations.aggregated_representation import ( AggregatedRepresentation, @@ -32,7 +33,9 @@ from systemds.scuro.modality.modality import Modality from systemds.scuro.representations.representation import ( CONTAINER_ARRAY, + FitRequirement, Representation, + RepresentationFitState, RepresentationStats, derive_stats, stats_bytes, @@ -97,10 +100,19 @@ def transform(self, modalities: List[Modality]): :param modalities: List of modalities used in the fusion :return: fused data """ + mods = self._prepare_modalities(modalities) + + return self.execute(mods) + + def _prepare_modalities(self, modalities: List[Modality]): mods = [] for modality in modalities: agg_modality = None - if not self.preserves_leading_axis and get_shape(modality.metadata) > 1: + if ( + not self.preserves_leading_axis + and modality.has_metadata() + and get_shape(modality.metadata) > 1 + ): agg_operator = AggregatedRepresentation() agg_modality = agg_operator.transform(modality) mods.append(agg_modality if agg_modality else modality) @@ -112,36 +124,58 @@ def transform(self, modalities: List[Modality]): for modality in mods: modality.pad(max_len=max_len) - return self.execute(mods) + return mods - def transform_with_training(self, modalities: List[Modality], task): - fusion_train_indices = task.fusion_train_indices + @staticmethod + def _validate_indexed_modalities(inputs): + if not isinstance(inputs, (list, tuple)) or not inputs: + raise TypeError( + "Fusion fitting expects a non-empty sequence of IndexedModality inputs" + ) + if not all(isinstance(item, IndexedModality) for item in inputs): + raise TypeError("Fusion fitting expects IndexedModality inputs") + + instance_ids = inputs[0].instance_ids + if any(item.instance_ids != instance_ids for item in inputs[1:]): + raise ValueError("Fusion inputs must have aligned instance IDs") + return instance_ids + + def fit(self, inputs, labels=None) -> RepresentationFitState: + if self.fit_requirement is not FitRequirement.SUPERVISED: + return super().fit(inputs, labels) + instance_ids = self._validate_indexed_modalities(inputs) + if labels is None: + raise ValueError(f"{self.name} requires labels for fitting") + if len(labels) != len(instance_ids): + raise ValueError( + f"Expected {len(instance_ids)} labels for {self.name}, got {len(labels)}" + ) - train_modalities = [] - for modality in modalities: - train_data = [ - d for i, d in enumerate(modality.data) if i in fusion_train_indices - ] - train_modality = TransformedModality(modality, self) - train_modality.data = list(train_data) - train_modalities.append(train_modality) - - transformed_train = self.execute( - train_modalities, task.labels[fusion_train_indices] + modalities = self._prepare_modalities([item.modality for item in inputs]) + self.execute(modalities, np.asarray(labels)) + return RepresentationFitState( + self.fit_requirement, + fit_instance_ids=instance_ids, + payload=self, ) - all_other_indices = [ - i for i in range(len(modalities[0].data)) if i not in fusion_train_indices - ] - transformed_other = self.transform_data(modalities, all_other_indices) - - transformed_data = np.zeros( - (len(modalities[0].data), transformed_train.shape[1]) + def transform_fitted(self, inputs, fit_state, **kwargs): + instance_ids = self._validate_indexed_modalities(inputs) + if ( + not isinstance(fit_state, RepresentationFitState) + or fit_state.fit_requirement is not self.fit_requirement + or not isinstance(fit_state.payload, type(self)) + ): + raise ValueError(f"{self.name} requires a compatible fitted state") + + modalities = fit_state.payload._prepare_modalities( + [item.modality for item in inputs] ) - transformed_data[fusion_train_indices] = transformed_train - transformed_data[all_other_indices] = transformed_other - - return transformed_data + transformed = TransformedModality( + modalities[0], fit_state.payload, ModalityType.EMBEDDING + ) + transformed.data = fit_state.payload.apply_representation(modalities) + return IndexedModality(transformed, instance_ids) def transform_data(self, modalities: List[Modality], indices=None): val_modalities = [] diff --git a/src/main/python/systemds/scuro/representations/mlp_learned_dim_reduction.py b/src/main/python/systemds/scuro/representations/mlp_learned_dim_reduction.py index 9c24dfdb4db..25f449cce11 100644 --- a/src/main/python/systemds/scuro/representations/mlp_learned_dim_reduction.py +++ b/src/main/python/systemds/scuro/representations/mlp_learned_dim_reduction.py @@ -72,7 +72,7 @@ def __init__(self, output_dim=256, batch_size=32, learning_rate=0.001, epochs=5) self.epochs = epochs self.model = None - def execute_with_training(self, data, labels): + def execute(self, data, labels=None): if labels is None: raise ValueError("MLP labels requires labels for training") @@ -89,13 +89,11 @@ def execute_with_training(self, data, labels): self.num_classes = len(np.unique(y)) input_dim = X.shape[1] - self.device = get_device_for_model(self.model, memory_factor=1.5) - self.model = self.model.to(self.device) self.is_trained = False - self.model = self._build_model(input_dim, self.output_dim, self.num_classes).to( - self.device - ) + self.model = self._build_model(input_dim, self.output_dim, self.num_classes) + self.device = get_device_for_model(self.model, memory_factor=1.5) + self.model = self.model.to(self.device) if self.is_multilabel: criterion = nn.BCEWithLogitsLoss() else: @@ -108,7 +106,7 @@ def execute_with_training(self, data, labels): else: y_tensor = torch.LongTensor(y) - sample = data[0] if data else "" + sample = data[0] if len(data) > 0 else "" self.batch_size = compute_batch_size( model=self.model, device=self.device, From 304ce04789b7bd8c7fde77d050ef201ab48e19fd Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Tue, 8 Sep 2026 16:10:45 +0200 Subject: [PATCH 06/15] add feature dim limit and statistics --- .../scuro/drsearch/hyperparameter_tuner.py | 12 +- .../scuro/drsearch/multimodal_ga_optimizer.py | 137 ++++++++++++- .../systemds/scuro/drsearch/node_executor.py | 52 ++++- .../systemds/scuro/drsearch/node_scheduler.py | 48 +++++ .../scuro/drsearch/representation_dag.py | 15 +- .../scuro/drsearch/split_aware_evaluator.py | 7 + .../python/systemds/scuro/drsearch/task.py | 60 ++++++ .../scuro/drsearch/unimodal_optimizer.py | 8 +- .../systemds/scuro/models/training_stats.py | 184 +++++++++++++++++ .../scuro/representations/aggregate.py | 6 - .../systemds/scuro/representations/fusion.py | 76 ++++--- .../systemds/scuro/representations/lstm.py | 192 +++++++++++++++--- .../multimodal_attention_fusion.py | 116 ++++++++++- 13 files changed, 819 insertions(+), 94 deletions(-) create mode 100644 src/main/python/systemds/scuro/models/training_stats.py diff --git a/src/main/python/systemds/scuro/drsearch/hyperparameter_tuner.py b/src/main/python/systemds/scuro/drsearch/hyperparameter_tuner.py index 529e38032de..43bedeeab0c 100644 --- a/src/main/python/systemds/scuro/drsearch/hyperparameter_tuner.py +++ b/src/main/python/systemds/scuro/drsearch/hyperparameter_tuner.py @@ -528,9 +528,7 @@ def visit_node(node_id): ) baseline = ( baseline_params, - HyperparamEvaluationScores( - baseline_raw_scores, self.scoring_metric - ), + HyperparamEvaluationScores(baseline_raw_scores, self.scoring_metric), ) if not hyperparams: @@ -958,15 +956,11 @@ def objective(trial: optuna.Trial) -> float: if np.isnan(val_score): raise optuna.TrialPruned() - evaluation_scores = HyperparamEvaluationScores( - scores, self.scoring_metric - ) + evaluation_scores = HyperparamEvaluationScores(scores, self.scoring_metric) seen[self._config_key(params)] = (params, evaluation_scores) trial_results.append((params, evaluation_scores)) - val_folds = evaluation_scores.val_fold_scores.get( - self.scoring_metric, [] - ) + val_folds = evaluation_scores.val_fold_scores.get(self.scoring_metric, []) if val_folds is not None and len(val_folds) > 1: lam = 0.5 robust_val = np.mean(val_folds) - lam * np.std(val_folds, ddof=1) 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 62bf203f4f8..ba170bccd01 100644 --- a/src/main/python/systemds/scuro/drsearch/multimodal_ga_optimizer.py +++ b/src/main/python/systemds/scuro/drsearch/multimodal_ga_optimizer.py @@ -22,6 +22,7 @@ from __future__ import annotations import copy +import gc import os import pickle import random @@ -33,6 +34,7 @@ from itertools import chain from typing import Any, Dict, List, Optional, Tuple, Union +import numpy as np from deap import base, tools from systemds.scuro.drsearch.modality_shared_memory import ( @@ -55,10 +57,13 @@ from systemds.scuro.representations.aggregated_representation import ( AggregatedRepresentation, ) +from systemds.scuro.representations.representation import RepresentationStats from systemds.scuro.utils.schema_helpers import get_shape Tree = Union[int, Tuple["Tree", "Tree"]] +DEFAULT_MAX_FUSED_FEATURE_DIM = 100_000 + def _collect_internal_paths(tree: Tree, path: str = "") -> List[str]: if isinstance(tree, int): @@ -184,6 +189,16 @@ def _failure_fitness(objective_specs: List[ObjectiveSpec]) -> Tuple[float, ...]: ) +def _release_cuda_memory() -> None: + gc.collect() + try: + import torch + + if torch.cuda.is_available(): + torch.cuda.empty_cache() + except (ImportError, RuntimeError): + pass + def _evaluate_genome_body( dag: RepresentationDag, @@ -262,10 +277,20 @@ def _evaluate_genome_body( def _dispatch_genome_evaluation(payload, _gpu_id): dag, task, modalities, objective_specs = payload - fitness, result = _evaluate_genome_body(dag, task, modalities, objective_specs) - if fitness is None: - fitness = _failure_fitness(objective_specs) - return fitness, result + try: + try: + fitness, result = _evaluate_genome_body( + dag, task, modalities, objective_specs + ) + if fitness is None: + fitness = _failure_fitness(objective_specs) + return fitness, result + except Exception: + return _failure_fitness(objective_specs), { + "_evaluation_error": traceback.format_exc() + } + finally: + _release_cuda_memory() _WORKER_DISPATCH = {"genome": _dispatch_genome_evaluation} @@ -303,11 +328,14 @@ def __init__( hall_of_fame_size: int = 5, allow_repeated_modalities: bool = False, threads_per_worker: Optional[int] = None, + max_feature_dim: Optional[int] = DEFAULT_MAX_FUSED_FEATURE_DIM, ): self.modalities = modalities self.tasks = tasks self.debug = debug self.allow_repeated_modalities = allow_repeated_modalities + self.max_feature_dim = None if max_feature_dim is None else int(max_feature_dim) + self.oversized_candidates: Dict[str, int] = {} self.min_modalities = max(1, min_modalities) requested_max = max_modalities or len(modalities) @@ -426,6 +454,7 @@ def _optimize( self._current_task_name = task_name self.optimization_results.setdefault(task_name, []) self.evaluation_errors.setdefault(task_name, 0) + self.oversized_candidates.setdefault(task_name, 0) self._eval_counter = 0 self._search_start = time.perf_counter() @@ -471,7 +500,8 @@ def _optimize( print( f"[GA] task={task_name} gen={gen} {debug_msg} " f"no_improve={no_improve} " - f"errors={self.evaluation_errors.get(task_name, 0)}" + f"errors={self.evaluation_errors.get(task_name, 0)} " + f"oversized={self.oversized_candidates.get(task_name, 0)}" ) stagnated = ( @@ -670,6 +700,10 @@ def _evaluate_individuals_parallel( if sig in pending_followers: pending_followers[sig].append(ind) continue + reason = self._oversized_reason(task_name, genome) + if reason is not None: + ind.fitness.values = self._reject_oversized(task_name, genome, reason) + continue dag = self._genome_to_dag(genome) pending_followers[sig] = [] pending_work.append((ind, dag, sig)) @@ -692,7 +726,13 @@ def _evaluate_individuals_parallel( ind, dag, sig = jobs.pop(jr.job_id) if jr.ok: fitness, payload = jr.value - error = None + error = ( + payload.get("_evaluation_error") + if isinstance(payload, dict) + else None + ) + if error is not None: + payload = None else: fitness = _failure_fitness(self.objective_specs) payload = None @@ -719,9 +759,7 @@ def _record_evaluation( self.evaluation_errors.get(task_name, 0) + 1 ) if self.debug and error is not None: - print( - f"[GA] genome evaluation failed for task={task_name}:\n{error}" - ) + print(f"[GA] genome evaluation failed for task={task_name}:\n{error}") return result = FusionSearchResult( @@ -936,6 +974,81 @@ def norm(t: Tree): tuple(sorted((p, c.__name__) for p, c in g.fusion_ops.items())), ) + def _leaf_stats(self, task_name: str, leaf: Tuple[Any, int]): + """Shape of one candidate leaf representation, or None if unknown.""" + modality_id, representation_index = leaf + try: + modality = self.k_best_representations[task_name][modality_id][ + representation_index + ] + except (KeyError, IndexError, TypeError): + return None + + data = getattr(modality, "data", None) + if data is None or len(data) == 0: + return None + try: + sample = np.asarray(data[0]) + except (TypeError, ValueError): + return None + if sample.dtype == object: + return None + shape = tuple(sample.shape) if sample.ndim else (1,) + return RepresentationStats(len(data), shape, output_shape_is_known=True) + + def _predicted_feature_dim( + self, task_name: str, genome: DagGenome + ) -> Optional[int]: + if task_name is None: + return None + + leaf_stats = [self._leaf_stats(task_name, leaf) for leaf in genome.leaves] + if any(stats is None for stats in leaf_stats): + return None + + def walk(subtree: Tree, path: str): + if isinstance(subtree, int): + return leaf_stats[subtree] + left = walk(subtree[0], path + "L") + right = walk(subtree[1], path + "R") + if left is None or right is None: + return None + return genome.fusion_ops[path]().get_output_stats([left, right]) + + try: + root = walk(genome.tree, "") + except Exception: + return None + if root is None or not root.output_shape: + return None + return int(np.prod(root.output_shape)) + + def _oversized_reason(self, task_name: str, genome: DagGenome) -> Optional[str]: + if not self.max_feature_dim: + return None + predicted = self._predicted_feature_dim(task_name, genome) + if predicted is None or predicted <= self.max_feature_dim: + return None + return ( + f"fused representation would have {predicted:,} features, above the " + f"configured maximum of {self.max_feature_dim:,}; candidate rejected " + "before execution" + ) + + def _reject_oversized( + self, task_name: str, genome: DagGenome, reason: str + ) -> Tuple[float, ...]: + self.oversized_candidates[task_name] = ( + self.oversized_candidates.get(task_name, 0) + 1 + ) + if self.debug: + print(f"[GA] task={task_name} skipped candidate: {reason}") + fitness = _failure_fitness(self.objective_specs) + self._fitness_cache.setdefault(task_name, {})[ + self._genome_signature(genome) + ] = fitness + return fitness + def _evaluate_genome(self, genome: DagGenome, task: Task) -> Tuple[float, ...]: task_name = task.model.name sig = self._genome_signature(genome) @@ -943,6 +1056,10 @@ def _evaluate_genome(self, genome: DagGenome, task: Task) -> Tuple[float, ...]: if sig in cache: return cache[sig] + reason = self._oversized_reason(task_name, genome) + if reason is not None: + return self._reject_oversized(task_name, genome, reason) + dag = self._genome_to_dag(genome) modalities = list( chain.from_iterable(self.k_best_representations[task_name].values()) @@ -958,6 +1075,8 @@ def _evaluate_genome(self, genome: DagGenome, task: Task) -> Tuple[float, ...]: except Exception: fitness = _failure_fitness(self.objective_specs) payload, error = None, traceback.format_exc() + finally: + _release_cuda_memory() self._record_evaluation(task_name, dag, payload, error) cache[sig] = fitness diff --git a/src/main/python/systemds/scuro/drsearch/node_executor.py b/src/main/python/systemds/scuro/drsearch/node_executor.py index bc5cf3f6b43..a707db1cd24 100644 --- a/src/main/python/systemds/scuro/drsearch/node_executor.py +++ b/src/main/python/systemds/scuro/drsearch/node_executor.py @@ -79,9 +79,14 @@ def _place_on_device(obj: Any, gpu_id: Optional[int]) -> None: obj.device = device except (AttributeError, RuntimeError): pass - model = getattr(obj, "model", None) + + model = obj.__dict__.get("model") if model is not None and hasattr(model, "to"): - obj.model = model.to(device) + try: + obj.model = model.to(device) + except (AttributeError, RuntimeError): + # A read-only "model" means the operator owns its own placement. + pass def _run_gpu_op(fn, gpu_id: Optional[int]): @@ -461,7 +466,11 @@ def __init__( resume=False, ) self._memory_usage_data: Dict[str, Any] = {} - self.statistics = {"worker_stats": {}, "node_stats": {}} + self.statistics = {"worker_stats": {}, "node_stats": {}, "reuse": {}} + self._cache_consumptions = 0 + self._cache_misses = 0 + self._nodes_executed = 0 + self._peak_cache_bytes = 0 self._eval_counter = 0 self._nodes_executed = 0 self._search_start = ( @@ -585,7 +594,7 @@ def _submit_node( gpu_id = node.gpu_id parent_ids = self.scheduler.get_valid_parents(node_id) parent_results = ( - [self._result_cache.get(pid) for pid in parent_ids] if parent_ids else None + [self._get_cached_parent(pid) for pid in parent_ids] if parent_ids else None ) if self._is_task_node(node): @@ -668,7 +677,17 @@ def _gpu_id(entry): allow_gpu_worker_for_cpu=allow_gpu_worker_for_cpu, ) + def _get_cached_parent(self, parent_id: str): + try: + result = self._result_cache.get(parent_id) + except KeyError: + self._cache_misses += 1 + raise + self._cache_consumptions += 1 + return result + def _record_stats(self, node_id: str, pid: int, start_time: float, end_time: float): + self._nodes_executed += 1 node_stats = self.statistics["node_stats"] worker_stats = self.statistics["worker_stats"] node_stats[node_id] = {"start_time": start_time, "end_time": end_time} @@ -831,6 +850,7 @@ def _handle_modality_result( shm_bytes=shm_bytes, ) after_bytes = self._result_cache.get_memory_total_memory_usage() + self._peak_cache_bytes = max(self._peak_cache_bytes, after_bytes) self.scheduler.update_cpu_memory_in_use(after_bytes - before_bytes) self.scheduler.complete_node(node_id) @@ -1034,7 +1054,31 @@ def run(self) -> Dict[str, Any]: self._task_results, self._discard_report() ) + self.statistics["reuse"] = self._reuse_report() + return { "task_results": list(self._task_results.values()), "statistics": self.statistics, } + + def _reuse_report(self) -> Dict[str, Any]: + report = dict(getattr(self.scheduler, "reuse_stats", {}) or {}) + slots = report.get("dag_node_slots", 0) + saved = report.get("nodes_saved_by_merging", 0) + report.update( + { + "nodes_executed": self._nodes_executed, + "cache_consumptions": self._cache_consumptions, + "cache_misses": self._cache_misses, + "peak_cache_bytes": self._peak_cache_bytes, + "slot_reduction_ratio": (saved / slots) if slots else 0.0, + "failed_nodes": len(getattr(self.scheduler, "failed_nodes", []) or []), + "blocked_memory_nodes": len( + getattr(self.scheduler, "blocked_memory_nodes_perm", []) or [] + ), + "cpu_fallback_nodes": len( + getattr(self.scheduler, "cpu_fallback_nodes", []) or [] + ), + } + ) + return report diff --git a/src/main/python/systemds/scuro/drsearch/node_scheduler.py b/src/main/python/systemds/scuro/drsearch/node_scheduler.py index 8398c82a81b..44f7dfef9ff 100644 --- a/src/main/python/systemds/scuro/drsearch/node_scheduler.py +++ b/src/main/python/systemds/scuro/drsearch/node_scheduler.py @@ -569,13 +569,40 @@ def _release_execution_memory(self, node_id: str, gpu_id: int) -> None: 0, self.memory_stats["cpu_in_flight"] - reserved ) + @staticmethod + def _operation_label(node) -> str: + """A readable name for what a node computes. + + `operation` is a class on representation nodes but an instance on some + others, and task roots carry their kind in `parameters` instead, so a + bare `.__name__` silently buckets both under "unknown" -- which hides + exactly the nodes whose sharing is worth reporting. + """ + parameters = getattr(node, "parameters", None) or {} + kind = parameters.get("_node_kind") + if kind == "task": + return "task" + operation = getattr(node, "operation", None) + if operation is None: + return "leaf" if not getattr(node, "inputs", None) else "unknown" + name = getattr(operation, "__name__", None) + if name: + return name + return type(operation).__name__ + def _get_nodes_from_dags( self, dags: List[RepresentationDag] ) -> List[RepresentationNode]: nodes = [] + dag_node_slots = 0 + slot_counts = defaultdict(int) + node_operations = {} for dag in dags: self.roots.add(dag.root_node_id) for node in dag.nodes: + dag_node_slots += 1 + slot_counts[node.node_id] += 1 + node_operations[node.node_id] = self._operation_label(node) self.mapping[node.node_id] = node if not node.inputs: self.leaves.add(node.node_id) @@ -605,6 +632,27 @@ def _get_nodes_from_dags( self.remaining_children = { node_id: len(self.children.get(node_id, set())) for node_id in self.mapping } + + saved_by_operation = defaultdict(int) + for node_id, count in slot_counts.items(): + if count > 1: + saved_by_operation[node_operations.get(node_id, "unknown")] += count - 1 + self.reuse_stats = { + "n_dags": len(dags), + "dag_node_slots": dag_node_slots, + "unique_nodes": len(self.mapping), + "nodes_saved_by_merging": dag_node_slots - len(self.mapping), + "shared_nodes": sum( + 1 + for node_id in self.mapping + if len(self.children.get(node_id, set())) > 1 + ), + "max_fanout": max( + (len(children) for children in self.children.values()), default=0 + ), + "saved_by_operation": dict(saved_by_operation), + "n_leaves": len(self.leaves), + } return nodes def _get_unresolved_parents(self) -> Dict[str, int]: diff --git a/src/main/python/systemds/scuro/drsearch/representation_dag.py b/src/main/python/systemds/scuro/drsearch/representation_dag.py index de14b462735..e0c39537ae0 100644 --- a/src/main/python/systemds/scuro/drsearch/representation_dag.py +++ b/src/main/python/systemds/scuro/drsearch/representation_dag.py @@ -484,8 +484,7 @@ def compose_dag_leaves( ) id_map = { - node.node_id: f"{leaf.node_id}__{node.node_id}" - for node in source_dag.nodes + node.node_id: f"{leaf.node_id}__{node.node_id}" for node in source_dag.nodes } for source_node in source_dag.nodes: cloned = copy.deepcopy(source_node) @@ -495,16 +494,16 @@ def compose_dag_leaves( replacement_roots[leaf.node_id] = id_map[source_dag.root_node_id] internal_nodes = [node for node in dag.nodes if node.inputs] - internal_ids = { - node.node_id: f"fusion__{node.node_id}" for node in internal_nodes - } + internal_ids = {node.node_id: f"fusion__{node.node_id}" for node in internal_nodes} for node in internal_nodes: cloned = copy.deepcopy(node) cloned.node_id = internal_ids[node.node_id] cloned.inputs = [ - replacement_roots[input_id] - if input_id in replacement_roots - else internal_ids[input_id] + ( + replacement_roots[input_id] + if input_id in replacement_roots + else internal_ids[input_id] + ) for input_id in node.inputs ] composed_nodes.append(cloned) 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 12c11cee44d..39c6e1be21c 100644 --- a/src/main/python/systemds/scuro/drsearch/split_aware_evaluator.py +++ b/src/main/python/systemds/scuro/drsearch/split_aware_evaluator.py @@ -24,6 +24,7 @@ import numpy as np +from systemds.scuro.models.training_stats import TrainingWorkAccumulator from systemds.scuro.drsearch.representation_dag import ( RepresentationDag, pushdown_aggregation_for_node, @@ -238,6 +239,7 @@ def evaluate(self) -> SplitEvaluationResult: test_scores = PerformanceMeasure("test", self.task.performance_measures) self.task.training_time = [] self.task.inference_time = [] + training_work = TrainingWorkAccumulator() representation_time = 0.0 task_time = 0.0 @@ -256,6 +258,7 @@ def evaluate(self) -> SplitEvaluationResult: train_X = self._partition_data(transformed_train) evaluation_X = self._partition_data(transformed_evaluation) test_X = self._partition_data(transformed_test) + self.task.validate_input_memory(train_X) train_y = np.asarray(self.task.labels)[list(split.fit_indices)] evaluation_y = np.asarray(self.task.labels)[list(split.evaluation_indices)] test_y = np.asarray(self.task.labels)[list(self.task.test_indices)] @@ -279,9 +282,12 @@ def evaluate(self) -> SplitEvaluationResult: val_scores.add_scores(val_score[0]) test_scores.add_scores(test_score[0]) + training_work.add(model) + if hasattr(model, "clean_up"): model.clean_up() + self.task.training_work = training_work self.task.last_run_timing = { "train_time_per_fold_s": list(self.task.training_time), "test_inference_time_per_fold_s": list(self.task.inference_time), @@ -296,6 +302,7 @@ def evaluate(self) -> SplitEvaluationResult: else 0.0 ), "n_test_instances": len(self.task.test_indices), + **training_work.as_dict(), } return SplitEvaluationResult( diff --git a/src/main/python/systemds/scuro/drsearch/task.py b/src/main/python/systemds/scuro/drsearch/task.py index 164830c4042..05cbb761c88 100644 --- a/src/main/python/systemds/scuro/drsearch/task.py +++ b/src/main/python/systemds/scuro/drsearch/task.py @@ -23,6 +23,7 @@ import time from typing import List from systemds.scuro.models.model import Model +from systemds.scuro.models.training_stats import TrainingWorkAccumulator import numpy as np from sklearn.model_selection import train_test_split from systemds.scuro.drsearch.evaluation_split import EvaluationSplit @@ -108,6 +109,7 @@ def __init__( self.inference_time = [] self.training_time = [] self.last_run_timing = {} + self.training_work = TrainingWorkAccumulator() self.expected_dim = 1 self.performance_measures = performance_measures self.train_scores = PerformanceMeasure("train", performance_measures) @@ -208,6 +210,59 @@ def estimate_relative_cost(self, input_stats) -> float: max(1, self.kfold) * passes * max(1, n_fold_train) * max(1, feature_dim) ) + def validate_input_memory(self, data, gpu_memory_fraction: float = 0.9) -> None: + if data is None or len(data) == 0: + return + + if isinstance(data, np.ndarray): + num_instances = int(data.shape[0]) if data.ndim else 1 + output_shape = tuple(data.shape[1:]) if data.ndim > 1 else (1,) + dtype = data.dtype + else: + sample = np.asarray(data[0]) + num_instances = len(data) + output_shape = tuple(sample.shape) if sample.ndim else (1,) + dtype = sample.dtype + + estimate = self.estimate_peak_memory_bytes( + RepresentationStats( + num_instances, + output_shape, + output_shape_is_known=True, + dtype=dtype, + ) + ) + required = int(estimate.get("gpu_peak_bytes", 0)) + if required <= 0: + return + + try: + import torch + + if not torch.cuda.is_available(): + return + device = getattr(self.model, "device", torch.device("cuda")) + device_index = ( + device.index if getattr(device, "index", None) is not None else 0 + ) + free, total = torch.cuda.mem_get_info(device_index) + cached = torch.cuda.memory_reserved( + device_index + ) - torch.cuda.memory_allocated(device_index) + except (ImportError, RuntimeError): + return + + available = int((free + max(cached, 0)) * gpu_memory_fraction) + if required > available: + feature_dim = int(np.prod(output_shape)) if output_shape else 0 + raise MemoryError( + f"{self.model.name} estimated CUDA peak " + f"{required / 1024**3:.2f} GiB for input_dim={feature_dim:,}, but " + f"only {free / 1024**3:.2f} GiB is free " + f"({total / 1024**3:.2f} GiB total). Rejecting this candidate " + "before allocating CUDA tensors." + ) + def _create_cv_splits(self): train_labels = [self.labels[i] for i in self.train_indices] train_labels_array = np.array(train_labels) @@ -273,6 +328,7 @@ def run(self, data): :param data: The aligned data used in the prediction process :return: the validation accuracy """ + self.validate_input_memory(data) self._reset_params() model = self.create_model() @@ -303,6 +359,7 @@ def run(self, data): float(np.mean(self.inference_time)) if self.inference_time else 0.0 ), "n_test_instances": len(self.test_indices) if self.test_indices else 0, + **self.training_work.as_dict(), } return [ self.train_scores.compute_averages(), @@ -398,6 +455,7 @@ def _measure_single_sample_latency( def _reset_params(self): self.inference_time = [] self.training_time = [] + self.training_work = TrainingWorkAccumulator() self.train_scores = PerformanceMeasure("train", self.performance_measures) self.val_scores = PerformanceMeasure("val", self.performance_measures) self.test_scores = PerformanceMeasure("test", self.performance_measures) @@ -415,6 +473,8 @@ def _run_fold(self, model, train_X, train_y, val_X, val_y, test_X, test_y): self.inference_time.append(test_end - test_start) self.val_scores.add_scores(val_score[0]) self.test_scores.add_scores(test_score[0]) + # No-op for the non-neural models that share this machinery. + self.training_work.add(model) @staticmethod def _gather_by_indices(values, indices): diff --git a/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py b/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py index fabe6f1e6fc..218feaf0d92 100644 --- a/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py +++ b/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py @@ -356,13 +356,9 @@ def _process_modality(self, modality, skip_remaining: int = 0, scheduler=None): ) start_time = time.perf_counter() - fitted_dags = [ - dag for dag in dags_with_pushdown if dag_requires_fitting(dag) - ] + fitted_dags = [dag for dag in dags_with_pushdown if dag_requires_fitting(dag)] stateless_dag_ids = { - dag.dag_id - for dag in dags_with_pushdown - if not dag_requires_fitting(dag) + dag.dag_id for dag in dags_with_pushdown if not dag_requires_fitting(dag) } stateless_task_dags = [ dag diff --git a/src/main/python/systemds/scuro/models/training_stats.py b/src/main/python/systemds/scuro/models/training_stats.py new file mode 100644 index 00000000000..80d8816bf4a --- /dev/null +++ b/src/main/python/systemds/scuro/models/training_stats.py @@ -0,0 +1,184 @@ +# ------------------------------------------------------------- +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# +# ------------------------------------------------------------- +import time + + +def _sync(device): + if getattr(device, "type", None) != "cuda": + return + try: + import torch + + if torch.cuda.is_available(): + torch.cuda.synchronize() + except Exception: + pass + + +class _Phase: + """Context manager accumulating elapsed seconds into one stats field.""" + + __slots__ = ("_stats", "_field", "_start") + + def __init__(self, stats, field): + self._stats = stats + self._field = field + self._start = 0.0 + + def __enter__(self): + _sync(self._stats.device) + self._start = time.perf_counter() + return self._stats + + def __exit__(self, *exc): + _sync(self._stats.device) + setattr( + self._stats, + self._field, + getattr(self._stats, self._field) + (time.perf_counter() - self._start), + ) + return False + + +class TrainingStats: + """Counters and phase timers for one model instance (i.e. one CV fold).""" + + FIELDS = ( + "fits", + "epochs_run", + "train_steps", + "train_samples", + "setup_s", + "train_loop_s", + "eval_calls", + "eval_batches", + "eval_samples", + "eval_s", + ) + + __slots__ = FIELDS + ("device", "_accounted") + + def __init__(self, device=None): + self.device = device + self._accounted = None + for field in self.FIELDS: + setattr(self, field, 0 if not field.endswith("_s") else 0.0) + + # -- phases ------------------------------------------------------------ + def setup(self): + """Model construction, scaling, tensor materialisation, loader build.""" + return _Phase(self, "setup_s") + + def training(self): + """The epoch/step loop only.""" + return _Phase(self, "train_loop_s") + + def evaluation(self): + """A forward-only pass: predict(), test(), or the train-set scoring.""" + self.eval_calls += 1 + return _Phase(self, "eval_s") + + # -- counters ---------------------------------------------------------- + def count_step(self, batch_samples=0): + self.train_steps += 1 + self.train_samples += int(batch_samples) + + def count_epoch(self): + self.epochs_run += 1 + + def count_eval_batch(self, batch_samples=0): + self.eval_batches += 1 + self.eval_samples += int(batch_samples) + + def count_fit(self): + self.fits += 1 + + def as_dict(self): + return {field: getattr(self, field) for field in self.FIELDS} + + def take_delta(self): + current = self.as_dict() + previous = self._accounted or {field: 0 for field in self.FIELDS} + self._accounted = current + return {field: current[field] - previous[field] for field in self.FIELDS} + + +class TrainingStatsMixin: + @property + def training_stats(self): + stats = getattr(self, "_training_stats", None) + if stats is None: + stats = TrainingStats(getattr(self, "device", None)) + self._training_stats = stats + + stats.device = getattr(self, "device", None) + return stats + + def reset_training_stats(self): + self._training_stats = TrainingStats(getattr(self, "device", None)) + return self._training_stats + + +class TrainingWorkAccumulator: + + def __init__(self): + self.totals = {field: 0 for field in TrainingStats.FIELDS} + self.folds = 0 + self.observed = False + + def add(self, model): + stats = getattr(model, "training_stats", None) + if stats is None or not hasattr(stats, "take_delta"): + return + values = stats.take_delta() + if not values.get("fits"): + return + self.observed = True + self.folds += 1 + for field, value in values.items(): + self.totals[field] += value + + def as_dict(self): + """Flat, JSON/CSV-friendly totals, prefixed so they cannot collide.""" + if not self.observed: + return {} + totals = self.totals + fixed_s = totals["setup_s"] + totals["eval_s"] + train_s = totals["train_loop_s"] + instrumented = fixed_s + train_s + return { + "work_folds": self.folds, + "work_fits": totals["fits"], + "work_epochs_run": totals["epochs_run"], + "work_train_steps": totals["train_steps"], + "work_train_samples": totals["train_samples"], + "work_setup_s": totals["setup_s"], + "work_train_loop_s": train_s, + "work_eval_s": totals["eval_s"], + "work_eval_calls": totals["eval_calls"], + "work_eval_batches": totals["eval_batches"], + "work_eval_samples": totals["eval_samples"], + "work_fixed_s": fixed_s, + "work_fixed_fraction": (fixed_s / instrumented) if instrumented else None, + "work_s_per_train_step": ( + train_s / totals["train_steps"] if totals["train_steps"] else None + ), + } diff --git a/src/main/python/systemds/scuro/representations/aggregate.py b/src/main/python/systemds/scuro/representations/aggregate.py index afbd6b346ea..366bb9f3ce9 100644 --- a/src/main/python/systemds/scuro/representations/aggregate.py +++ b/src/main/python/systemds/scuro/representations/aggregate.py @@ -20,7 +20,6 @@ # ------------------------------------------------------------- import numpy as np -from systemds.scuro.modality.type import ModalityType from systemds.scuro.representations import utils @@ -105,11 +104,6 @@ def execute(self, modality, aggregate_dim=(0,), squeeze_singleton=True): data.append([]) if isinstance(instance, np.ndarray) or isinstance(instance, list): if ( - modality.modality_type == ModalityType.IMAGE - or modality.modality_type == ModalityType.VIDEO - ) and instance.ndim > 2: - aggregated_data = instance.flatten() - elif ( squeeze_singleton and isinstance(instance, np.ndarray) and instance.ndim == 2 diff --git a/src/main/python/systemds/scuro/representations/fusion.py b/src/main/python/systemds/scuro/representations/fusion.py index 496fc2fba1f..c3c83230b75 100644 --- a/src/main/python/systemds/scuro/representations/fusion.py +++ b/src/main/python/systemds/scuro/representations/fusion.py @@ -30,7 +30,7 @@ ) from systemds.scuro.modality.transformed import TransformedModality -from systemds.scuro.modality.modality import Modality +from systemds.scuro.modality.modality import Modality, selective_copy_metadata from systemds.scuro.representations.representation import ( CONTAINER_ARRAY, FitRequirement, @@ -43,6 +43,8 @@ ) from systemds.scuro.utils.schema_helpers import get_shape +DEFAULT_MAX_ALIGNMENT_PAD_RATIO = None + class Fusion(Representation): def __init__(self, name, parameters=None): @@ -59,6 +61,7 @@ def __init__(self, name, parameters=None): self.preserves_leading_axis = False self.output_modality_type = ModalityType.EMBEDDING self.data_type = np.float32 + self.max_alignment_pad_ratio = DEFAULT_MAX_ALIGNMENT_PAD_RATIO @staticmethod def _as_stats_list(input_stats) -> List[RepresentationStats]: @@ -113,19 +116,47 @@ def _prepare_modalities(self, modalities: List[Modality]): and modality.has_metadata() and get_shape(modality.metadata) > 1 ): - agg_operator = AggregatedRepresentation() + agg_operator = AggregatedRepresentation( + target_dimensions=1, aggregate_leading=True + ) agg_modality = agg_operator.transform(modality) mods.append(agg_modality if agg_modality else modality) if self.needs_alignment: + mods = [self._detached_copy(modality) for modality in mods] for modality in mods: self._normalize_for_fusion(modality) max_len = self.get_max_embedding_size(mods) + self._check_alignment_padding(mods, max_len) for modality in mods: modality.pad(max_len=max_len) return mods + @staticmethod + def _detached_copy(modality: Modality) -> Modality: + detached = copy.copy(modality) + detached.metadata = selective_copy_metadata(modality.metadata) + return detached + + def _check_alignment_padding(self, modalities: List[Modality], max_len: int): + if not self.max_alignment_pad_ratio or max_len <= 0: + return + + for modality in modalities: + width = self._embedding_width(modality) + if width <= 0: + continue + if max_len > width * self.max_alignment_pad_ratio: + raise ValueError( + f"{self.name} would have to zero-pad modality " + f"{modality.modality_id} from {width:,} to {max_len:,} " + f"features to align it, exceeding the maximum padding ratio " + f"of {self.max_alignment_pad_ratio:g}x. Aggregate or reduce " + "the operands to comparable widths before fusing them " + "elementwise." + ) + @staticmethod def _validate_indexed_modalities(inputs): if not isinstance(inputs, (list, tuple)) or not inputs: @@ -216,6 +247,24 @@ def _normalize_for_fusion(modality: Modality): elif not modality.has_metadata() or len(modality.metadata) <= 1: modality.data = arr.reshape(1, -1).copy() + @staticmethod + def _embedding_width(m: Modality) -> int: + data = m.data + if isinstance(data, memoryview): + data = np.array(data) + arr = np.asarray(data) + if arr.dtype == object: + return 0 + if arr.ndim == 1: + if m.has_metadata() and len(m.metadata) == 1: + return arr.shape[0] + if not m.has_metadata() or len(m.metadata) <= 1: + return arr.shape[0] + return 0 + if arr.ndim >= 2: + return arr.shape[1] + return 0 + def get_max_embedding_size(self, modalities: List[Modality]): """ Computes the maximum embedding size from a given list of modalities @@ -223,25 +272,4 @@ def get_max_embedding_size(self, modalities: List[Modality]): :return: maximum embedding size """ - max_size = 0 - for m in modalities: - data = m.data - if isinstance(data, memoryview): - data = np.array(data) - arr = np.asarray(data) - if arr.dtype == object: - continue - if arr.ndim == 1: - if m.has_metadata() and len(m.metadata) == 1: - emb_size = arr.shape[0] - elif not m.has_metadata() or len(m.metadata) <= 1: - emb_size = arr.shape[0] - else: - continue - elif arr.ndim >= 2: - emb_size = arr.shape[1] - else: - continue - if emb_size > max_size: - max_size = emb_size - return max_size + return max((self._embedding_width(m) for m in modalities), default=0) diff --git a/src/main/python/systemds/scuro/representations/lstm.py b/src/main/python/systemds/scuro/representations/lstm.py index 346c8e992db..2086b199446 100644 --- a/src/main/python/systemds/scuro/representations/lstm.py +++ b/src/main/python/systemds/scuro/representations/lstm.py @@ -30,7 +30,11 @@ from systemds.scuro.modality.modality import Modality from systemds.scuro.representations.fusion import Fusion -from systemds.scuro.representations.representation import FitRequirement +from systemds.scuro.representations.representation import ( + FitRequirement, + RepresentationStats, + stats_bytes, +) from systemds.scuro.drsearch.operator_registry import register_fusion_operator @@ -47,6 +51,8 @@ def __init__( learning_rate=0.001, epochs=20, batch_size=32, + max_input_dim=16384, + max_parameter_bytes=1024**3, params=None, ): parameters = { @@ -67,6 +73,8 @@ def __init__( learning_rate = params.get("learning_rate", learning_rate) epochs = params.get("epochs", epochs) batch_size = params.get("batch_size", batch_size) + max_input_dim = params.get("max_input_dim", max_input_dim) + max_parameter_bytes = params.get("max_parameter_bytes", max_parameter_bytes) self.width = int(width) self.depth = int(depth) @@ -74,9 +82,14 @@ def __init__( self.learning_rate = float(learning_rate) self.epochs = int(epochs) self.batch_size = int(batch_size) + self.max_input_dim = None if max_input_dim is None else int(max_input_dim) + self.max_parameter_bytes = ( + None if max_parameter_bytes is None else int(max_parameter_bytes) + ) self.device = get_device() self.needs_training = True - self.needs_alignment = True + self.needs_alignment = False + self.preserves_leading_axis = True self.model = None self.input_dim = None self.num_classes = None @@ -95,35 +108,89 @@ def _set_random_seeds(self, seed=42): torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False - def _prepare_data(self, modalities: List[Modality]) -> np.ndarray: - processed_modalities = [] - - for modality in modalities: - try: - data = np.array(modality.data) - except: - max_len = -1 - for md in modality.metadata: - if max_len < md["data_layout"]["shape"][0]: - max_len = md["data_layout"]["shape"][0] - data = np.zeros((len(modality.data), max_len)) - for i, d in enumerate(modality.data): - data[i, : len(d)] = d - - modality.data = data + @staticmethod + def _shape_description(modality: Modality) -> str: + data = modality.data + if isinstance(data, np.ndarray): + shape = data.shape + else: + shapes = [np.asarray(sample).shape for sample in data[:3]] + shape = f"ragged samples {shapes}" + return ( + f"id={modality.modality_id}, type={modality.modality_type}, " + f"shape={shape}" + ) + + @staticmethod + def _as_sequence_array(modality: Modality) -> np.ndarray: + try: + data = np.asarray(modality.data) + except ValueError: + data = np.asarray(modality.data, dtype=object) + if data.dtype != object: if data.ndim == 1: - data = data.reshape(-1, 1, 1) - elif data.ndim == 2: - data = data.reshape(data.shape[0], 1, data.shape[1]) - elif data.ndim == 3: - pass - else: + return data.reshape(-1, 1, 1) + if data.ndim == 2: + return data.reshape(data.shape[0], 1, data.shape[1]) + if data.ndim == 3: + return data + raise ValueError( + "LSTM fusion accepts per-instance vectors or sequences with " + f"shape (instances, features) or (instances, sequence, features); " + f"got {LSTM._shape_description(modality)}. Pool spatial dimensions " + "or apply a feature representation before LSTM fusion." + ) + + samples = [] + feature_dim = None + for index, sample in enumerate(modality.data): + sample = np.asarray(sample) + if sample.ndim == 1: + sample = sample.reshape(1, -1) + elif sample.ndim != 2: + raise ValueError( + "Ragged LSTM inputs must contain vectors or (sequence, features) " + f"matrices; sample {index} has shape {sample.shape} in " + f"{LSTM._shape_description(modality)}." + ) + if feature_dim is None: + feature_dim = sample.shape[1] + elif sample.shape[1] != feature_dim: raise ValueError( - f"Unsupported data shape: {data.shape}. Expected 1D, 2D, or 3D arrays." + "Ragged LSTM sequences must have a fixed feature dimension; " + f"expected {feature_dim}, got {sample.shape[1]} for sample {index}." ) + samples.append(sample) + + if not samples: + raise ValueError("LSTM fusion requires at least one instance") + + max_seq_len = max(sample.shape[0] for sample in samples) + result = np.zeros((len(samples), max_seq_len, feature_dim), dtype=np.float32) + for index, sample in enumerate(samples): + result[index, : sample.shape[0]] = sample + return result - processed_modalities.append(data) + def _prepare_data(self, modalities: List[Modality]) -> np.ndarray: + if not modalities: + raise ValueError("LSTM fusion requires at least one modality") + + processed_modalities = [ + self._as_sequence_array(modality) for modality in modalities + ] + num_instances = processed_modalities[0].shape[0] + if any(data.shape[0] != num_instances for data in processed_modalities[1:]): + shapes = ", ".join(self._shape_description(m) for m in modalities) + raise ValueError(f"LSTM inputs have different instance counts: {shapes}") + + input_dim = sum(data.shape[2] for data in processed_modalities) + if self.max_input_dim is not None and input_dim > self.max_input_dim: + shapes = ", ".join(self._shape_description(m) for m in modalities) + raise ValueError( + f"LSTM input feature dimension {input_dim:,} exceeds " + f"the configured maximum {self.max_input_dim:,}. Inputs: {shapes}" + ) max_seq_len = max(mod.shape[1] for mod in processed_modalities) @@ -135,9 +202,70 @@ 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) + def _parameter_bytes(self, input_dim: int, num_classes: int) -> int: + directions = 2 + gates = 4 + parameter_count = 0 + layer_input_dim = input_dim + for _ in range(self.depth): + per_direction = ( + gates * self.width * layer_input_dim + + gates * self.width * self.width + + 2 * gates * self.width + ) + parameter_count += directions * per_direction + layer_input_dim = directions * self.width + parameter_count += directions * self.width * num_classes + num_classes + return parameter_count * np.dtype(np.float32).itemsize + + def _validate_model_size(self, input_dim: int, num_classes: int) -> None: + parameter_bytes = self._parameter_bytes(input_dim, num_classes) + if ( + self.max_parameter_bytes is not None + and parameter_bytes > self.max_parameter_bytes + ): + raise ValueError( + f"LSTM parameters require approximately " + f"{parameter_bytes / 1024**3:.2f} GiB, exceeding the configured " + f"maximum {self.max_parameter_bytes / 1024**3:.2f} GiB " + f"(input_dim={input_dim:,}, width={self.width}, depth={self.depth})." + ) + + def get_output_stats(self, input_stats) -> RepresentationStats: + stats_list = self._as_stats_list(input_stats) + num_instances = max((stats.num_instances for stats in stats_list), default=0) + return RepresentationStats( + num_instances, + (2 * self.width,), + output_shape_is_known=all( + stats.output_shape_is_known for stats in stats_list + ), + dtype=np.float32, + ) + + def estimate_peak_memory_bytes(self, input_stats) -> dict: + stats_list = self._as_stats_list(input_stats) + input_dim = 0 + max_seq_len = 1 + for stats in stats_list: + shape = tuple(stats.output_shape) + if len(shape) == 1: + input_dim += shape[0] + elif len(shape) == 2: + max_seq_len = max(max_seq_len, shape[0]) + input_dim += shape[1] + else: + return {"cpu_peak_bytes": 0, "gpu_peak_bytes": 2**63 - 1} + parameter_bytes = self._parameter_bytes(input_dim, 2) + input_bytes = sum(stats_bytes(stats) for stats in stats_list) + batch_activations = self.batch_size * max_seq_len * 2 * self.width * 4 + return { + "cpu_peak_bytes": input_bytes + parameter_bytes, + "gpu_peak_bytes": 4 * parameter_bytes + 3 * batch_activations, + } + def _build_model(self, input_dim: int, num_classes: int) -> nn.Module: class LSTMClassifier(nn.Module): @@ -189,6 +317,7 @@ def execute(self, modalities: List[Modality], labels: np.ndarray = None): self.num_classes = len(np.unique(y)) self.input_dim = X.shape[2] + self._validate_model_size(self.input_dim, self.num_classes) self.model = self._build_model(self.input_dim, self.num_classes) self.model = self.model.to(self.device) @@ -254,6 +383,11 @@ def apply_representation(self, modalities: List[Modality]) -> np.ndarray: raise ValueError("Model must be trained before applying representation") X = self._prepare_data(modalities) + if X.shape[2] != self.input_dim: + raise ValueError( + f"LSTM was fitted with input_dim={self.input_dim:,}, but received " + f"input_dim={X.shape[2]:,}." + ) self.model = self.model.to(self.device) @@ -279,7 +413,11 @@ def set_model_state(self, state: Dict[str, Any]): self.input_dim = state["input_dim"] self.num_classes = state["num_classes"] self.is_multilabel = state.get("is_multilabel", False) + self.width = int(state.get("width", self.width)) + self.depth = int(state.get("depth", self.depth)) + self.dropout_rate = float(state.get("dropout_rate", self.dropout_rate)) + self._validate_model_size(self.input_dim, self.num_classes) self.model = self._build_model(self.input_dim, self.num_classes) self.model.load_state_dict(state["state_dict"]) self.is_trained = True 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 b972d0a12d5..162892a7faa 100644 --- a/src/main/python/systemds/scuro/representations/multimodal_attention_fusion.py +++ b/src/main/python/systemds/scuro/representations/multimodal_attention_fusion.py @@ -29,7 +29,11 @@ from systemds.scuro.drsearch.operator_registry import register_fusion_operator from systemds.scuro.modality.modality import Modality from systemds.scuro.representations.fusion import Fusion -from systemds.scuro.representations.representation import FitRequirement +from systemds.scuro.representations.representation import ( + FitRequirement, + RepresentationStats, + stats_bytes, +) from systemds.scuro.utils.static_variables import get_device @@ -45,6 +49,8 @@ def __init__( batch_size=32, num_epochs=20, learning_rate=0.001, + max_input_dim=16384, + max_parameter_bytes=1024**3, params=None, ): parameters = { @@ -70,6 +76,13 @@ def __init__( self.batch_size = int(params.get("batch_size", self.batch_size)) self.num_epochs = int(params.get("num_epochs", self.num_epochs)) self.learning_rate = float(params.get("learning_rate", self.learning_rate)) + max_input_dim = params.get("max_input_dim", max_input_dim) + max_parameter_bytes = params.get("max_parameter_bytes", max_parameter_bytes) + + self.max_input_dim = None if max_input_dim is None else int(max_input_dim) + self.max_parameter_bytes = ( + None if max_parameter_bytes is None else int(max_parameter_bytes) + ) self.needs_training = True self.needs_alignment = True @@ -127,6 +140,101 @@ def _prepare_data(self, modalities: List[Modality]): return inputs, input_dimensions, max_sequence_length + def _parameter_bytes( + self, + input_dimensions: Dict[str, int], + max_sequence_length: int, + num_classes: int, + ) -> int: + hidden = self.hidden_dim + projections = sum((dim + 1) * hidden for dim in input_dimensions.values()) + positional = max_sequence_length * hidden + attention = 2 * (4 * hidden * hidden + 4 * hidden) + modality_attention = hidden + 1 + final_projection = hidden * hidden + hidden + head_width = hidden // 2 + classification = hidden * head_width + head_width + classification += head_width * num_classes + num_classes + return int( + ( + projections + + positional + + attention + + modality_attention + + final_projection + + classification + ) + * np.dtype(np.float32).itemsize + ) + + def _validate_model_size( + self, + input_dimensions: Dict[str, int], + max_sequence_length: int, + num_classes: int, + ) -> None: + total_input_dim = sum(input_dimensions.values()) + if self.max_input_dim is not None and total_input_dim > self.max_input_dim: + raise ValueError( + f"AttentionFusion input feature dimension {total_input_dim:,} exceeds " + f"the configured maximum {self.max_input_dim:,}; per-modality " + f"dimensions: {input_dimensions}. Apply aggregation or dimensionality " + "reduction before attention fusion." + ) + + parameter_bytes = self._parameter_bytes( + input_dimensions, max_sequence_length, num_classes + ) + if ( + self.max_parameter_bytes is not None + and parameter_bytes > self.max_parameter_bytes + ): + raise ValueError( + "AttentionFusion parameters require approximately " + f"{parameter_bytes / 1024**3:.2f} GiB, exceeding the configured " + f"maximum {self.max_parameter_bytes / 1024**3:.2f} GiB." + ) + + def get_output_stats(self, input_stats) -> RepresentationStats: + stats_list = self._as_stats_list(input_stats) + num_instances = max((stats.num_instances for stats in stats_list), default=0) + return RepresentationStats( + num_instances, + (self.hidden_dim,), + output_shape_is_known=all( + stats.output_shape_is_known for stats in stats_list + ), + dtype=np.float32, + ) + + def estimate_peak_memory_bytes(self, input_stats) -> dict: + stats_list = self._as_stats_list(input_stats) + dimensions = {} + max_sequence_length = 1 + for index, stats in enumerate(stats_list): + shape = tuple(stats.output_shape) + if len(shape) == 1: + dimensions[f"modality_{index}"] = shape[0] + elif len(shape) == 2: + max_sequence_length = max(max_sequence_length, shape[0]) + dimensions[f"modality_{index}"] = shape[1] + else: + return {"cpu_peak_bytes": 0, "gpu_peak_bytes": 2**63 - 1} + + parameter_bytes = self._parameter_bytes(dimensions, max_sequence_length, 2) + input_bytes = sum(stats_bytes(stats) for stats in stats_list) + batch_activations = ( + self.batch_size + * max_sequence_length + * max(1, len(stats_list)) + * self.hidden_dim + * np.dtype(np.float32).itemsize + ) + return { + "cpu_peak_bytes": input_bytes + parameter_bytes, + "gpu_peak_bytes": 4 * parameter_bytes + 4 * batch_activations, + } + def execute(self, modalities: List[Modality], labels: np.ndarray = None): if labels is None: raise ValueError("Attention fusion requires labels for training") @@ -145,6 +253,9 @@ def execute(self, modalities: List[Modality], labels: np.ndarray = None): self.input_dim = input_dimensions self.max_sequence_length = max_sequence_length + self._validate_model_size( + self.input_dim, self.max_sequence_length, self.num_classes + ) self.encoder = MultiModalAttentionFusion( self.input_dim, @@ -309,6 +420,9 @@ def set_model_state(self, state: Dict[str, Any]): self.max_sequence_length = state["max_sequence_length"] self.num_classes = state["num_classes"] self.is_multilabel = state.get("is_multilabel", False) + self._validate_model_size( + self.input_dim, self.max_sequence_length, self.num_classes + ) self.encoder = MultiModalAttentionFusion( self.input_dim, From 20ceb18c2458cfd7d5d6868a60efc1208025f610 Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Tue, 8 Sep 2026 21:15:03 +0200 Subject: [PATCH 07/15] improve metadata --- .../systemds/scuro/drsearch/worker_pool.py | 21 ++++++++----- .../systemds/scuro/models/training_stats.py | 2 +- .../representations/physiological_window.py | 30 +++++++++++------- .../systemds/scuro/utils/schema_helpers.py | 31 ++++++++++++++++++- 4 files changed, 63 insertions(+), 21 deletions(-) diff --git a/src/main/python/systemds/scuro/drsearch/worker_pool.py b/src/main/python/systemds/scuro/drsearch/worker_pool.py index 752efdb6eed..4e7a112dd5a 100644 --- a/src/main/python/systemds/scuro/drsearch/worker_pool.py +++ b/src/main/python/systemds/scuro/drsearch/worker_pool.py @@ -41,7 +41,7 @@ def _resolve_thread_count(num_threads: int) -> int: - explicit = os.environ.get("OMP_NUM_THREADS") + explicit = os.environ.get("SCURO_THREADS_PER_WORKER") if explicit: try: num_threads = int(explicit) @@ -199,8 +199,11 @@ def _spawn_worker(self, physical_gpu_id: Optional[int]) -> None: set_thread_env_before_spawn(self._threads_per_worker) job_q = self._ctx.Queue() previous_visible = os.environ.get("CUDA_VISIBLE_DEVICES") + visible_id = physical_gpu_id + if physical_gpu_id is not None and previous_visible: + visible_id = previous_visible.split(",")[physical_gpu_id].strip() os.environ["CUDA_VISIBLE_DEVICES"] = ( - "" if physical_gpu_id is None else str(physical_gpu_id) + "" if visible_id is None else str(visible_id) ) p = self._ctx.Process( target=_worker_main, @@ -209,15 +212,17 @@ def _spawn_worker(self, physical_gpu_id: Optional[int]) -> None: self._result_q, self._dispatch, self._threads_per_worker, - physical_gpu_id, + visible_id, ), daemon=True, ) - p.start() - if previous_visible is None: - os.environ.pop("CUDA_VISIBLE_DEVICES", None) - else: - os.environ["CUDA_VISIBLE_DEVICES"] = previous_visible + try: + p.start() + finally: + if previous_visible is None: + os.environ.pop("CUDA_VISIBLE_DEVICES", None) + else: + os.environ["CUDA_VISIBLE_DEVICES"] = previous_visible self._workers[p.pid] = { "process": p, "job_q": job_q, diff --git a/src/main/python/systemds/scuro/models/training_stats.py b/src/main/python/systemds/scuro/models/training_stats.py index 80d8816bf4a..bee9d20f7a2 100644 --- a/src/main/python/systemds/scuro/models/training_stats.py +++ b/src/main/python/systemds/scuro/models/training_stats.py @@ -128,7 +128,7 @@ def training_stats(self): if stats is None: stats = TrainingStats(getattr(self, "device", None)) self._training_stats = stats - + stats.device = getattr(self, "device", None) return stats diff --git a/src/main/python/systemds/scuro/representations/physiological_window.py b/src/main/python/systemds/scuro/representations/physiological_window.py index 338cfdd2201..c6d32e3fe2b 100644 --- a/src/main/python/systemds/scuro/representations/physiological_window.py +++ b/src/main/python/systemds/scuro/representations/physiological_window.py @@ -32,6 +32,22 @@ ) +def _event_peak_indices(normalized, threshold, min_distance): + x = np.asarray(normalized).reshape(len(normalized)) + x = x.astype(np.result_type(x.dtype.type(0), threshold), copy=False) + candidates = ( + np.flatnonzero((x[1:-1] > threshold) & (x[1:-1] > x[:-2]) & (x[1:-1] > x[2:])) + + 1 + ) + peaks = [] + index = 0 + while index < len(candidates): + peak = int(candidates[index]) + peaks.append(peak) + index = int(np.searchsorted(candidates, peak + min_distance)) + return peaks + + def _estimate_windowed_output_stats( window_obj, input_stats, estimated_num_windows, representative_window_length ): @@ -229,17 +245,9 @@ def execute(self, modality): for signal in modality.data: normalized = (signal - np.mean(signal)) / (np.std(signal) + 1e-6) - peaks = [] - last_peak = -self.min_distance - for i in range(1, len(normalized) - 1): - if ( - normalized[i] > self.event_threshold - and normalized[i] > normalized[i - 1] - and normalized[i] > normalized[i + 1] - and i - last_peak >= self.min_distance - ): - peaks.append(i) - last_peak = i + peaks = _event_peak_indices( + normalized, self.event_threshold, self.min_distance + ) windows = [signal[peaks[i] : peaks[i + 1]] for i in range(len(peaks) - 1)] diff --git a/src/main/python/systemds/scuro/utils/schema_helpers.py b/src/main/python/systemds/scuro/utils/schema_helpers.py index 929e9c7f4f3..a122d8eae08 100644 --- a/src/main/python/systemds/scuro/utils/schema_helpers.py +++ b/src/main/python/systemds/scuro/utils/schema_helpers.py @@ -21,7 +21,7 @@ import numpy as np -def create_timestamps(frequency, sample_length, start_datetime=None): +def _materialize_timestamps(frequency, sample_length, start_datetime=None): start_time = ( start_datetime if start_datetime is not None @@ -36,6 +36,35 @@ def create_timestamps(frequency, sample_length, start_datetime=None): return timestamps.astype(np.int64) +class DerivedTimestamps(np.ndarray): + + def __new__(cls, frequency, sample_length, start_datetime=None): + obj = _materialize_timestamps(frequency, sample_length, start_datetime).view( + cls + ) + obj._generator = (frequency, sample_length, start_datetime) + return obj + + def __array_finalize__(self, obj): + self._generator = None + + def __reduce__(self): + generator = getattr(self, "_generator", None) + if generator is not None: + return (DerivedTimestamps, generator) + return super().__reduce__() + + def __deepcopy__(self, memo): + generator = getattr(self, "_generator", None) + if generator is not None: + return DerivedTimestamps(*generator) + return super().__deepcopy__(memo) + + +def create_timestamps(frequency, sample_length, start_datetime=None): + return DerivedTimestamps(frequency, sample_length, start_datetime) + + def calculate_new_frequency(new_length, old_length, old_frequency): duration = old_length / old_frequency new_frequency = new_length / duration From 93b7f82d70e287f71d9bf1b499f4df18a64c3072 Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Wed, 9 Sep 2026 06:58:10 +0200 Subject: [PATCH 08/15] set default threads per worker --- .../systemds/scuro/drsearch/multimodal_ga_optimizer.py | 9 ++++++--- .../python/systemds/scuro/drsearch/node_executor.py | 9 ++++++--- src/main/python/systemds/scuro/drsearch/worker_pool.py | 10 ++++++++++ 3 files changed, 22 insertions(+), 6 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 ba170bccd01..0c26c9ebcb4 100644 --- a/src/main/python/systemds/scuro/drsearch/multimodal_ga_optimizer.py +++ b/src/main/python/systemds/scuro/drsearch/multimodal_ga_optimizer.py @@ -53,7 +53,11 @@ SplitAwareDAGEvaluator, dag_requires_fitting, ) -from systemds.scuro.drsearch.worker_pool import PersistentWorkerPool, create_mp_context +from systemds.scuro.drsearch.worker_pool import ( + PersistentWorkerPool, + create_mp_context, + default_threads_per_worker, +) from systemds.scuro.representations.aggregated_representation import ( AggregatedRepresentation, ) @@ -410,13 +414,12 @@ def __init__( self.elite_size = max(0, min(elite_size, self.population_size - 1)) self.max_workers = max(1, max_workers) self.batch_size = max(1, batch_size or self.max_workers) - cpu_count = os.cpu_count() or 1 self.threads_per_worker = max( 1, ( threads_per_worker if threads_per_worker is not None - else cpu_count // self.max_workers + else default_threads_per_worker() ), ) self.early_stopping_patience = early_stopping_patience diff --git a/src/main/python/systemds/scuro/drsearch/node_executor.py b/src/main/python/systemds/scuro/drsearch/node_executor.py index a707db1cd24..32795b8084a 100644 --- a/src/main/python/systemds/scuro/drsearch/node_executor.py +++ b/src/main/python/systemds/scuro/drsearch/node_executor.py @@ -39,7 +39,11 @@ RepresentationNode, ) from systemds.scuro.drsearch.task import PerformanceMeasure -from systemds.scuro.drsearch.worker_pool import PersistentWorkerPool, create_mp_context +from systemds.scuro.drsearch.worker_pool import ( + PersistentWorkerPool, + create_mp_context, + default_threads_per_worker, +) from systemds.scuro.representations.aggregated_representation import ( AggregatedRepresentation, ) @@ -487,8 +491,7 @@ def __init__( self._owns_pool = worker_pool is None if worker_pool is None: - cpu_count = os.cpu_count() or 1 - threads_per_worker = max(1, cpu_count // max(1, self.max_num_workers)) + threads_per_worker = default_threads_per_worker() worker_pool = PersistentWorkerPool( self.max_num_workers, _WORKER_DISPATCH, diff --git a/src/main/python/systemds/scuro/drsearch/worker_pool.py b/src/main/python/systemds/scuro/drsearch/worker_pool.py index 4e7a112dd5a..9969b67e626 100644 --- a/src/main/python/systemds/scuro/drsearch/worker_pool.py +++ b/src/main/python/systemds/scuro/drsearch/worker_pool.py @@ -50,6 +50,10 @@ def _resolve_thread_count(num_threads: int) -> int: return max(1, int(num_threads)) +def default_threads_per_worker() -> int: + return _resolve_thread_count(1) + + def set_thread_env_before_spawn(num_threads: int) -> None: num_threads = _resolve_thread_count(num_threads) for var in _THREAD_ENV_VARS: @@ -196,6 +200,7 @@ def gpu_worker_devices(self) -> List[int]: ) def _spawn_worker(self, physical_gpu_id: Optional[int]) -> None: + previous_threads = {var: os.environ.get(var) for var in _THREAD_ENV_VARS} set_thread_env_before_spawn(self._threads_per_worker) job_q = self._ctx.Queue() previous_visible = os.environ.get("CUDA_VISIBLE_DEVICES") @@ -223,6 +228,11 @@ def _spawn_worker(self, physical_gpu_id: Optional[int]) -> None: os.environ.pop("CUDA_VISIBLE_DEVICES", None) else: os.environ["CUDA_VISIBLE_DEVICES"] = previous_visible + for var, value in previous_threads.items(): + if value is None: + os.environ.pop(var, None) + else: + os.environ[var] = value self._workers[p.pid] = { "process": p, "job_q": job_q, From fce8e0a4419e36de0f554c9a019cf227dbfb0b01 Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Wed, 9 Sep 2026 07:07:31 +0200 Subject: [PATCH 09/15] add split to w2v --- .../scuro/representations/word2vec.py | 62 +++++++++++++++---- 1 file changed, 49 insertions(+), 13 deletions(-) diff --git a/src/main/python/systemds/scuro/representations/word2vec.py b/src/main/python/systemds/scuro/representations/word2vec.py index 0954a17e290..16cbd414fbf 100644 --- a/src/main/python/systemds/scuro/representations/word2vec.py +++ b/src/main/python/systemds/scuro/representations/word2vec.py @@ -20,8 +20,13 @@ # ------------------------------------------------------------- import numpy as np from systemds.scuro.dataloader.text_loader import TextStats +from systemds.scuro.modality.indexed import IndexedModality from systemds.scuro.modality.transformed import TransformedModality -from systemds.scuro.representations.representation import RepresentationStats +from systemds.scuro.representations.representation import ( + FitRequirement, + RepresentationFitState, + RepresentationStats, +) from systemds.scuro.representations.unimodal import UnimodalRepresentation from systemds.scuro.representations.utils import save_embeddings from gensim.models import Word2Vec @@ -29,7 +34,6 @@ from systemds.scuro.modality.type import ModalityType from systemds.scuro.drsearch.operator_registry import register_representation -import nltk def get_embedding(sentence, model): @@ -43,6 +47,8 @@ def get_embedding(sentence, model): @register_representation(ModalityType.TEXT) class W2V(UnimodalRepresentation): + fit_requirement = FitRequirement.UNSUPERVISED + def __init__(self, vector_size=128, min_count=1, output_file=None, params=None): parameters = { "vector_size": [64, 128, 256, 512, 1024], @@ -79,21 +85,51 @@ def estimate_peak_memory_bytes(self, input_stats: TextStats) -> dict: model_estimate = 100 * 1024 * 1024 return {"cpu_peak_bytes": model_estimate + output_bytes, "gpu_peak_bytes": 0} - def transform(self, modality, params=None): - transformed_modality = TransformedModality(modality, self) - t = [list(tokenize(s.lower())) for s in modality.data] + def fit(self, inputs, labels=None) -> RepresentationFitState: + if not isinstance(inputs, IndexedModality): + raise TypeError("W2V.fit expects an IndexedModality") + model = Word2Vec( - sentences=t, + sentences=[list(tokenize(s.lower())) for s in inputs.modality.data], vector_size=self.vector_size, min_count=self.min_count, ) - embeddings = [] - for sentences in modality.data: - tokens = list(tokenize(sentences.lower())) - embeddings.append(np.array(get_embedding(tokens, model)).astype(np.float32)) + return RepresentationFitState( + self.fit_requirement, + fit_instance_ids=inputs.instance_ids, + payload=model, + ) + + def transform_fitted(self, inputs, fit_state, aggregation=None): + if not isinstance(inputs, IndexedModality): + raise TypeError("W2V.transform_fitted expects an IndexedModality") + if ( + not isinstance(fit_state, RepresentationFitState) + or fit_state.fit_requirement is not self.fit_requirement + or type(fit_state.payload) is not Word2Vec + ): + raise ValueError("W2V requires a compatible fitted state") + + model = fit_state.payload + transformed_modality = TransformedModality(inputs.modality, self) + embeddings = np.asarray( + [ + get_embedding(list(tokenize(sentence.lower())), model) + for sentence in inputs.modality.data + ], + dtype=self.data_type, + ).reshape(-1, model.vector_size) if self.output_file is not None: - save_embeddings(np.array(embeddings), self.output_file) + save_embeddings(embeddings, self.output_file) transformed_modality.data_type = self.data_type - transformed_modality.data = np.array(embeddings) - return transformed_modality + transformed_modality.data = embeddings + if aggregation is not None: + transformed_modality.data = aggregation.execute(transformed_modality.data) + return IndexedModality(transformed_modality, inputs.instance_ids) + + def transform(self, modality, params=None, aggregation=None): + inputs = IndexedModality(modality, range(len(modality.data))) + return self.transform_fitted( + inputs, self.fit(inputs), aggregation=aggregation + ).modality From e84f2a5db7ae6731bacaa208dd2bbb716116f882 Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Thu, 10 Sep 2026 10:18:15 +0200 Subject: [PATCH 10/15] handle gc --- .../scuro/drsearch/modality_shared_memory.py | 17 +------ .../systemds/scuro/drsearch/node_executor.py | 28 +++++++---- .../scuro/modality/unimodal_modality.py | 47 ++++++++++++++----- .../systemds/scuro/representations/wav2vec.py | 9 +++- 4 files changed, 63 insertions(+), 38 deletions(-) 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 57a5048cda3..b7472a35a6d 100644 --- a/src/main/python/systemds/scuro/drsearch/modality_shared_memory.py +++ b/src/main/python/systemds/scuro/drsearch/modality_shared_memory.py @@ -20,18 +20,11 @@ # ------------------------------------------------------------- from typing import Any, List, Tuple import numpy as np -from multiprocessing import shared_memory, resource_tracker +from multiprocessing import shared_memory SHARED_MEMORY_MIN_BYTES = 1 * 1024 * 1024 -def _untrack(shm: shared_memory.SharedMemory) -> None: - try: - resource_tracker.unregister(shm._name, "shared_memory") - except Exception: - pass - - def unlink_shm(name: str) -> None: try: shm = shared_memory.SharedMemory(name=name) @@ -55,7 +48,6 @@ def __init__( def _ensure_attached(self): if self._shm is None: self._shm = shared_memory.SharedMemory(name=self.shm_name) - _untrack(self._shm) def __len__(self): return len(self.offsets) @@ -110,7 +102,6 @@ def __init__( def _ensure_attached(self): if self._shm is None: self._shm = shared_memory.SharedMemory(name=self.shm_name) - _untrack(self._shm) self._buffer = np.ndarray( (self.total_elems,), dtype=self._dtype, buffer=self._shm.buf ) @@ -166,7 +157,6 @@ def __init__(self, shm_name: str, dtype_str: str, shape: tuple): def _ensure_attached(self): if self._shm is None: self._shm = shared_memory.SharedMemory(name=self.shm_name) - _untrack(self._shm) self._arr = np.ndarray(self.shape, dtype=self._dtype, buffer=self._shm.buf) self._arr.setflags(write=False) @@ -236,7 +226,6 @@ def __init__( def _ensure_attached(self): if self._shm is None: self._shm = shared_memory.SharedMemory(name=self.shm_name) - _untrack(self._shm) self._buffer = np.ndarray( (self.total_elems,), dtype=self._dtype, buffer=self._shm.buf ) @@ -358,7 +347,6 @@ def add_shared_memory_candidate(data: Any, resident_bytes: int = 0) -> bool: resident_bytes, max(2 * 1024 * 1024, len(offsets) * 64) ) shm.close() - _untrack(shm) return data, shm.name, data_nbytes, resident_bytes elif _is_shared_ndarray_candidate(data): arr = data @@ -373,7 +361,6 @@ def add_shared_memory_candidate(data: Any, resident_bytes: int = 0) -> bool: resident_bytes = min(resident_bytes, 2 * 1024 * 1024) shm.close() - _untrack(shm) return data, shm.name, data_nbytes, resident_bytes elif _is_nested_shared_memory_candidate(data): leaves: List[np.ndarray] = [] @@ -409,7 +396,6 @@ def add_shared_memory_candidate(data: Any, resident_bytes: int = 0) -> bool: resident_bytes, max(2 * 1024 * 1024, len(offsets) * 64) ) shm.close() - _untrack(shm) return data, shm.name, data_nbytes, resident_bytes elif _is_string_list_shared_memory_candidate(data): encoded = [s.encode("utf-8") for s in data] @@ -429,7 +415,6 @@ def add_shared_memory_candidate(data: Any, resident_bytes: int = 0) -> bool: resident_bytes, max(2 * 1024 * 1024, len(str_offsets) * 32) ) shm.close() - _untrack(shm) return data, shm.name, data_nbytes, resident_bytes return None, None, 0, resident_bytes diff --git a/src/main/python/systemds/scuro/drsearch/node_executor.py b/src/main/python/systemds/scuro/drsearch/node_executor.py index 32795b8084a..d42a096d77f 100644 --- a/src/main/python/systemds/scuro/drsearch/node_executor.py +++ b/src/main/python/systemds/scuro/drsearch/node_executor.py @@ -93,7 +93,7 @@ def _place_on_device(obj: Any, gpu_id: Optional[int]) -> None: pass -def _run_gpu_op(fn, gpu_id: Optional[int]): +def _run_gpu_op(fn, gpu_id: Optional[int], cleanup_after: bool = True): if gpu_id is None or not torch.cuda.is_available(): return fn() try: @@ -106,7 +106,8 @@ def _run_gpu_op(fn, gpu_id: Optional[int]): return fn() raise finally: - cleanup_gpu(gpu_id) + if cleanup_after: + cleanup_gpu(gpu_id) _WORKER_OP_CACHE: Dict[str, Any] = {} @@ -207,7 +208,7 @@ def _run_node_op(): if DEBUG: input_resident = sum(estimate_modality_bytes(m) for m in input_mods) result, measurement = measure_memory_during( - lambda: _run_gpu_op(_run_node_op, gpu_id), + lambda: _run_gpu_op(_run_node_op, gpu_id, cleanup_after=False), input_resident_bytes=input_resident, sample_s=0.01, ) @@ -215,7 +216,7 @@ def _run_node_op(): torch.cuda.max_memory_allocated(device) if gpu_id is not None else 0 ) else: - result = _run_gpu_op(_run_node_op, gpu_id) + result = _run_gpu_op(_run_node_op, gpu_id, cleanup_after=False) shm_name, resident_bytes, shm_bytes, actual_stats = _offload_to_shared_memory( result @@ -487,6 +488,7 @@ def __init__( self._job_retained_shm: Dict[int, List[str]] = {} self._leaf_shm_names: List[str] = [] self._leaf_shm_owners: Dict[str, Any] = {} + self._leaf_original_data: Dict[str, Any] = {} self._task_results: Dict[str, ResultEntry] = {} self._owns_pool = worker_pool is None @@ -551,6 +553,13 @@ def _load_leaf_modalities(self) -> None: if self._loads_in_chunks(modality): continue 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 attempts = 0 while True: @@ -585,7 +594,7 @@ def _cleanup_leaf_shared_memory(self) -> None: close = getattr(modality._data, "close", None) if callable(close): close() - modality._data = None + modality._data = self._leaf_original_data.pop(shm_name, None) unlink_shm(shm_name) self._leaf_shm_names = [] self._leaf_shm_owners = {} @@ -652,6 +661,8 @@ def _submit_leaf_batch( 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() def _gpu_id(entry): @@ -660,6 +671,8 @@ def _gpu_id(entry): ready.sort(key=lambda entry: _gpu_id(entry) is None) for entry in ready: + if not self._pool.has_idle_worker: + break gpu_id = _gpu_id(entry) allow_gpu_worker_for_cpu = gpu_id is None if not self._pool.has_idle_worker_for( @@ -1035,17 +1048,14 @@ def _is_task_node(node: RepresentationNode) -> bool: def run(self) -> Dict[str, Any]: self._task_results = {} - self._load_leaf_modalities() - try: - self._fill_pipeline() + self._load_leaf_modalities() while self._job_units or not self.scheduler.is_finished(): self._fill_pipeline() if not self._job_units: continue jr = self._pool.wait() self._process_result(jr) - self._fill_pipeline() finally: if self._owns_pool: self._pool.shutdown() diff --git a/src/main/python/systemds/scuro/modality/unimodal_modality.py b/src/main/python/systemds/scuro/modality/unimodal_modality.py index c23faa06e09..e91a30f7891 100644 --- a/src/main/python/systemds/scuro/modality/unimodal_modality.py +++ b/src/main/python/systemds/scuro/modality/unimodal_modality.py @@ -20,7 +20,8 @@ # ------------------------------------------------------------- import copy from concurrent.futures import ThreadPoolExecutor, as_completed -import gc +from contextlib import nullcontext +from functools import partial import time import numpy as np from systemds.scuro.dataloader.base_loader import BaseLoader @@ -222,12 +223,20 @@ def apply_representations( start = ( time.time() ) # TODO: should be repalced in unimodal_representation.transform - with ThreadPoolExecutor( - max_workers=len(representations) if parallel else 1 + with ( + ThreadPoolExecutor(max_workers=len(representations)) + if parallel and len(representations) > 1 + else nullcontext() ) as executor: time_s = time.time() - for _ in self.iter_raw_data_chunks(reset=True): + chunks = ( + (None,) + if not self.data_loader.is_chunked and self.has_data() + else self.iter_raw_data_chunks(reset=True) + ) + for _ in chunks: representations_futures = {} + serial_transforms = [] for ( representation_key, representation, @@ -235,14 +244,31 @@ def apply_representations( ) in representation_specs: if representation_key in failed_representations: continue - future = executor.submit( - representation.transform, self, rep_aggregation + if executor is None: + serial_transforms.append( + ( + representation_key, + partial( + representation.transform, self, rep_aggregation + ), + ) + ) + else: + future = executor.submit( + representation.transform, self, rep_aggregation + ) + representations_futures[future] = representation_key + completed = ( + ( + (representations_futures[future], future.result) + for future in as_completed(representations_futures) ) - representations_futures[future] = representation_key - for future in as_completed(representations_futures.keys()): - representation_key = representations_futures.get(future) + if executor is not None + else serial_transforms + ) + for representation_key, transform in completed: try: - transformed_chunk = future.result() + transformed_chunk = transform() except Exception as e: failed_representations[representation_key] = e continue @@ -281,7 +307,6 @@ def apply_representations( transformed_modalities_per_representation[ representation_key ].self_contained = representation.self_contained - gc.collect() self.failed_representations = failed_representations return transformed_modalities_per_representation diff --git a/src/main/python/systemds/scuro/representations/wav2vec.py b/src/main/python/systemds/scuro/representations/wav2vec.py index d50465ca5c0..bee23af4dce 100644 --- a/src/main/python/systemds/scuro/representations/wav2vec.py +++ b/src/main/python/systemds/scuro/representations/wav2vec.py @@ -177,7 +177,9 @@ def estimate_peak_memory_bytes(self, input_stats) -> dict: frames = max(1, int(np.ceil(signal_len / stride))) model_resident = 420 * 1024 * 1024 # ~420 MB - activation_bytes = int(frames * hidden * 4 * 24) + 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) io_temp = int(signal_len * 4 * 4) + 16 * 1024 * 1024 output_bytes = n * 512 * 4 @@ -186,7 +188,10 @@ def estimate_peak_memory_bytes(self, input_stats) -> dict: (model_resident + activation_bytes + io_temp + output_bytes) * 1.25 ) - gpu_peak = 0 + gpu_peak = int( + (model_resident + activation_bytes + attention_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} From 676b9a6b3105e6b5ce40848416f2449440f3dcc2 Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Sat, 12 Sep 2026 21:12:26 +0200 Subject: [PATCH 11/15] simplify init --- src/main/python/systemds/scuro/__init__.py | 252 +++++++----------- .../systemds/scuro/dataloader/audio_loader.py | 29 +- .../systemds/scuro/dataloader/base_loader.py | 47 ++++ .../systemds/scuro/dataloader/image_loader.py | 20 +- .../systemds/scuro/dataloader/json_loader.py | 7 +- .../scuro/dataloader/tabular_loader.py | 13 +- .../scuro/dataloader/timeseries_loader.py | 19 +- .../systemds/scuro/dataloader/video_loader.py | 7 +- .../systemds/scuro/drsearch/node_executor.py | 6 +- .../scuro/drsearch/operator_registry.py | 75 +++++- .../python/systemds/scuro/drsearch/ranking.py | 86 +++--- .../python/systemds/scuro/drsearch/task.py | 156 +++++++++-- .../scuro/drsearch/unimodal_optimizer.py | 86 +++++- .../systemds/scuro/drsearch/worker_pool.py | 37 ++- .../python/systemds/scuro/modality/type.py | 3 +- .../systemds/scuro/models/training_stats.py | 10 +- 16 files changed, 602 insertions(+), 251 deletions(-) diff --git a/src/main/python/systemds/scuro/__init__.py b/src/main/python/systemds/scuro/__init__.py index bea9fab2d1f..58116b81bc4 100644 --- a/src/main/python/systemds/scuro/__init__.py +++ b/src/main/python/systemds/scuro/__init__.py @@ -18,158 +18,113 @@ # under the License. # # ------------------------------------------------------------- -from systemds.scuro.dataloader.base_loader import BaseLoader -from systemds.scuro.dataloader.audio_loader import AudioLoader -from systemds.scuro.dataloader.image_loader import ImageLoader -from systemds.scuro.dataloader.video_loader import VideoLoader -from systemds.scuro.dataloader.text_loader import TextLoader -from systemds.scuro.dataloader.json_loader import JSONLoader -from systemds.scuro.dataloader.tabular_loader import TabularLoader -from systemds.scuro.representations.representation import ( - FitRequirement, - Representation, - RepresentationFitState, -) -from systemds.scuro.representations.aggregate import Aggregation -from systemds.scuro.representations.aggregated_representation import ( - AggregatedRepresentation, -) -from systemds.scuro.representations.average import Average -from systemds.scuro.representations.bert import ( - Bert, - RoBERTa, - DistillBERT, - ALBERT, - ELECTRA, -) -from systemds.scuro.representations.bow import BoW -from systemds.scuro.representations.concatenation import Concatenation -from systemds.scuro.representations.context import Context -from systemds.scuro.representations.fusion import Fusion -from systemds.scuro.representations.glove import GloVe -from systemds.scuro.representations.lstm import LSTM -from systemds.scuro.representations.max import RowMax -from systemds.scuro.representations.mel_spectrogram import MelSpectrogram -from systemds.scuro.representations.multimodal_attention_fusion import ( - AttentionFusion, -) -from systemds.scuro.representations.timeseries_representations import ( - Mean, - Max, - Min, - Kurtosis, - Skew, - Std, - RMS, - ACF, - FrequencyMagnitude, - SpectralCentroid, - Quantile, - ZeroCrossingRate, - BandpowerFFT, -) -from systemds.scuro.representations.mfcc import MFCC -from systemds.scuro.representations.hadamard import Hadamard -from systemds.scuro.representations.optical_flow import OpticalFlow -from systemds.scuro.representations.representation import Representation -from systemds.scuro.representations.representation_dataloader import NPY -from systemds.scuro.representations.representation_dataloader import JSON -from systemds.scuro.representations.representation_dataloader import Pickle -from systemds.scuro.representations.resnet import ResNet -from systemds.scuro.representations.openface import OpenFace - -try: - from systemds.scuro.representations.image_bind import ImageBind -except ImportError as _imagebind_import_error: # pragma: no cover - import warnings as _warnings +from importlib import import_module - _warnings.warn( - f"ImageBind representation unavailable, it is excluded from the search " - f"({_imagebind_import_error})." - ) - ImageBind = None - -from systemds.scuro.representations.spectrogram import Spectrogram -from systemds.scuro.representations.sum import Sum -from systemds.scuro.representations.swin_video_transformer import SwinVideoTransformer -from systemds.scuro.representations.tabular_features import TabularFeatures -from systemds.scuro.representations.tfidf import TfIdf -from systemds.scuro.representations.unimodal import UnimodalRepresentation -from systemds.scuro.representations.wav2vec import Wav2Vec -from systemds.scuro.representations.window_aggregation import ( - WindowAggregation, - DynamicWindow, - StaticWindow, -) -from systemds.scuro.representations.physiological_window import ( - AdaptiveWindow, - PhysiologicalEventWindow, -) -from systemds.scuro.representations.word2vec import W2V -from systemds.scuro.representations.x3d import X3D -from systemds.scuro.representations.color_histogram import ColorHistogram -from systemds.scuro.models.model import Model -from systemds.scuro.models.discrete_model import DiscreteModel -from systemds.scuro.modality.joined import JoinedModality -from systemds.scuro.modality.joined_transformed import JoinedTransformedModality -from systemds.scuro.modality.modality import Modality -from systemds.scuro.modality.indexed import IndexedModality -from systemds.scuro.utils.identifier import Identifier -from systemds.scuro.modality.transformed import TransformedModality -from systemds.scuro.modality.type import ModalityType -from systemds.scuro.modality.unimodal_modality import UnimodalModality -from systemds.scuro.drsearch.task import Task -from systemds.scuro.drsearch.evaluation_split import EvaluationSplit -from systemds.scuro.drsearch.operator_registry import Registry -from systemds.scuro.drsearch.optimization_data import OptimizationData -from systemds.scuro.representations.covarep_audio_features import ( - RMSE, - Spectral, - ZeroCrossing, - Pitch, -) -from systemds.scuro.drsearch.multimodal_optimizer import MultimodalOptimizer -from systemds.scuro.drsearch.unimodal_optimizer import UnimodalOptimizer -from systemds.scuro.representations.vgg import VGG19 -from systemds.scuro.representations.clip import CLIPText, CLIPVisual -from systemds.scuro.representations.text_context import ( - SentenceBoundarySplit, - OverlappingSplit, -) -from systemds.scuro.representations.text_context_with_indices import ( - SentenceBoundarySplitIndices, - OverlappingSplitIndices, -) -from systemds.scuro.representations.elmo import ELMoRepresentation -from systemds.scuro.representations.dimensionality_reduction import ( - DimensionalityReduction, -) -from systemds.scuro.representations.mlp_averaging import MLPAveraging -from systemds.scuro.representations.mlp_learned_dim_reduction import ( - MLPLearnedDimReduction, +_EXPORT_GROUPS = ( + ("dataloader.base_loader", "BaseLoader"), + ("dataloader.audio_loader", "AudioLoader"), + ("dataloader.image_loader", "ImageLoader"), + ("dataloader.video_loader", "VideoLoader"), + ("dataloader.text_loader", "TextLoader"), + ("dataloader.json_loader", "JSONLoader"), + ("dataloader.tabular_loader", "TabularLoader"), + ( + "representations.representation", + "FitRequirement Representation RepresentationFitState", + ), + ("representations.aggregate", "Aggregation"), + ("representations.aggregated_representation", "AggregatedRepresentation"), + ("representations.average", "Average"), + ("representations.bert", "Bert RoBERTa DistillBERT ALBERT ELECTRA"), + ("representations.bow", "BoW"), + ("representations.concatenation", "Concatenation"), + ("representations.context", "Context"), + ("representations.fusion", "Fusion"), + ("representations.glove", "GloVe"), + ("representations.lstm", "LSTM"), + ("representations.max", "RowMax"), + ("representations.mel_spectrogram", "MelSpectrogram"), + ("representations.multimodal_attention_fusion", "AttentionFusion"), + ( + "representations.timeseries_representations", + "Mean Max Min Kurtosis Skew Std RMS ACF FrequencyMagnitude SpectralCentroid Quantile ZeroCrossingRate BandpowerFFT", + ), + ("representations.mfcc", "MFCC"), + ("representations.hadamard", "Hadamard"), + ("representations.optical_flow", "OpticalFlow"), + ("representations.representation_dataloader", "NPY JSON Pickle"), + ("representations.resnet", "ResNet"), + ("representations.openface", "OpenFace"), + ("representations.image_bind", "ImageBind"), + ("representations.spectrogram", "Spectrogram"), + ("representations.sum", "Sum"), + ("representations.swin_video_transformer", "SwinVideoTransformer"), + ("representations.tabular_features", "TabularFeatures"), + ("representations.tfidf", "TfIdf"), + ("representations.unimodal", "UnimodalRepresentation"), + ("representations.wav2vec", "Wav2Vec"), + ( + "representations.window_aggregation", + "WindowAggregation DynamicWindow StaticWindow", + ), + ("representations.physiological_window", "AdaptiveWindow PhysiologicalEventWindow"), + ("representations.word2vec", "W2V"), + ("representations.x3d", "X3D"), + ("representations.color_histogram", "ColorHistogram"), + ("models.model", "Model"), + ("models.discrete_model", "DiscreteModel"), + ("modality.joined", "JoinedModality"), + ("modality.joined_transformed", "JoinedTransformedModality"), + ("modality.modality", "Modality"), + ("modality.indexed", "IndexedModality"), + ("utils.identifier", "Identifier"), + ("modality.transformed", "TransformedModality"), + ("modality.type", "ModalityType"), + ("modality.unimodal_modality", "UnimodalModality"), + ("drsearch.task", "Task"), + ("drsearch.evaluation_split", "EvaluationSplit"), + ("drsearch.operator_registry", "Registry"), + ("drsearch.optimization_data", "OptimizationData"), + ("drsearch.multimodal_optimizer", "MultimodalOptimizer"), + ("drsearch.unimodal_optimizer", "UnimodalOptimizer"), + ("representations.covarep_audio_features", "RMSE Spectral ZeroCrossing Pitch"), + ("representations.vgg", "VGG19"), + ("representations.clip", "CLIPText CLIPVisual"), + ("representations.text_context", "SentenceBoundarySplit OverlappingSplit"), + ( + "representations.text_context_with_indices", + "SentenceBoundarySplitIndices OverlappingSplitIndices", + ), + ("representations.elmo", "ELMoRepresentation"), + ("representations.dimensionality_reduction", "DimensionalityReduction"), + ("representations.mlp_averaging", "MLPAveraging"), + ("representations.mlp_learned_dim_reduction", "MLPLearnedDimReduction"), + ( + "representations.physiological_representations", + "SDNN RMSSD pNN RRPerMinute HRVBandPower HRVVLF HRVLF HRVHF HRVLFHF PoincareSD1 PoincareSD2 SCLSlope SCLDynamicRange SCRPeaksPerMinute SCRAverageAmplitude SCRAverageDuration", + ), ) +_EXPORTS = { + name: (module_name, name) + for module_name, names in _EXPORT_GROUPS + for name in names.split() +} -from systemds.scuro.representations.physiological_representations import ( - SDNN, - RMSSD, - pNN, - RRPerMinute, - HRVBandPower, - HRVVLF, - HRVLF, - HRVHF, - HRVLFHF, - PoincareSD1, - PoincareSD2, - # SampleEntropy, - # DFAAlpha, - SCLSlope, - SCLDynamicRange, - SCRPeaksPerMinute, - SCRAverageAmplitude, - SCRAverageDuration, -) +def __getattr__(name): + try: + module_name, attribute = _EXPORTS[name] + except KeyError as exc: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from exc + try: + value = getattr(import_module(f"{__name__}.{module_name}"), attribute) + except ImportError: + if name != "ImageBind": + raise + value = None + globals()[name] = value + return value + __all__ = [ "BaseLoader", @@ -262,7 +217,6 @@ "MLPAveraging", "MLPLearnedDimReduction", "DimensionalityReduction", - "MultimodalGAPymooOptimizer", "SDNN", "RMSSD", "pNN", diff --git a/src/main/python/systemds/scuro/dataloader/audio_loader.py b/src/main/python/systemds/scuro/dataloader/audio_loader.py index 1c7ae4f3a85..b5ddcb07677 100644 --- a/src/main/python/systemds/scuro/dataloader/audio_loader.py +++ b/src/main/python/systemds/scuro/dataloader/audio_loader.py @@ -19,6 +19,7 @@ # # ------------------------------------------------------------- from dataclasses import dataclass +import math import os from typing import List, Optional, Union import librosa @@ -56,7 +57,12 @@ def __init__( ) self.normalize = normalize self.load_data_from_file = load - self.stats = self.get_stats(source_path) + files = self.get_file_names(self.indices) + self.stats = self.cached_stats( + files, + {"dtype": str(data_type), "normalize": normalize}, + lambda: self.get_stats(source_path), + ) def extract(self, file: str, index: Optional[Union[str, List[str]]] = None): self.file_sanity_check(file) @@ -80,15 +86,20 @@ def get_stats(self, source_path: str): max_length = 0 avg_length = 0 num_instances = 0 - for file in os.listdir(source_path): - file_name = file.split(".")[0] - if file_name not in self.indices: - continue - self.file_sanity_check(source_path + file) - audio, sr = librosa.load(source_path + file, dtype=self._data_type) + import soundfile + + for file in self.get_file_names(self.indices): + self.file_sanity_check(file) + try: + info = soundfile.info(file) + sr = 22050 + length = math.ceil(info.frames * sr / info.samplerate) + except (RuntimeError, TypeError): + audio, sr = librosa.load(file, dtype=self._data_type) + length = audio.shape[0] num_instances += 1 sampling_rate = max(sampling_rate, sr) - max_length = max(max_length, audio.shape[0]) - avg_length += audio.shape[0] + max_length = max(max_length, length) + avg_length += length avg_length /= num_instances return AudioStats(sampling_rate, max_length, avg_length, num_instances, True) diff --git a/src/main/python/systemds/scuro/dataloader/base_loader.py b/src/main/python/systemds/scuro/dataloader/base_loader.py index e58a755838c..9deaa666826 100644 --- a/src/main/python/systemds/scuro/dataloader/base_loader.py +++ b/src/main/python/systemds/scuro/dataloader/base_loader.py @@ -18,7 +18,10 @@ # under the License. # # ------------------------------------------------------------- +import hashlib +import json import os +import pickle from abc import ABC, abstractmethod from collections.abc import Sequence from typing import Callable, Iterator, List, Optional, Tuple, Union @@ -77,6 +80,50 @@ def __init__( self.stats = None self.chunk_size = chunk_size + def cached_stats(self, files, options, compute): + if os.environ.get("SCURO_DISABLE_STATS_CACHE", "0") == "1": + return compute() + if isinstance(files, (str, os.PathLike)): + files = [files] + fingerprint = [] + for file in files: + path = os.path.realpath(os.fspath(file)) + stat = os.stat(path) + fingerprint.append((path, stat.st_size, stat.st_mtime_ns)) + key = json.dumps( + { + "schema": 2, + "loader": f"{type(self).__module__}.{type(self).__qualname__}", + "files": fingerprint, + "options": options, + }, + sort_keys=True, + default=str, + ).encode() + digest = hashlib.sha256(key).hexdigest() + cache_root = os.environ.get("SCURO_STATS_CACHE_DIR") + if cache_root is None: + cache_root = os.path.join( + os.environ.get("XDG_CACHE_HOME", os.path.expanduser("~/.cache")), + "scuro", + "loader_stats", + ) + cache_path = os.path.join(cache_root, f"{digest}.pkl") + try: + with open(cache_path, "rb") as cache_file: + return pickle.load(cache_file) + except (FileNotFoundError, EOFError, OSError, pickle.PickleError): + stats = compute() + try: + os.makedirs(cache_root, exist_ok=True) + temporary_path = f"{cache_path}.{os.getpid()}.tmp" + with open(temporary_path, "wb") as cache_file: + pickle.dump(stats, cache_file, protocol=pickle.HIGHEST_PROTOCOL) + os.replace(temporary_path, cache_path) + except OSError: + pass + return stats + @property def chunk_size(self): return self._chunk_size diff --git a/src/main/python/systemds/scuro/dataloader/image_loader.py b/src/main/python/systemds/scuro/dataloader/image_loader.py index 8da151a113e..35394448dfb 100644 --- a/src/main/python/systemds/scuro/dataloader/image_loader.py +++ b/src/main/python/systemds/scuro/dataloader/image_loader.py @@ -56,7 +56,12 @@ def __init__( ) self.load_data_from_file = load self._all_metadata = [] - self.stats = self.get_stats(source_path) + files = self.get_file_names(self.indices) + self.stats, self._all_metadata = self.cached_stats( + files, + {"dtype": str(data_type), "extension": ext}, + lambda: (self.get_stats(source_path), self._all_metadata), + ) def load(self): if self.chunk_size: @@ -102,11 +107,14 @@ def get_stats(self, source_path: str): for file in self.indices: path = os.path.join(source_path, f"{file}{self._ext}") self.file_sanity_check(path) - image = cv2.imread(path, cv2.IMREAD_COLOR) - if image is None: - raise ValueError(f"Could not read image at path: {path}") - image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) - height, width, channels = image.shape + from PIL import Image + + try: + with Image.open(path) as image: + width, height = image.size + except OSError as exc: + raise ValueError(f"Could not read image at path: {path}") from exc + channels = 3 self._all_metadata.append( self.modality_type.create_metadata(width, height, channels) ) diff --git a/src/main/python/systemds/scuro/dataloader/json_loader.py b/src/main/python/systemds/scuro/dataloader/json_loader.py index adb3f6aaf64..aaf11ef6d84 100644 --- a/src/main/python/systemds/scuro/dataloader/json_loader.py +++ b/src/main/python/systemds/scuro/dataloader/json_loader.py @@ -52,7 +52,12 @@ def __init__( source_path, indices, data_type, chunk_size, ModalityType.TEXT, ext ) self.field = field - self.stats = self.get_stats(source_path) + files = self.get_file_names(self.indices) + self.stats = self.cached_stats( + files, + {"field": field, "extension": ext}, + lambda: self.get_stats(source_path), + ) def extract(self, file: str, index: Optional[Union[str, List[str]]] = None): self.file_sanity_check(file) diff --git a/src/main/python/systemds/scuro/dataloader/tabular_loader.py b/src/main/python/systemds/scuro/dataloader/tabular_loader.py index 6652ba2b937..df131075ca1 100644 --- a/src/main/python/systemds/scuro/dataloader/tabular_loader.py +++ b/src/main/python/systemds/scuro/dataloader/tabular_loader.py @@ -52,7 +52,16 @@ def __init__( self.file_format = file_format.lower() if self.file_format != "npy": raise ValueError(f"Unsupported file format: {self.file_format}") - self.stats = self.get_stats(source_path) + files = [source_path + name + "." + self.file_format for name in self.indices] + self.stats = self.cached_stats( + files, + { + "dtype": str(data_type), + "features": feature_names, + "format": self.file_format, + }, + lambda: self.get_stats(source_path), + ) def extract(self, file: str, index: Optional[Union[str, List[str]]] = None): self.file_sanity_check(file) @@ -72,7 +81,7 @@ def get_stats(self, source_path: str) -> TabularStats: for file_name in self.indices: file = source_path + file_name + "." + self.file_format self.file_sanity_check(file) - data = np.load(file) + data = np.load(file, mmap_mode="r") num_features = max(num_features, int(np.prod(data.shape))) num_instances += 1 return TabularStats(num_instances, num_features, (num_features,)) diff --git a/src/main/python/systemds/scuro/dataloader/timeseries_loader.py b/src/main/python/systemds/scuro/dataloader/timeseries_loader.py index 685105e31e5..669022cda41 100644 --- a/src/main/python/systemds/scuro/dataloader/timeseries_loader.py +++ b/src/main/python/systemds/scuro/dataloader/timeseries_loader.py @@ -61,7 +61,18 @@ def __init__( self.channel_index = channel_index self.file_format = file_format.lower() self.channel_index = channel_index - self.stats = self.get_stats(source_path, sampling_rate) + files = [source_path + name + "." + self.file_format for name in self.indices] + self.stats = self.cached_stats( + files, + { + "dtype": str(data_type), + "signals": signal_names, + "sampling_rate": sampling_rate, + "format": self.file_format, + "channel_index": channel_index, + }, + lambda: self.get_stats(source_path, sampling_rate), + ) if self.file_format not in ["npy", "mat", "hdf5", "txt"]: raise ValueError(f"Unsupported file format: {self.file_format}") @@ -162,7 +173,11 @@ def get_stats(self, source_path: str, sampling_rate: int): avg_length = 0 for file_name in self.indices: file = source_path + file_name + "." + self.file_format - data = self._load_data(file) + data = ( + np.load(file, mmap_mode="r") + if self.file_format == "npy" + else self._load_data(file) + ) if self.channel_index is not None and data.ndim > 1: data = data[:, self.channel_index] max_length = max(max_length, data.shape[0]) diff --git a/src/main/python/systemds/scuro/dataloader/video_loader.py b/src/main/python/systemds/scuro/dataloader/video_loader.py index e2ee7b72d24..55ac34dc934 100644 --- a/src/main/python/systemds/scuro/dataloader/video_loader.py +++ b/src/main/python/systemds/scuro/dataloader/video_loader.py @@ -81,7 +81,12 @@ def __init__( self.fps = fps self.target_size = tuple(int(v) for v in target_size) if target_size else None self._all_metadata = [] - self.stats = self.get_stats(source_path) + files = self.get_file_names(self.indices) + self.stats, self._all_metadata = self.cached_stats( + files, + {"dtype": str(data_type), "fps": fps, "target_size": self.target_size}, + lambda: (self.get_stats(source_path), self._all_metadata), + ) def load(self): if self.chunk_size: diff --git a/src/main/python/systemds/scuro/drsearch/node_executor.py b/src/main/python/systemds/scuro/drsearch/node_executor.py index d42a096d77f..b30f0a5df07 100644 --- a/src/main/python/systemds/scuro/drsearch/node_executor.py +++ b/src/main/python/systemds/scuro/drsearch/node_executor.py @@ -503,7 +503,7 @@ def __init__( gpu_slots_per_device=int( os.environ.get("SCURO_GPU_SLOTS_PER_DEVICE", "1") ), - gpu_demand_fraction=self.scheduler.gpu_demand_fraction(), + gpu_demand_fraction=1.0, ) self._pool = worker_pool restrict_devices = getattr(self.scheduler, "restrict_gpu_devices", None) @@ -1046,6 +1046,10 @@ def _discard_report(self) -> Dict[str, Any]: def _is_task_node(node: RepresentationNode) -> bool: return bool(getattr(node, "parameters", {}).get("_node_kind") == "task") + def detach_worker_pool(self) -> PersistentWorkerPool: + self._owns_pool = False + return self._pool + def run(self) -> Dict[str, Any]: self._task_results = {} try: diff --git a/src/main/python/systemds/scuro/drsearch/operator_registry.py b/src/main/python/systemds/scuro/drsearch/operator_registry.py index d9ae798f867..1dff015a905 100644 --- a/src/main/python/systemds/scuro/drsearch/operator_registry.py +++ b/src/main/python/systemds/scuro/drsearch/operator_registry.py @@ -18,11 +18,80 @@ # under the License. # # ------------------------------------------------------------- -from typing import Union, List +from importlib import import_module +from typing import Iterable, Union, List import math from systemds.scuro.modality.type import ModalityType from systemds.scuro.representations.representation import Representation +_BUILTIN_MODULES = { + ModalityType.AUDIO: ( + "covarep_audio_features", + "mel_spectrogram", + "mfcc", + "spectrogram", + "timeseries_representations", + "wav2vec", + ), + ModalityType.TEXT: ("bert", "bow", "clip", "glove", "tfidf", "word2vec"), + ModalityType.IMAGE: ("clip", "color_histogram", "openface", "resnet", "vgg"), + ModalityType.VIDEO: ( + "clip", + "color_histogram", + "image_bind", + "openface", + "resnet", + "swin_video_transformer", + "vgg", + "x3d", + ), + ModalityType.TIMESERIES: ("timeseries_representations",), + ModalityType.PHYSIOLOGICAL: ( + "physiological_representations", + "timeseries_representations", + ), + ModalityType.EMBEDDING: ("tabular_features",), +} +_CONTEXT_MODULES = { + ModalityType.AUDIO: ("window_aggregation",), + ModalityType.TEXT: ("text_context", "text_context_with_indices"), + ModalityType.VIDEO: ("window_aggregation",), + ModalityType.TIMESERIES: ("window_aggregation",), + ModalityType.PHYSIOLOGICAL: ("physiological_window", "window_aggregation"), +} +_FUSION_MODULES = ( + "average", + "concatenation", + "hadamard", + "lstm", + "max", + "multimodal_attention_fusion", + "sum", +) +_DIMENSIONALITY_REDUCTION_MODULES = ("mlp_averaging",) +_OPTIONAL_MODULES = {"image_bind"} +_loaded_builtin_modules = set() + + +def load_builtin_operators( + modality_types: Iterable[ModalityType], include_fusion: bool = False +): + modules = set() + for modality_type in modality_types: + modules.update(_BUILTIN_MODULES.get(modality_type, ())) + modules.update(_CONTEXT_MODULES.get(modality_type, ())) + if modality_type == ModalityType.EMBEDDING: + modules.update(_DIMENSIONALITY_REDUCTION_MODULES) + if include_fusion: + modules.update(_FUSION_MODULES) + for module_name in sorted(modules - _loaded_builtin_modules): + try: + import_module(f"systemds.scuro.representations.{module_name}") + except ImportError: + if module_name not in _OPTIONAL_MODULES: + raise + _loaded_builtin_modules.add(module_name) + class Registry: """ @@ -116,6 +185,7 @@ def add_context_representation_operator( ) def get_representations(self, modality: ModalityType): + load_builtin_operators((modality,)) return self._representations[modality] def get_not_self_contained_representations(self, modality: ModalityType): @@ -126,14 +196,17 @@ def get_not_self_contained_representations(self, modality: ModalityType): return reps def get_context_operators(self, modality_type): + load_builtin_operators((modality_type,)) if modality_type not in self._context_operators.keys(): return [] return self._context_operators[modality_type] def get_dimensionality_reduction_operators(self, modality_type): + load_builtin_operators((modality_type,)) return self._dimensionality_reduction_operators.get(modality_type, []) def get_fusion_operators(self): + load_builtin_operators((), include_fusion=True) return type(self)._fusion_operators def get_fusion_operator_by_name(self, fusion_name): diff --git a/src/main/python/systemds/scuro/drsearch/ranking.py b/src/main/python/systemds/scuro/drsearch/ranking.py index 935bc565356..7b92ab0b130 100644 --- a/src/main/python/systemds/scuro/drsearch/ranking.py +++ b/src/main/python/systemds/scuro/drsearch/ranking.py @@ -24,18 +24,6 @@ import numpy as np -def _operator_signature(entry) -> frozenset: - """Set of operator class names in an entry's DAG, ignoring hyperparameters.""" - try: - return frozenset( - node.operation.__name__ - for node in entry.dag.nodes - if getattr(node, "operation", None) is not None - ) - except Exception: - return frozenset() - - def _dag_size(entry) -> int: try: return sum( @@ -51,8 +39,7 @@ def rank_by_robustness( entries: Iterable, *, performance_metric_name: str = "accuracy", - neighbourhood_weight: float = 0.5, - sharpness: int = 4, + standard_deviation_weight: float = 1.0, one_se_parsimony: bool = True, cache_scores: bool = True, score_attr: str = "robustness_score", @@ -61,6 +48,9 @@ def rank_by_robustness( if not entries: return [], [] + if not np.isfinite(standard_deviation_weight) or standard_deviation_weight < 0: + raise ValueError("standard_deviation_weight must be finite and non-negative") + def perf_of(entry): if entry is None: return None @@ -78,58 +68,50 @@ def perf_of(entry): if not indexed_entries: return [], [] + def variation_of(entry): + try: + fold_values = np.asarray( + entry.val_fold_scores[performance_metric_name], dtype=float + ) + fold_values = fold_values[np.isfinite(fold_values)] + except (AttributeError, KeyError, TypeError, ValueError): + fold_values = np.array([], dtype=float) + + if fold_values.size > 1: + return float(fold_values.std(ddof=1)), int(fold_values.size) + + try: + std = float(entry.val_score_stds[performance_metric_name]) + except (AttributeError, KeyError, TypeError, ValueError): + std = 0.0 + if not np.isfinite(std) or std < 0: + std = 0.0 + return std, max(1, int(fold_values.size)) + original_indices, entries, performance = zip(*indexed_entries) entries = list(entries) perf = np.array(performance, dtype=float) sizes = np.array([_dag_size(e) if e is not None else 0 for e in entries], float) - - smoothed = perf - if neighbourhood_weight > 0.0 and len(entries) > 1: - signatures = [ - _operator_signature(e) if e is not None else frozenset() for e in entries - ] - vocabulary = sorted({op for sig in signatures for op in sig}) - if vocabulary: - position = {op: i for i, op in enumerate(vocabulary)} - membership = np.zeros((len(entries), len(vocabulary)), dtype=np.float32) - for row, sig in enumerate(signatures): - for op in sig: - membership[row, position[op]] = 1.0 - intersection = membership @ membership.T - counts = membership.sum(1) - union = counts[:, None] + counts[None, :] - intersection - jaccard = np.divide( - intersection, - union, - out=np.zeros_like(intersection), - where=union > 0, - ) - weights = jaccard**sharpness - denominator = weights.sum(1) - neighbourhood = np.divide( - (weights * perf[None, :].astype(np.float32)).sum(1), - denominator, - out=perf.astype(np.float32).copy(), - where=denominator > 0, - ) - smoothed = ( - 1.0 - neighbourhood_weight - ) * perf + neighbourhood_weight * neighbourhood + variations = [variation_of(entry) for entry in entries] + stds = np.array([variation[0] for variation in variations], dtype=float) + fold_counts = np.array([variation[1] for variation in variations], dtype=int) + robustness = perf - standard_deviation_weight * stds if cache_scores: - for entry, score in zip(entries, smoothed): + for entry, score in zip(entries, robustness): if entry is not None: setattr(entry, score_attr, float(score)) if one_se_parsimony and len(entries) > 1: - standard_error = float(smoothed.std()) / np.sqrt(len(smoothed)) - threshold = float(smoothed.max()) - standard_error + best_index = int(np.argmax(robustness)) + standard_error = stds[best_index] / np.sqrt(fold_counts[best_index]) + threshold = float(robustness[best_index] - standard_error) keys = [ (True, -sz, float(s)) if s >= threshold else (False, 0.0, float(s)) - for s, sz in zip(smoothed, sizes) + for s, sz in zip(robustness, sizes) ] else: - keys = [(True, 0.0, float(s)) for s in smoothed] + keys = [(True, 0.0, float(s)) for s in robustness] local_indices = sorted(range(len(entries)), key=lambda i: keys[i], reverse=True) sorted_entries = [entries[i] for i in local_indices] diff --git a/src/main/python/systemds/scuro/drsearch/task.py b/src/main/python/systemds/scuro/drsearch/task.py index 05cbb761c88..33fe1b43fe8 100644 --- a/src/main/python/systemds/scuro/drsearch/task.py +++ b/src/main/python/systemds/scuro/drsearch/task.py @@ -25,7 +25,12 @@ from systemds.scuro.models.model import Model from systemds.scuro.models.training_stats import TrainingWorkAccumulator import numpy as np -from sklearn.model_selection import train_test_split +from sklearn.model_selection import ( + GroupKFold, + KFold, + StratifiedGroupKFold, + StratifiedKFold, +) from systemds.scuro.drsearch.evaluation_split import EvaluationSplit from systemds.scuro.representations.representation import RepresentationStats @@ -88,6 +93,9 @@ def __init__( kfold=5, measure_performance=True, performance_measures=["accuracy"], + groups=None, + allow_group_overlap=False, + cv_seed=11, ): """ Parent class for the prediction task that is performed on top of the aligned representation @@ -105,6 +113,9 @@ def __init__( self.train_indices = train_indices self.test_indices = val_indices 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.measure_performance = measure_performance self.inference_time = [] self.training_time = [] @@ -205,6 +216,11 @@ def estimate_relative_cost(self, input_stats) -> float: ) fold_sizes = [len(fold) for fold in self.cv_train_indices] n_fold_train = max(fold_sizes, default=len(self.train_indices or [])) + num_steps = getattr(self.model, "num_steps", None) + if num_steps is not None: + batch_size = max(1, int(getattr(self.model, "batch_size", 1) or 1)) + samples = max(1, int(num_steps)) * batch_size + return float(max(1, self.kfold) * samples * max(1, feature_dim)) passes = max(1, int(getattr(self.model, "epochs", 1) or 1)) return float( max(1, self.kfold) * passes * max(1, n_fold_train) * max(1, feature_dim) @@ -263,26 +279,101 @@ def validate_input_memory(self, data, gpu_memory_fraction: float = 0.9) -> None: "before allocating CUDA tensors." ) + @staticmethod + def _stratification_target(labels_array): + if labels_array.ndim == 1: + return labels_array + return None + + def _fold_candidates(self, y, n_units): + n_splits = max(2, min(self.kfold, int(n_units))) + stratify_on = self._stratification_target(y) + candidates = [] + if self.groups is not None: + if stratify_on is not None: + candidates.append( + ( + StratifiedGroupKFold( + n_splits=n_splits, shuffle=True, random_state=self.cv_seed + ), + stratify_on, + ) + ) + candidates.append((GroupKFold(n_splits=n_splits), None)) + else: + if stratify_on is not None: + candidates.append( + ( + StratifiedKFold( + n_splits=n_splits, shuffle=True, random_state=self.cv_seed + ), + stratify_on, + ) + ) + candidates.append( + ( + KFold(n_splits=n_splits, shuffle=True, random_state=self.cv_seed), + None, + ) + ) + return candidates + + def _check_train_test_groups(self): + if self.groups is None or not len(self.test_indices or []): + return + train_groups = set(self.groups[np.asarray(self.train_indices)].tolist()) + test_groups = set(self.groups[np.asarray(self.test_indices)].tolist()) + shared = train_groups & test_groups + if shared and not self.allow_group_overlap: + raise ValueError( + f"{self.name}: train and test share {len(shared)} group(s) " + f"(e.g. {sorted(map(str, shared))[:5]}). Use a group-disjoint " + "outer split, or pass allow_group_overlap=True to declare a " + "group-dependent protocol explicitly." + ) + self.train_test_group_overlap = sorted(map(str, shared)) + def _create_cv_splits(self): - train_labels = [self.labels[i] for i in self.train_indices] - train_labels_array = np.array(train_labels) + self._check_train_test_groups() + + train_indices_array = np.asarray(self.train_indices) + y = np.asarray([self.labels[i] for i in self.train_indices]) + fold_groups = None if self.groups is None else self.groups[train_indices_array] - train_indices_array = np.array(self.train_indices) + n_units = ( + len(train_indices_array) + if fold_groups is None + else len(set(fold_groups.tolist())) + ) + folds = None + reason = None + for splitter, stratify_on in self._fold_candidates(y, n_units): + try: + folds = list( + splitter.split(train_indices_array, stratify_on, groups=fold_groups) + ) + except ValueError as error: + reason = f"{type(splitter).__name__}: {error}" + continue + self.cv_stratified = stratify_on is not None + self.cv_fallback_reason = reason + break + if folds is None: + raise ValueError( + f"{self.name}: cannot build {self.kfold} folds over " + f"{len(train_indices_array)} training instances " + f"({n_units} unit(s)): {reason}" + ) + self.kfold = len(folds) + + self.cv_grouped = self.groups is not None self.cv_train_indices = [] self.cv_val_indices = [] - for fold_idx in range(self.kfold): - fold_train_indices_array, fold_val_indices_array, _, _ = train_test_split( - train_indices_array, - train_labels_array, - test_size=0.2, - shuffle=True, - random_state=11 + fold_idx, - ) - - fold_train_indices = fold_train_indices_array.tolist() - fold_val_indices = fold_val_indices_array.tolist() + for fold_idx, (fit_positions, val_positions) in enumerate(folds): + fold_train_indices = train_indices_array[fit_positions].tolist() + fold_val_indices = train_indices_array[val_positions].tolist() self.cv_train_indices.append(fold_train_indices) self.cv_val_indices.append(fold_val_indices) @@ -293,6 +384,16 @@ def _create_cv_splits(self): f"Fold {fold_idx}: Overlap detected between train and val indices: {overlap}" ) + if fold_groups is not None: + shared_groups = set(fold_groups[fit_positions].tolist()) & set( + fold_groups[val_positions].tolist() + ) + if shared_groups: + raise ValueError( + f"Fold {fold_idx}: groups appear in both the fit and the " + f"validation part: {sorted(map(str, shared_groups))[:5]}" + ) + self.evaluation_splits = tuple( EvaluationSplit(train_indices, val_indices, name=f"fold_{fold_idx}") for fold_idx, (train_indices, val_indices) in enumerate( @@ -300,6 +401,31 @@ def _create_cv_splits(self): ) ) + def split_report(self): + report = { + "task": self.name, + "n_train": len(self.train_indices), + "n_test": 0 if self.test_indices is None else len(self.test_indices), + "kfold": self.kfold, + "grouped": bool(getattr(self, "cv_grouped", False)), + "stratified": bool(getattr(self, "cv_stratified", False)), + "cv_seed": self.cv_seed, + } + if self.groups is not None: + report["n_groups_train"] = len( + set(self.groups[np.asarray(self.train_indices)].tolist()) + ) + if len(self.test_indices or []): + report["n_groups_test"] = len( + set(self.groups[np.asarray(self.test_indices)].tolist()) + ) + report["train_test_group_overlap"] = len( + getattr(self, "train_test_group_overlap", []) + ) + if getattr(self, "cv_fallback_reason", None): + report["stratification_fallback"] = self.cv_fallback_reason + return report + def create_model(self): """ Return a fresh, unfitted model instance. diff --git a/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py b/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py index 218feaf0d92..bd18a9beb52 100644 --- a/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py +++ b/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py @@ -86,6 +86,8 @@ def __init__( self.checkpoint_every = checkpoint_every self.resume = resume self.max_num_workers = max_num_workers + self._persistent_worker_pool = None + self._reuse_worker_pool = False self._checkpoint_manager = CheckpointManager( self.result_path or ".", "unimodal_checkpoint_", @@ -262,6 +264,16 @@ def optimize_parallel(self, n_workers=None): continue def optimize(self): + self._reuse_worker_pool = True + try: + return self._optimize_modalities() + finally: + if self._persistent_worker_pool is not None: + self._persistent_worker_pool.shutdown() + self._persistent_worker_pool = None + self._reuse_worker_pool = False + + def _optimize_modalities(self): if self.resume: self.resume_from_checkpoint() @@ -379,8 +391,11 @@ def _process_modality(self, modality, skip_remaining: int = 0, scheduler=None): max_num_workers=self.max_num_workers, result_path=self.result_path, enable_checkpointing=self.enable_checkpointing, + worker_pool=self._persistent_worker_pool, search_start=self._search_start, ) + if self._reuse_worker_pool and self._persistent_worker_pool is None: + self._persistent_worker_pool = node_executor.detach_worker_pool() exec_out = node_executor.run() task_results.extend(exec_out["task_results"]) statistics = exec_out["statistics"] @@ -875,6 +890,9 @@ def _summarize_windows(self, builder, node_id): ) +SELECTION_STRATEGIES = ("performance", "robustness", "tradeoff") + + class UnimodalResults: def __init__( self, @@ -884,6 +902,10 @@ def __init__( store_cache=True, k=-1, metric_name="accuracy", + selection_strategy="performance", + tradeoff_weights=(0.5, 0.5), + robustness_std_weight=1.0, + robustness_one_se_parsimony=True, ): self.modality_ids = [modality.modality_id for modality in modalities] self.task_names = [task.model.name for task in tasks] @@ -893,6 +915,12 @@ def __init__( self.store_cache = store_cache self.k = k self.metric_name = metric_name + self.set_selection_strategy( + selection_strategy, + tradeoff_weights=tradeoff_weights, + robustness_std_weight=robustness_std_weight, + robustness_one_se_parsimony=robustness_one_se_parsimony, + ) for modality in self.modality_ids: self.results[modality] = {task_name: [] for task_name in self.task_names} self.cache[modality] = {task_name: [] for task_name in self.task_names} @@ -989,6 +1017,51 @@ def print_results(self): for entry in self.results[modality][task_name]: print(f"{modality}_{task_name}: {entry}") + def set_selection_strategy( + self, + selection_strategy, + tradeoff_weights=None, + robustness_std_weight=None, + robustness_one_se_parsimony=None, + ): + if selection_strategy not in SELECTION_STRATEGIES: + raise ValueError( + f"unknown selection strategy {selection_strategy!r}, " + f"expected one of {SELECTION_STRATEGIES}" + ) + self.selection_strategy = selection_strategy + if tradeoff_weights is not None: + self.tradeoff_weights = tuple(tradeoff_weights) + if robustness_std_weight is not None: + self.robustness_std_weight = robustness_std_weight + if robustness_one_se_parsimony is not None: + self.robustness_one_se_parsimony = robustness_one_se_parsimony + return self + + def with_selection_strategy(self, selection_strategy, **strategy_kwargs): + view = copy.copy(self) + return view.set_selection_strategy(selection_strategy, **strategy_kwargs) + + def _rank_entries(self, task_results, performance_metric_name): + strategy = getattr(self, "selection_strategy", "performance") + if strategy == "robustness": + return rank_by_robustness( + task_results, + performance_metric_name=performance_metric_name, + standard_deviation_weight=getattr(self, "robustness_std_weight", 1.0), + one_se_parsimony=getattr(self, "robustness_one_se_parsimony", True), + ) + weights = ( + getattr(self, "tradeoff_weights", (0.5, 0.5)) + if strategy == "tradeoff" + else (1.0, 0.0) + ) + return rank_by_tradeoff( + task_results, + weights=weights, + performance_metric_name=performance_metric_name, + ) + def get_k_best_results( self, modality, @@ -1007,8 +1080,8 @@ def get_k_best_results( task_results = self.results[modality.modality_id][task.model.name] - results, sorted_indices = rank_by_tradeoff( - task_results, performance_metric_name=performance_metric_name + results, sorted_indices = self._rank_entries( + task_results, performance_metric_name ) results = results[: self.k] @@ -1018,6 +1091,11 @@ def get_k_best_results( task_cache = self.cache.get(modality.modality_id, {}).get( task.model.name, None ) + if ( + isinstance(task_cache, list) + and getattr(self, "selection_strategy", "performance") != "performance" + ): + task_cache = None if not task_cache: cache = [] for result in results: @@ -1063,7 +1141,7 @@ def get_k_most_robust_results( task, performance_metric_name, k=None, - neighbourhood_weight=0.5, + standard_deviation_weight=1.0, one_se_parsimony=True, ): task_results = self.results[modality.modality_id][task.model.name] @@ -1071,7 +1149,7 @@ def get_k_most_robust_results( results, sorted_indices = rank_by_robustness( task_results, performance_metric_name=performance_metric_name, - neighbourhood_weight=neighbourhood_weight, + standard_deviation_weight=standard_deviation_weight, one_se_parsimony=one_se_parsimony, ) diff --git a/src/main/python/systemds/scuro/drsearch/worker_pool.py b/src/main/python/systemds/scuro/drsearch/worker_pool.py index 9969b67e626..4ab50abfebb 100644 --- a/src/main/python/systemds/scuro/drsearch/worker_pool.py +++ b/src/main/python/systemds/scuro/drsearch/worker_pool.py @@ -22,7 +22,9 @@ import multiprocessing as mp import multiprocessing.connection as mp_connection import os +import queue import signal +import time from dataclasses import dataclass from typing import Any, Callable, Dict, List, Optional, Tuple @@ -94,12 +96,15 @@ def _worker_main( result_q, dispatch: Dict[str, Callable], num_threads: int, + ready_q, physical_gpu_id: Optional[int], ) -> None: os.environ["CUDA_VISIBLE_DEVICES"] = ( "" if physical_gpu_id is None else str(physical_gpu_id) ) _worker_initializer(num_threads) + if ready_q is not None: + ready_q.put(os.getpid()) while True: job = job_q.get() if job is None: @@ -165,6 +170,7 @@ def __init__( self.gpu_devices = list(dict.fromkeys(gpu_devices or [])) self.gpu_slots_per_device = max(1, int(gpu_slots_per_device)) self._result_q = self._ctx.Queue() + self._ready_q = self._ctx.Queue() self._job_counter = itertools.count() self._workers: Dict[int, Dict[str, Any]] = {} self._idle_pids: List[int] = [] @@ -188,6 +194,7 @@ def __init__( self._spawn_worker(gpu_id) for _ in range(self._cpu_worker_count): self._spawn_worker(None) + self._wait_for_workers_ready(n_workers) @property def gpu_worker_devices(self) -> List[int]: @@ -217,6 +224,7 @@ def _spawn_worker(self, physical_gpu_id: Optional[int]) -> None: self._result_q, self._dispatch, self._threads_per_worker, + getattr(self, "_ready_q", None), visible_id, ), daemon=True, @@ -238,7 +246,28 @@ def _spawn_worker(self, physical_gpu_id: Optional[int]) -> None: "job_q": job_q, "gpu_id": physical_gpu_id, } - self._mark_idle(p.pid) + + def _wait_for_workers_ready(self, count: int) -> None: + deadline = time.monotonic() + float( + os.environ.get("SCURO_WORKER_START_TIMEOUT_SEC", "300") + ) + ready = set() + while len(ready) < count: + for pid, worker in self._workers.items(): + if pid not in ready and not worker["process"].is_alive(): + raise RuntimeError(f"SCURO worker {pid} exited during startup") + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError( + "Timed out waiting for SCURO workers to become ready" + ) + try: + pid = self._ready_q.get(timeout=min(0.1, remaining)) + except queue.Empty: + continue + if pid in self._workers and pid not in ready: + ready.add(pid) + self._mark_idle(pid) def _mark_idle(self, pid: int) -> None: worker = self._workers.get(pid) @@ -366,6 +395,7 @@ def _replace_dead_worker(self, pid: int) -> Optional[_JobResult]: self._running.pop(failed_job_id, None) self._spawn_worker(physical_gpu_id) + self._wait_for_workers_ready(1) if failed_job_id is None: return None @@ -397,6 +427,11 @@ def shutdown(self) -> None: w["job_q"].join_thread() except Exception: pass + try: + self._ready_q.close() + self._ready_q.join_thread() + except Exception: + pass try: self._result_q.close() self._result_q.join_thread() diff --git a/src/main/python/systemds/scuro/modality/type.py b/src/main/python/systemds/scuro/modality/type.py index a648044562a..dc9856f9751 100644 --- a/src/main/python/systemds/scuro/modality/type.py +++ b/src/main/python/systemds/scuro/modality/type.py @@ -26,7 +26,6 @@ calculate_new_frequency, create_timestamps, ) -import torch # TODO: needs a way to define if data comes from a dataset with multiple instances or is like a streaming scenario where we only have one instance @@ -367,7 +366,7 @@ def get_data_layout(cls, data, data_is_single_instance): elif ( isinstance(data, list) or isinstance(data, np.ndarray) - or isinstance(data, torch.Tensor) + or type(data).__module__.split(".", 1)[0] == "torch" ): return DataLayout.NESTED_LEVEL diff --git a/src/main/python/systemds/scuro/models/training_stats.py b/src/main/python/systemds/scuro/models/training_stats.py index bee9d20f7a2..0124bc88da8 100644 --- a/src/main/python/systemds/scuro/models/training_stats.py +++ b/src/main/python/systemds/scuro/models/training_stats.py @@ -63,7 +63,7 @@ class TrainingStats: FIELDS = ( "fits", - "epochs_run", + "loader_passes", "train_steps", "train_samples", "setup_s", @@ -88,7 +88,7 @@ def setup(self): return _Phase(self, "setup_s") def training(self): - """The epoch/step loop only.""" + """The step loop only.""" return _Phase(self, "train_loop_s") def evaluation(self): @@ -101,8 +101,8 @@ def count_step(self, batch_samples=0): self.train_steps += 1 self.train_samples += int(batch_samples) - def count_epoch(self): - self.epochs_run += 1 + def count_pass(self): + self.loader_passes += 1 def count_eval_batch(self, batch_samples=0): self.eval_batches += 1 @@ -167,7 +167,7 @@ def as_dict(self): return { "work_folds": self.folds, "work_fits": totals["fits"], - "work_epochs_run": totals["epochs_run"], + "work_loader_passes": totals["loader_passes"], "work_train_steps": totals["train_steps"], "work_train_samples": totals["train_samples"], "work_setup_s": totals["setup_s"], From 6872518c978b0f0ddddccb5f5b8df6b1006aa589 Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Sat, 12 Sep 2026 21:42:17 +0200 Subject: [PATCH 12/15] add list of fields to json loader --- .../systemds/scuro/dataloader/json_loader.py | 23 +++++++++++-------- 1 file changed, 14 insertions(+), 9 deletions(-) diff --git a/src/main/python/systemds/scuro/dataloader/json_loader.py b/src/main/python/systemds/scuro/dataloader/json_loader.py index aaf11ef6d84..05f1330942b 100644 --- a/src/main/python/systemds/scuro/dataloader/json_loader.py +++ b/src/main/python/systemds/scuro/dataloader/json_loader.py @@ -43,7 +43,7 @@ def __init__( self, source_path: str, indices: List[str], - field: str, # TODO: make this a list so it is easier to get multiple fields from a json file. (i.e. Mustard: context + sentence) + field: Union[str, List[str]], data_type: Union[np.dtype, str] = str, chunk_size: Optional[int] = None, ext: str = ".json", @@ -59,6 +59,14 @@ def __init__( lambda: self.get_stats(source_path), ) + def _get_text(self, record: dict) -> str: + fields = [self.field] if isinstance(self.field, str) else self.field + parts = [] + for field in fields: + text = record[field] + parts.append(" ".join(text) if isinstance(text, list) else text) + return " ".join(part for part in parts if part) + def extract(self, file: str, index: Optional[Union[str, List[str]]] = None): self.file_sanity_check(file) with open(file) as f: @@ -68,11 +76,10 @@ def extract(self, file: str, index: Optional[Union[str, List[str]]] = None): index = [index] for idx in index: try: - text = json_file[idx][self.field] + text = self._get_text(json_file[idx]) except: - text = json_file[self.field] + text = self._get_text(json_file) - text = " ".join(text) if isinstance(text, list) else text self.data.append(text) self.metadata.append( self.modality_type.create_metadata(len(text), text) | json_file[idx] @@ -90,11 +97,10 @@ def get_stats(self, source_path: str): json_file = json.load(f) for id in self.indices: try: - text = json_file[id][self.field] + text = self._get_text(json_file[id]) except: - text = json_file[self.field] + text = self._get_text(json_file) - text = " ".join(text) if isinstance(text, list) else text num_instances += 1 max_length = max(max_length, len(text)) avg_length += len(text) @@ -102,9 +108,8 @@ def get_stats(self, source_path: str): for id in self.indices: with open(os.path.join(source_path, f"{id}{self._ext}")) as f: json_file = json.load(f) - text = json_file[self.field] + text = self._get_text(json_file) - text = " ".join(text) if isinstance(text, list) else text num_instances += 1 max_length = max(max_length, len(text)) # number of characters max_words = max(max_words, len(text.split())) From f7a54b9e0735cd4f1694560c84a5fc13634ca313 Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Sun, 13 Sep 2026 18:34:53 +0200 Subject: [PATCH 13/15] limit dimentionality of fusion ops --- .../systemds/scuro/dataloader/json_loader.py | 9 +-- .../scuro/drsearch/multimodal_ga_optimizer.py | 58 ++++++++++++++- .../scuro/drsearch/split_aware_evaluator.py | 52 +++++++++++++ src/main/python/systemds/scuro/models/dnn.py | 73 +++++++++++++++++++ .../systemds/scuro/representations/fusion.py | 36 ++++++++- .../systemds/scuro/representations/lstm.py | 24 +++++- .../systemds/scuro/representations/max.py | 18 +++++ .../multimodal_attention_fusion.py | 19 ++++- 8 files changed, 272 insertions(+), 17 deletions(-) create mode 100644 src/main/python/systemds/scuro/models/dnn.py diff --git a/src/main/python/systemds/scuro/dataloader/json_loader.py b/src/main/python/systemds/scuro/dataloader/json_loader.py index 05f1330942b..d2802894ecb 100644 --- a/src/main/python/systemds/scuro/dataloader/json_loader.py +++ b/src/main/python/systemds/scuro/dataloader/json_loader.py @@ -75,14 +75,13 @@ def extract(self, file: str, index: Optional[Union[str, List[str]]] = None): if isinstance(index, str): index = [index] for idx in index: - try: - text = self._get_text(json_file[idx]) - except: - text = self._get_text(json_file) + # One file keyed by instance id, or one file per instance. + record = json_file[idx] if idx in json_file else json_file + text = self._get_text(record) self.data.append(text) self.metadata.append( - self.modality_type.create_metadata(len(text), text) | json_file[idx] + self.modality_type.create_metadata(len(text), text) | record ) def get_stats(self, source_path: str): 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 0c26c9ebcb4..79a04ca262c 100644 --- a/src/main/python/systemds/scuro/drsearch/multimodal_ga_optimizer.py +++ b/src/main/python/systemds/scuro/drsearch/multimodal_ga_optimizer.py @@ -61,6 +61,7 @@ from systemds.scuro.representations.aggregated_representation import ( AggregatedRepresentation, ) +from systemds.scuro.representations.fusion import FusionInputTooLargeError from systemds.scuro.representations.representation import RepresentationStats from systemds.scuro.utils.schema_helpers import get_shape @@ -289,6 +290,8 @@ def _dispatch_genome_evaluation(payload, _gpu_id): if fitness is None: fitness = _failure_fitness(objective_specs) return fitness, result + except FusionInputTooLargeError as error: + return _failure_fitness(objective_specs), {"_oversized": str(error)} except Exception: return _failure_fitness(objective_specs), { "_evaluation_error": traceback.format_exc() @@ -729,6 +732,13 @@ def _evaluate_individuals_parallel( ind, dag, sig = jobs.pop(jr.job_id) if jr.ok: fitness, payload = jr.value + if isinstance(payload, dict) and "_oversized" in payload: + ind.fitness.values = self._reject_oversized( + task_name, None, payload["_oversized"], sig + ) + for follower in pending_followers.pop(sig, []): + follower.fitness.values = ind.fitness.values + continue error = ( payload.get("_evaluation_error") if isinstance(payload, dict) @@ -1026,7 +1036,45 @@ def walk(subtree: Tree, path: str): return None return int(np.prod(root.output_shape)) + def _operator_input_violation( + self, task_name: str, genome: DagGenome + ) -> Optional[str]: + """First fusion node whose operands its operator cannot take, if any. + + Checked per node, since an operator's limit applies to its own inputs + (LSTM and AttentionFusion build a model over the operand width), not + to the width of the final fused representation. + """ + if task_name is None: + return None + leaf_stats = [self._leaf_stats(task_name, leaf) for leaf in genome.leaves] + + def walk(subtree: Tree, path: str): + if isinstance(subtree, int): + return leaf_stats[subtree], None + left, reason = walk(subtree[0], path + "L") + if reason is not None: + return None, reason + right, reason = walk(subtree[1], path + "R") + if reason is not None: + return None, reason + if left is None or right is None: + return None, None + operator = genome.fusion_ops[path]() + reason = operator.input_size_violation([left, right]) + if reason is not None: + return None, reason + return operator.get_output_stats([left, right]), None + + try: + return walk(genome.tree, "")[1] + except Exception: + return None + def _oversized_reason(self, task_name: str, genome: DagGenome) -> Optional[str]: + reason = self._operator_input_violation(task_name, genome) + if reason is not None: + return f"{reason}; candidate rejected before execution" if not self.max_feature_dim: return None predicted = self._predicted_feature_dim(task_name, genome) @@ -1039,7 +1087,7 @@ def _oversized_reason(self, task_name: str, genome: DagGenome) -> Optional[str]: ) def _reject_oversized( - self, task_name: str, genome: DagGenome, reason: str + self, task_name: str, genome: Optional[DagGenome], reason: str, signature=None ) -> Tuple[float, ...]: self.oversized_candidates[task_name] = ( self.oversized_candidates.get(task_name, 0) + 1 @@ -1047,9 +1095,9 @@ def _reject_oversized( if self.debug: print(f"[GA] task={task_name} skipped candidate: {reason}") fitness = _failure_fitness(self.objective_specs) - self._fitness_cache.setdefault(task_name, {})[ - self._genome_signature(genome) - ] = fitness + if signature is None: + signature = self._genome_signature(genome) + self._fitness_cache.setdefault(task_name, {})[signature] = fitness return fitness def _evaluate_genome(self, genome: DagGenome, task: Task) -> Tuple[float, ...]: @@ -1075,6 +1123,8 @@ def _evaluate_genome(self, genome: DagGenome, task: Task) -> Tuple[float, ...]: error = None if fitness is None: fitness = _failure_fitness(self.objective_specs) + except FusionInputTooLargeError as too_large: + return self._reject_oversized(task_name, genome, str(too_large), sig) except Exception: fitness = _failure_fitness(self.objective_specs) payload, error = None, traceback.format_exc() 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 39c6e1be21c..ebeecc733d6 100644 --- a/src/main/python/systemds/scuro/drsearch/split_aware_evaluator.py +++ b/src/main/python/systemds/scuro/drsearch/split_aware_evaluator.py @@ -132,6 +132,57 @@ def _apply_stateless(self, operation, inputs, node): return self._wrap(transformed, inputs.instance_ids) + @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 = [] + for item in result: + data = item.modality.data + if isinstance(data, np.ndarray): + array = data + elif isinstance(data, list) and data and all( + isinstance(d, np.ndarray) for d in data + ): + if len({d.shape for d in data}) != 1: + return result + array = None + else: + return result + if array is not None and array.dtype == object: + return result + arrays.append(array) + shapes = [ + a.shape[1:] if a is not None else item.modality.data[0].shape + for a, item in zip(arrays, result) + ] + if len({len(shape) for shape in shapes}) != 1 or not shapes[0]: + return result + target = tuple(max(dims) for dims in zip(*shapes)) + for array, shape, item in zip(arrays, shapes, result): + if shape == target: + continue + modality = item.modality + source = array if array is not None else np.stack(modality.data) + padded = np.zeros((source.shape[0], *target), dtype=source.dtype) + padded[(slice(None),) + tuple(slice(0, dim) for dim in shape)] = source + modality.data = padded if array is not None else list(padded) + for metadata in modality.metadata or []: + layout = metadata.get("data_layout") if isinstance(metadata, dict) else None + if isinstance(layout, dict) and len(tuple(layout.get("shape") or ())) == len(target): + layout["shape"] = target + return result + @staticmethod def _leaf_input(partition, node): if len(partition) == 1: @@ -204,6 +255,7 @@ def execute_node(node_id): self._apply_stateless(operation, partition, node) for partition in input_partitions ) + result = self._align_partition_shapes(result) cache[node_id] = result return result diff --git a/src/main/python/systemds/scuro/models/dnn.py b/src/main/python/systemds/scuro/models/dnn.py new file mode 100644 index 00000000000..71d9b4785aa --- /dev/null +++ b/src/main/python/systemds/scuro/models/dnn.py @@ -0,0 +1,73 @@ +# ------------------------------------------------------------- +# +# 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. +# +# ------------------------------------------------------------- +from systemds.scuro.models.model import Model + + +class DNN(Model): + """Base class for the neural task models, carrying their training budget. + + The budget is a number of optimizer steps -- mini-batches -- and not a + number of epochs. An epoch is not a comparable unit of work: one epoch is + two mini-batches on a StressID fold and over two hundred on an MM-IMDb + fold, so an epoch count says nothing on its own about how much optimisation + a model got, and a budget expressed in epochs cannot be compared across + datasets at all. A step is the same unit everywhere (at a fixed batch + size), so `num_steps` means the same thing on every dataset, and a budget + smaller than one pass over the data is expressible rather than rounded up. + + `num_steps` is None only for the older models that still run their own + epoch loop and do not use `train_batches`. + """ + + def __init__(self, name, num_steps=None): + super().__init__(name) + self.num_steps = num_steps + + def drop_last_batch(self, n_samples, batch_size): + """Whether the loader should drop an uneven final batch. + + Every step has to be the same amount of work, or the budget stops being + a unit of work at all; a short tail batch also breaks BatchNorm when it + holds a single sample. Dropping it is only safe while at least one full + batch remains, which is not a given on the small folds (StressID's are + a handful of batches wide). + """ + return int(n_samples) >= int(batch_size) * 2 + + def train_batches(self, loader, stats=None): + """Yield mini-batches until the step budget is spent. + + The loader is re-iterated as often as the budget needs, so a budget of + more than one pass reshuffles between passes exactly as an epoch loop + would, and a budget of less than one pass simply stops part-way through + the first one. + """ + remaining = max(0, int(self.num_steps or 0)) + if not len(loader): + return + while remaining > 0: + if stats is not None: + stats.count_pass() + for batch in loader: + yield batch + remaining -= 1 + if remaining == 0: + return diff --git a/src/main/python/systemds/scuro/representations/fusion.py b/src/main/python/systemds/scuro/representations/fusion.py index c3c83230b75..5cbe128cae0 100644 --- a/src/main/python/systemds/scuro/representations/fusion.py +++ b/src/main/python/systemds/scuro/representations/fusion.py @@ -19,7 +19,7 @@ # # ------------------------------------------------------------- import copy -from typing import List +from typing import List, Optional import numpy as np @@ -46,6 +46,14 @@ DEFAULT_MAX_ALIGNMENT_PAD_RATIO = None +class FusionInputTooLargeError(ValueError): + """The operands are too wide for this fusion operator's model. + + Raised before any training, so a search can discard the candidate as + oversized instead of counting it as a failed evaluation. + """ + + class Fusion(Representation): def __init__(self, name, parameters=None): """ @@ -84,6 +92,32 @@ def _pre_aggregated_stats(self, stats: RepresentationStats) -> RepresentationSta def _fusion_input_stats(self, input_stats) -> List[RepresentationStats]: return [self._pre_aggregated_stats(s) for s in self._as_stats_list(input_stats)] + def input_size_violation(self, input_stats) -> Optional[str]: + """Why these operands are too wide for this operator, or None. + + Predicted from shapes alone so that a search can reject a candidate + before executing it. Operators without an input-width limit accept + everything. + """ + return None + + def _operand_widths(self, input_stats) -> Optional[List[int]]: + """Per-operand feature width as _prepare_modalities hands it to execute. + + The width is the last axis of each operand after pre-aggregation; + operators that need alignment zero-pad every operand to the widest. + None if a shape is unknown or has more than two axes. + """ + widths = [] + for stats in self._fusion_input_stats(input_stats): + shape = tuple(stats.output_shape or ()) + if not stats.output_shape_is_known or len(shape) not in (1, 2): + return None + widths.append(int(shape[-1])) + if self.needs_alignment and widths: + widths = [max(widths)] * len(widths) + return widths + def _raw_input_bytes(self, input_stats) -> int: return sum( stats_bytes(s, quantile=0.95) for s in self._as_stats_list(input_stats) diff --git a/src/main/python/systemds/scuro/representations/lstm.py b/src/main/python/systemds/scuro/representations/lstm.py index 2086b199446..a8a5c2a0519 100644 --- a/src/main/python/systemds/scuro/representations/lstm.py +++ b/src/main/python/systemds/scuro/representations/lstm.py @@ -24,12 +24,12 @@ import torch from torch import nn from torch.utils.data import DataLoader, TensorDataset -from typing import List, Dict, Any +from typing import List, Dict, Any, Optional from systemds.scuro.utils.static_variables import get_device import numpy as np from systemds.scuro.modality.modality import Modality -from systemds.scuro.representations.fusion import Fusion +from systemds.scuro.representations.fusion import Fusion, FusionInputTooLargeError from systemds.scuro.representations.representation import ( FitRequirement, RepresentationStats, @@ -187,7 +187,7 @@ def _prepare_data(self, modalities: List[Modality]) -> np.ndarray: input_dim = sum(data.shape[2] for data in processed_modalities) if self.max_input_dim is not None and input_dim > self.max_input_dim: shapes = ", ".join(self._shape_description(m) for m in modalities) - raise ValueError( + raise FusionInputTooLargeError( f"LSTM input feature dimension {input_dim:,} exceeds " f"the configured maximum {self.max_input_dim:,}. Inputs: {shapes}" ) @@ -226,13 +226,29 @@ def _validate_model_size(self, input_dim: int, num_classes: int) -> None: self.max_parameter_bytes is not None and parameter_bytes > self.max_parameter_bytes ): - raise ValueError( + raise FusionInputTooLargeError( f"LSTM parameters require approximately " f"{parameter_bytes / 1024**3:.2f} GiB, exceeding the configured " f"maximum {self.max_parameter_bytes / 1024**3:.2f} GiB " f"(input_dim={input_dim:,}, width={self.width}, depth={self.depth})." ) + def input_size_violation(self, input_stats) -> Optional[str]: + widths = self._operand_widths(input_stats) + if widths is None: + return None + input_dim = sum(widths) + if self.max_input_dim is not None and input_dim > self.max_input_dim: + return ( + f"LSTM input feature dimension {input_dim:,} would exceed the " + f"configured maximum {self.max_input_dim:,} (operand widths {widths})" + ) + try: + self._validate_model_size(input_dim, 2) + except FusionInputTooLargeError as error: + return str(error) + return None + def get_output_stats(self, input_stats) -> RepresentationStats: stats_list = self._as_stats_list(input_stats) num_instances = max((stats.num_instances for stats in stats_list), default=0) diff --git a/src/main/python/systemds/scuro/representations/max.py b/src/main/python/systemds/scuro/representations/max.py index 2dadc497c78..2e0df4c464e 100644 --- a/src/main/python/systemds/scuro/representations/max.py +++ b/src/main/python/systemds/scuro/representations/max.py @@ -23,6 +23,7 @@ import numpy as np from systemds.scuro.modality.modality import Modality +from systemds.scuro.representations.representation import RepresentationStats from systemds.scuro.representations.fusion import Fusion from systemds.scuro.drsearch.operator_registry import register_fusion_operator @@ -46,3 +47,20 @@ def execute( ): fused_data = np.maximum.reduce([m.data for m in modalities]) return fused_data + + def get_output_stats(self, input_stats_list) -> RepresentationStats: + # Elementwise over operands aligned to the widest, as in Average. + stats_list = self._fusion_input_stats(input_stats_list) + if not stats_list: + return RepresentationStats(0, (0,)) + + num_instances = max(s.num_instances for s in stats_list) + rank = len(stats_list[0].output_shape) + if rank > 0 and all(len(s.output_shape) == rank for s in stats_list): + output_shape = tuple( + max(s.output_shape[d] for s in stats_list) for d in range(rank) + ) + else: + output_shape = max(stats_list, key=self._stats_num_elements).output_shape + output_shape_is_known = all(s.output_shape_is_known for s in stats_list) + return RepresentationStats(num_instances, output_shape, output_shape_is_known) 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 162892a7faa..aecd2abbc33 100644 --- a/src/main/python/systemds/scuro/representations/multimodal_attention_fusion.py +++ b/src/main/python/systemds/scuro/representations/multimodal_attention_fusion.py @@ -28,7 +28,7 @@ import numpy as np from systemds.scuro.drsearch.operator_registry import register_fusion_operator from systemds.scuro.modality.modality import Modality -from systemds.scuro.representations.fusion import Fusion +from systemds.scuro.representations.fusion import Fusion, FusionInputTooLargeError from systemds.scuro.representations.representation import ( FitRequirement, RepresentationStats, @@ -175,7 +175,7 @@ def _validate_model_size( ) -> None: total_input_dim = sum(input_dimensions.values()) if self.max_input_dim is not None and total_input_dim > self.max_input_dim: - raise ValueError( + raise FusionInputTooLargeError( f"AttentionFusion input feature dimension {total_input_dim:,} exceeds " f"the configured maximum {self.max_input_dim:,}; per-modality " f"dimensions: {input_dimensions}. Apply aggregation or dimensionality " @@ -189,12 +189,25 @@ def _validate_model_size( self.max_parameter_bytes is not None and parameter_bytes > self.max_parameter_bytes ): - raise ValueError( + raise FusionInputTooLargeError( "AttentionFusion parameters require approximately " f"{parameter_bytes / 1024**3:.2f} GiB, exceeding the configured " f"maximum {self.max_parameter_bytes / 1024**3:.2f} GiB." ) + def input_size_violation(self, input_stats) -> Optional[str]: + widths = self._operand_widths(input_stats) + if widths is None: + return None + dimensions = {f"modality_{i}": width for i, width in enumerate(widths)} + try: + # Sequence length and class count are not known from shapes; the + # width limit is what wide operands hit first. + self._validate_model_size(dimensions, 1, 2) + except FusionInputTooLargeError as error: + return str(error) + return None + def get_output_stats(self, input_stats) -> RepresentationStats: stats_list = self._as_stats_list(input_stats) num_instances = max((stats.num_instances for stats in stats_list), default=0) From 871d117c627992268120f4dac66bc049e67c2042 Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Mon, 28 Sep 2026 11:58:30 +0200 Subject: [PATCH 14/15] add batched video execution --- .../systemds/scuro/dataloader/video_loader.py | 18 +- .../scuro/drsearch/hyperparameter_tuner.py | 214 +++++++++++-- .../scuro/drsearch/multimodal_optimizer.py | 7 +- .../systemds/scuro/drsearch/node_executor.py | 6 +- .../systemds/scuro/drsearch/node_scheduler.py | 46 ++- .../scuro/drsearch/operator_registry.py | 2 +- .../scuro/drsearch/representation_dag.py | 17 + .../scuro/drsearch/split_aware_evaluator.py | 40 ++- .../scuro/drsearch/trial_result_cache.py | 212 ++++++++++++ .../scuro/drsearch/unimodal_optimizer.py | 45 ++- .../python/systemds/scuro/modality/indexed.py | 9 +- .../scuro/modality/unimodal_modality.py | 18 +- .../systemds/scuro/representations/bow.py | 4 + .../representations/covarep_audio_features.py | 19 +- .../systemds/scuro/representations/elmo.py | 5 + .../scuro/representations/openface.py | 277 ++++++++++++---- .../scuro/representations/spectrogram.py | 3 + .../representations/swin_video_transformer.py | 101 +++--- .../text_context_with_indices.py | 3 + .../systemds/scuro/representations/utils.py | 30 +- .../scuro/representations/video_batching.py | 301 ++++++++++++++++++ .../systemds/scuro/representations/x3d.py | 162 +++++----- .../systemds/scuro/utils/memory_utility.py | 18 +- 23 files changed, 1263 insertions(+), 294 deletions(-) create mode 100644 src/main/python/systemds/scuro/drsearch/trial_result_cache.py create mode 100644 src/main/python/systemds/scuro/representations/video_batching.py diff --git a/src/main/python/systemds/scuro/dataloader/video_loader.py b/src/main/python/systemds/scuro/dataloader/video_loader.py index 55ac34dc934..1f499a09b3e 100644 --- a/src/main/python/systemds/scuro/dataloader/video_loader.py +++ b/src/main/python/systemds/scuro/dataloader/video_loader.py @@ -40,6 +40,7 @@ class VideoStats: num_instances: int num_total_instances: int shape_variance: float = 0.0 + lengths: Optional[Tuple[int, ...]] = None @property def output_shape(self): @@ -84,9 +85,15 @@ def __init__( files = self.get_file_names(self.indices) self.stats, self._all_metadata = self.cached_stats( files, - {"dtype": str(data_type), "fps": fps, "target_size": self.target_size}, + { + "dtype": str(data_type), + "fps": fps, + "target_size": self.target_size, + "frame_count": "decodable", + }, lambda: (self.get_stats(source_path), self._all_metadata), ) + self.stats.lengths = tuple(int(md["length"]) for md in self._all_metadata) def load(self): if self.chunk_size: @@ -106,6 +113,13 @@ def _frame_interval(self, source_fps: float) -> int: return max(1, int(round(source_fps / self.fps))) return 1 + @staticmethod + def _decodable_frame_count(cap) -> int: + count = 0 + while cap.grab(): + count += 1 + return count + def _stored_length(self, source_length: int, source_fps: float) -> int: interval = self._frame_interval(source_fps) return int(math.ceil(source_length / interval)) if source_length > 0 else 0 @@ -193,10 +207,10 @@ def get_stats(self, source_path: str): if not cap.isOpened(): raise ValueError(f"Could not read video at path: {file}") try: - source_length = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) source_fps = cap.get(cv2.CAP_PROP_FPS) width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + source_length = self._decodable_frame_count(cap) finally: cap.release() diff --git a/src/main/python/systemds/scuro/drsearch/hyperparameter_tuner.py b/src/main/python/systemds/scuro/drsearch/hyperparameter_tuner.py index 43bedeeab0c..2871c60340c 100644 --- a/src/main/python/systemds/scuro/drsearch/hyperparameter_tuner.py +++ b/src/main/python/systemds/scuro/drsearch/hyperparameter_tuner.py @@ -40,7 +40,21 @@ ) from systemds.scuro.drsearch.task import PerformanceMeasure import pickle +import torch from systemds.scuro.utils.checkpointing import CheckpointManager +from systemds.scuro.drsearch.node_executor import _place_on_device +from systemds.scuro.drsearch.trial_result_cache import ( + TrialNodeCache, + TrialResultCache, + trial_node_signatures, +) + +# torch.OutOfMemoryError only exists on newer torch; RuntimeError covers the rest. +_CUDA_OOM_ERRORS = tuple( + e + for e in (getattr(torch, "OutOfMemoryError", None), torch.cuda.OutOfMemoryError) + if e is not None +) or (RuntimeError,) def _wandb_safe_tag(s: str, max_len: int = 64) -> str: @@ -337,7 +351,8 @@ def __init__( save_results: bool = False, debug: bool = False, checkpoint_every: Optional[int] = None, - resume: bool = True, + checkpoint_dir: Optional[str] = None, + resume: bool = False, random_state: int = 42, exhaustive_threshold: int = 256, local_search_patience: int = 3, @@ -372,7 +387,7 @@ def __init__( self.local_search_patience = max(1, local_search_patience) self._rng = random.Random(self.random_state) self._checkpoint_manager = CheckpointManager( - os.getcwd(), + checkpoint_dir or os.getcwd(), "hyperparam_checkpoint_", checkpoint_every=self.checkpoint_every, resume=self.resume, @@ -445,15 +460,14 @@ def resume_from_checkpoint(self): self.optimization_results.results = results def tune_unimodal_representations(self, max_eval_per_rep: Optional[int] = None): - # self.resume_from_checkpoint() + self.resume_from_checkpoint() for task in self.tasks: reps = self.k_best_representations[task.model.name] - skip_remaining = 0 - # skip_remaining = self._checkpoint_manager.skip_remaining_by_key.get( - # task.model.name, 0 - # ) - # if skip_remaining >= len(reps): - # continue + skip_remaining = self._checkpoint_manager.skip_remaining_by_key.get( + task.model.name, 0 + ) + if skip_remaining >= len(reps): + continue chunk_size = self.checkpoint_every or len(reps) for start_idx in range(skip_remaining, len(reps), chunk_size): @@ -486,6 +500,24 @@ def tune_unimodal_representations(self, max_eval_per_rep: Optional[int] = None): def tune_dag_representation( self, dag, root_node_id, task, max_evals=None, mm_opt=False + ): + self._trial_cache = TrialResultCache.from_env() + try: + return self._tune_dag_representation( + dag, root_node_id, task, max_evals, mm_opt + ) + finally: + if self._trial_cache is not None: + self.logger.info( + f"trial cache for {dag.get_represntation_names()} / " + f"{task.model.name}: {self._trial_cache.summary()}" + ) + print(f"[hpo-cache] {task.model.name}: {self._trial_cache.summary()}") + self._trial_cache.clear() + self._trial_cache = None + + def _tune_dag_representation( + self, dag, root_node_id, task, max_evals=None, mm_opt=False ): hyperparams = {} reps = [] @@ -502,7 +534,7 @@ def visit_node(node_id): visit_node(input_id) visited.add(node_id) if node.operation is not None: - params = self.__get_params_for_node(node) + params = self.__get_params_for_node(node, dag) if params: hyperparams[node_id] = params reps.append(node.operation) @@ -609,7 +641,22 @@ def get_score(result): return best_result - def __get_params_for_node(self, node: RepresentationNode) -> Dict[str, Any]: + def _leaf_input_stats(self, node: RepresentationNode, dag) -> Any: + if dag is None or len(node.inputs) != 1: + return None + try: + input_node = dag.get_node_by_id(node.inputs[0]) + except Exception: + return None + if input_node is None or input_node.operation is not None: + return None + modalities = self.get_modalities_by_id([input_node.modality_id]) + get_stats = getattr(modalities[0], "get_stats", None) if modalities else None + return get_stats() if callable(get_stats) else None + + def __get_params_for_node( + self, node: RepresentationNode, dag=None + ) -> Dict[str, Any]: try: if node.parameters: op = node.operation(params=node.parameters) @@ -622,6 +669,13 @@ def __get_params_for_node(self, node: RepresentationNode) -> Dict[str, Any]: return None params = copy.deepcopy(op.parameters) + input_stats = self._leaf_input_stats(node, dag) + if input_stats is not None: + for name, values in list(params.items()): + if isinstance(values, list): + params[name] = ( + op.filter_parameter_domain(name, values, input_stats) or values + ) if node.parameters: if inspect.isclass(node.parameters.get("aggregation_function")): params["aggregation_function"] = node.parameters["aggregation_function"] @@ -1008,30 +1062,128 @@ def _get_cached_modalities_for_task(self, task, modality_ids): ) return cached_modalities + def _trial_node_cache(self, dag, params): + store = getattr(self, "_trial_cache", None) + if store is None: + return None + return TrialNodeCache(store, trial_node_signatures(dag, params)) + + def _evaluate_dag_once( + self, + dag, + params, + node_order, + modality_ids, + task, + modalities_override=None, + on_cpu=False, + node_cache=None, + ): + dag_copy = copy.deepcopy(dag) + for node_id in node_order: + node = dag_copy.get_node_by_id(node_id) + if node.operation is not None: + node.parameters = _apply_trial_params_to_node(node, params) + if on_cpu and node.operation is not None: + _place_on_device(node.operation, None) + + modalities = ( + modalities_override + if modalities_override is not None + else self.get_modalities_by_id(modality_ids) + ) + if on_cpu: + _place_on_device(getattr(task, "model", task), None) + + if dag_requires_fitting(dag_copy): + evaluation = SplitAwareDAGEvaluator( + dag_copy, modalities, task, node_cache=node_cache + ).evaluate() + return params, evaluation.scores + modified_modality = dag_copy.execute(modalities, task, node_cache=node_cache) + return params, task.run(modified_modality.data) + + def _evaluate_dag_trial( + self, + dag, + params, + node_order, + modality_ids, + task, + modalities_override=None, + on_cpu=False, + ): + node_cache = self._trial_node_cache(dag, params) + if node_cache is None: + return self._evaluate_dag_once( + dag, + params, + node_order, + modality_ids, + task, + modalities_override, + on_cpu=on_cpu, + ) + try: + return self._evaluate_dag_once( + dag, + params, + node_order, + modality_ids, + task, + modalities_override, + on_cpu=on_cpu, + node_cache=node_cache, + ) + except _CUDA_OOM_ERRORS: + raise + except Exception as error: + if node_cache.hits == 0: + raise + self.logger.warning( + f"Trial with {node_cache.hits} cached result(s) failed ({error}); " + "repeating it without the cache" + ) + return self._evaluate_dag_once( + dag, + params, + node_order, + modality_ids, + task, + modalities_override, + on_cpu=on_cpu, + ) + def evaluate_dag_config( self, dag, params, node_order, modality_ids, task, modalities_override=None ): try: - dag_copy = copy.deepcopy(dag) - for node_id in node_order: - node = dag_copy.get_node_by_id(node_id) - if node.operation is not None: - node.parameters = _apply_trial_params_to_node(node, params) - - modalities = ( - modalities_override - if modalities_override is not None - else self.get_modalities_by_id(modality_ids) - ) - if dag_requires_fitting(dag_copy): - evaluation = SplitAwareDAGEvaluator( - dag_copy, modalities, task - ).evaluate() - return params, evaluation.scores - modified_modality = dag_copy.execute(modalities, task) - score = task.run(modified_modality.data) - - return params, score + try: + return self._evaluate_dag_trial( + dag, params, node_order, modality_ids, task, modalities_override + ) + except _CUDA_OOM_ERRORS as oom: + # A trial that does not fit the card used to be recorded as nan, which + # silently shrinks the tuning the stage is supposed to perform. Fall back + # to the CPU for this trial, the same way the node executor does when no + # GPU has room. Same configuration, same budget, different device. + self.logger.warning( + f"Trial ran out of GPU memory ({oom}); retrying on the CPU" + ) + if torch.cuda.is_available(): + torch.cuda.empty_cache() + result = self._evaluate_dag_trial( + dag, + params, + node_order, + modality_ids, + task, + modalities_override, + on_cpu=True, + ) + if torch.cuda.is_available(): + torch.cuda.empty_cache() + return result except Exception as e: import traceback diff --git a/src/main/python/systemds/scuro/drsearch/multimodal_optimizer.py b/src/main/python/systemds/scuro/drsearch/multimodal_optimizer.py index c146a6485a2..54cddfd142b 100644 --- a/src/main/python/systemds/scuro/drsearch/multimodal_optimizer.py +++ b/src/main/python/systemds/scuro/drsearch/multimodal_optimizer.py @@ -24,7 +24,7 @@ import itertools import threading from dataclasses import dataclass -from typing import List, Dict, Any, Generator +from typing import List, Dict, Any, Generator, Optional from systemds.scuro.drsearch.task import Task, PerformanceMeasure from systemds.scuro.drsearch.representation_dag import ( RepresentationDag, @@ -126,7 +126,8 @@ def __init__( max_modalities: int = None, metric: str = "accuracy", checkpoint_every: int = None, - resume: bool = True, + checkpoint_dir: Optional[str] = None, + resume: bool = False, ): self.modalities = modalities self.tasks = tasks @@ -146,7 +147,7 @@ def __init__( self.checkpoint_every = checkpoint_every self.resume = resume self._checkpoint_manager = CheckpointManager( - os.getcwd(), + checkpoint_dir or os.getcwd(), "multimodal_checkpoint_", checkpoint_every=self.checkpoint_every, resume=self.resume, diff --git a/src/main/python/systemds/scuro/drsearch/node_executor.py b/src/main/python/systemds/scuro/drsearch/node_executor.py index b30f0a5df07..3219b37f279 100644 --- a/src/main/python/systemds/scuro/drsearch/node_executor.py +++ b/src/main/python/systemds/scuro/drsearch/node_executor.py @@ -444,7 +444,11 @@ def __init__( search_start: Optional[float] = None, ): self.enable_checkpointing = enable_checkpointing - available_total_cpu = cpu_memory_budget_bytes() + available_total_cpu = cpu_memory_budget_bytes( + getattr(getattr(modalities[0], "modality_type", None), "name", None) + if modalities + else None + ) self.dags = dags self.scheduler = MemoryAwareNodeScheduler( dags, modalities, tasks, available_total_cpu diff --git a/src/main/python/systemds/scuro/drsearch/node_scheduler.py b/src/main/python/systemds/scuro/drsearch/node_scheduler.py index 44f7dfef9ff..78ef5b953fa 100644 --- a/src/main/python/systemds/scuro/drsearch/node_scheduler.py +++ b/src/main/python/systemds/scuro/drsearch/node_scheduler.py @@ -40,6 +40,10 @@ _MAX_GPU_SCHEDULE_ATTEMPTS = int(os.environ.get("SCURO_MAX_GPU_SCHEDULE_ATTEMPTS", "3")) +def _exclusive_over_budget_enabled() -> bool: + return os.environ.get("SCURO_EXCLUSIVE_OVER_BUDGET", "0") == "1" + + class MemoryAwareNodeScheduler: def __init__( @@ -110,6 +114,7 @@ def __init__( }, } self._cpu_reserved_nodes: Dict[str, int] = {} + self.exclusive_memory_nodes: Dict[str, str] = {} self._initialized = False def _total_cpu_in_use(self) -> float: @@ -184,6 +189,16 @@ def get_runnable(self) -> List[RepresentationNode]: def _get_runnable_nodes(self) -> List[str]: runnable_nodes = list(self._candidates) + # With exclusive admission a full budget no longer stops the run, so + # the cache of intermediates must be drained deliberately: while the + # CPU budget is exceeded, nodes whose completion does not grow the + # cache (task roots, or last children releasing at least their own + # output) run before nodes that would cache another output. + draining_first = ( + _exclusive_over_budget_enabled() + and self._total_cpu_in_use() >= self.memory_budget["cpu"] + ) + def _score(node_id: str): release_bytes = 0 for parent_id in self.parents.get(node_id, set()): @@ -192,7 +207,12 @@ def _score(node_id: str): and self.remaining_children.get(parent_id, 0) == 1 ): release_bytes += self.node_resources[parent_id][0] + adds_bytes = ( + self.node_resources[node_id][0] if self.children.get(node_id) else 0 + ) + grows_cache = draining_first and adds_bytes > release_bytes return ( + grows_cache, -self.node_priorities.get(node_id, 0.0), -release_bytes, node_id not in self.roots, @@ -430,7 +450,10 @@ def not_enough_memory(self) -> bool: return False for node_id in self._candidates: cpu_mem, gpu_mem = self.node_resources[node_id] - if cpu_mem > self.memory_budget["cpu"] - self._total_cpu_in_use(): + if ( + cpu_mem > self.memory_budget["cpu"] - self._total_cpu_in_use() + and not _exclusive_over_budget_enabled() + ): self.deadlock_reason = ( f"node {node_id} needs {cpu_mem / 1024**3:.2f} GB CPU but only " f"{(self.memory_budget['cpu'] - self._total_cpu_in_use()) / 1024**3:.2f} " @@ -461,11 +484,26 @@ def _check_memory_constraints( ) -> bool: cpu_mem, gpu_mem = self.node_resources[node_id] gpu_id = None + free_cpu = self.memory_budget["cpu"] - self._total_cpu_in_use() - pending_bytes if ( - cpu_mem - > self.memory_budget["cpu"] - self._total_cpu_in_use() - pending_bytes + cpu_mem > free_cpu + and _exclusive_over_budget_enabled() + and ( + not self.running_nodes + and not self._ready_set + and pending_bytes == 0 + and self.memory_stats["cpu_in_flight"] == 0 + ) ): - if cpu_mem > self.memory_budget["cpu"]: + self.exclusive_memory_nodes[node_id] = ( + f"estimated {cpu_mem / 1024**3:.2f} GB with {free_cpu / 1024**3:.2f} GB " + f"of the CPU budget free; run alone" + ) + elif cpu_mem > free_cpu: + if ( + cpu_mem > self.memory_budget["cpu"] + and not _exclusive_over_budget_enabled() + ): self.blocked_memory_nodes_perm.append(node_id) self.blocked_memory_reasons[node_id] = ( f"estimated CPU peak {cpu_mem / 1024**3:.2f} GB exceeds the " diff --git a/src/main/python/systemds/scuro/drsearch/operator_registry.py b/src/main/python/systemds/scuro/drsearch/operator_registry.py index 1dff015a905..e06a602a676 100644 --- a/src/main/python/systemds/scuro/drsearch/operator_registry.py +++ b/src/main/python/systemds/scuro/drsearch/operator_registry.py @@ -36,10 +36,10 @@ ModalityType.TEXT: ("bert", "bow", "clip", "glove", "tfidf", "word2vec"), ModalityType.IMAGE: ("clip", "color_histogram", "openface", "resnet", "vgg"), ModalityType.VIDEO: ( + "openface", "clip", "color_histogram", "image_bind", - "openface", "resnet", "swin_video_transformer", "vgg", diff --git a/src/main/python/systemds/scuro/drsearch/representation_dag.py b/src/main/python/systemds/scuro/drsearch/representation_dag.py index e0c39537ae0..e8366cefc4b 100644 --- a/src/main/python/systemds/scuro/drsearch/representation_dag.py +++ b/src/main/python/systemds/scuro/drsearch/representation_dag.py @@ -36,6 +36,7 @@ DimensionalityReduction, ) from systemds.scuro.utils.identifier import get_op_id, get_node_id +import time from collections import OrderedDict, defaultdict, deque @@ -330,9 +331,25 @@ def execute( rep_cache: Dict[Any, TransformedModality] = None, consumer_count: Dict[str, int] = None, gpu_id: int = None, + node_cache=None, ) -> Union[Dict[str, TransformedModality], TransformedModality]: def execute_node(node_id: str, task) -> TransformedModality: + if ( + node_cache is None + or node_id == self.root_node_id + or not self.get_node_by_id(node_id).inputs + ): + return compute_node(node_id, task) + cached = node_cache.get(node_id) + if cached is not None: + return cached + start = time.perf_counter() + result = compute_node(node_id, task) + node_cache.put(node_id, result, time.perf_counter() - start) + return result + + def compute_node(node_id: str, task) -> TransformedModality: if external_cache is not None: cached = external_cache.get(node_id) if cached is not None: 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 ebeecc733d6..97de3da3a58 100644 --- a/src/main/python/systemds/scuro/drsearch/split_aware_evaluator.py +++ b/src/main/python/systemds/scuro/drsearch/split_aware_evaluator.py @@ -63,8 +63,9 @@ def dag_requires_fitting(dag: RepresentationDag) -> bool: class SplitAwareDAGEvaluator: - def __init__(self, dag: RepresentationDag, modalities, task): + def __init__(self, dag: RepresentationDag, modalities, task, node_cache=None): self.dag = dag + self.node_cache = node_cache self.modalities = ( list(modalities) if isinstance(modalities, (list, tuple)) else [modalities] ) @@ -151,8 +152,10 @@ def _align_partition_shapes(result): data = item.modality.data if isinstance(data, np.ndarray): array = data - elif isinstance(data, list) and data and all( - isinstance(d, np.ndarray) for d in data + elif ( + isinstance(data, list) + and data + and all(isinstance(d, np.ndarray) for d in data) ): if len({d.shape for d in data}) != 1: return result @@ -178,8 +181,12 @@ def _align_partition_shapes(result): padded[(slice(None),) + tuple(slice(0, dim) for dim in shape)] = source modality.data = padded if array is not None else list(padded) for metadata in modality.metadata or []: - layout = metadata.get("data_layout") if isinstance(metadata, dict) else None - if isinstance(layout, dict) and len(tuple(layout.get("shape") or ())) == len(target): + layout = ( + metadata.get("data_layout") if isinstance(metadata, dict) else None + ) + if isinstance(layout, dict) and len( + tuple(layout.get("shape") or ()) + ) == len(target): layout["shape"] = target return result @@ -206,6 +213,10 @@ def _leaf_input(partition, node): def _execute_partitions(self, *partitions): cache = {} + fitted = set() + partition_key = tuple( + tuple(item.instance_ids for item in partition) for partition in partitions + ) def execute_node(node_id): if node_id in cache: @@ -220,6 +231,17 @@ def execute_node(node_id): return result input_results = [execute_node(input_id) for input_id in node.inputs] + reusable = ( + self.node_cache is not None + and node_id != self.dag.root_node_id + and not any(input_id in fitted for input_id in node.inputs) + ) + if reusable: + cached = self.node_cache.get(node_id, partition_key) + if cached is not None: + cache[node_id] = cached + return cached + start = time.perf_counter() input_partitions = tuple( tuple(result[i] for result in input_results) for i in range(len(partitions)) @@ -231,6 +253,8 @@ def execute_node(node_id): requirement = getattr(operation, "fit_requirement", FitRequirement.NONE) if requirement is not FitRequirement.NONE: + fitted.add(node_id) + reusable = False labels = None fit_input = input_partitions[0] fit_ids = ( @@ -257,6 +281,12 @@ def execute_node(node_id): ) result = self._align_partition_shapes(result) + if any(input_id in fitted for input_id in node.inputs): + fitted.add(node_id) + if reusable: + self.node_cache.put( + node_id, result, time.perf_counter() - start, partition_key + ) cache[node_id] = result return result diff --git a/src/main/python/systemds/scuro/drsearch/trial_result_cache.py b/src/main/python/systemds/scuro/drsearch/trial_result_cache.py new file mode 100644 index 00000000000..9298dbe46c2 --- /dev/null +++ b/src/main/python/systemds/scuro/drsearch/trial_result_cache.py @@ -0,0 +1,212 @@ +# ------------------------------------------------------------- +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# +# ------------------------------------------------------------- + +import copy +import os +import threading +import time +from collections import OrderedDict +from typing import Any, Dict, Hashable, Optional, Tuple + +import numpy as np + +from systemds.scuro.modality.indexed import IndexedModality + +_GB = 1024**3 + + +def _nbytes(value: Any) -> int: + if value is None: + return 0 + if isinstance(value, np.ndarray): + if value.dtype == object: + return sum(_nbytes(item) for item in value.ravel()) + return int(value.nbytes) + if isinstance(value, IndexedModality): + return _nbytes(value.modality) + 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()) + if hasattr(value, "_data"): + return _nbytes(value._data) + _nbytes(getattr(value, "metadata", None)) + if hasattr(value, "nbytes"): + try: + return int(value.nbytes) + except (TypeError, ValueError): + return 0 + return 0 + + +def _clone(value: Any) -> Any: + if isinstance(value, IndexedModality): + return IndexedModality(_clone(value.modality), value.instance_ids) + if isinstance(value, tuple): + return tuple(_clone(item) for item in value) + if hasattr(value, "_data"): + clone = copy.copy(value) + clone._data = copy.deepcopy(value._data) + clone.metadata = copy.deepcopy(value.metadata) + return clone + return copy.deepcopy(value) + + +class TrialResultCache: + def __init__(self, max_bytes: int, min_compute_seconds: float = 1.0): + self.max_bytes = max(0, int(max_bytes)) + self.min_compute_seconds = min_compute_seconds + self._entries: "OrderedDict[Hashable, Tuple[Any, int, float]]" = OrderedDict() + self._lock = threading.Lock() + self.bytes = 0 + self.peak_bytes = 0 + self.hits = 0 + self.misses = 0 + self.stored = 0 + self.evictions = 0 + self.skipped_too_large = 0 + self.saved_seconds = 0.0 + + @classmethod + def from_env(cls) -> Optional["TrialResultCache"]: + budget_gb = float(os.environ.get("SCURO_HPO_CACHE_GB", "16")) + if budget_gb <= 0: + return None + min_seconds = float(os.environ.get("SCURO_HPO_CACHE_MIN_SECONDS", "1")) + return cls(int(budget_gb * _GB), min_seconds) + + def get(self, key: Hashable) -> Optional[Any]: + with self._lock: + entry = self._entries.get(key) + if entry is None: + self.misses += 1 + return None + self._entries.move_to_end(key) + value, _, compute_seconds = entry + self.hits += 1 + self.saved_seconds += compute_seconds + return _clone(value) + + def put(self, key: Hashable, value: Any, compute_seconds: float) -> None: + if value is None or compute_seconds < self.min_compute_seconds: + return + size = _nbytes(value) + if size > self.max_bytes: + self.skipped_too_large += 1 + return + stored = _clone(value) + with self._lock: + if key in self._entries: + return + while self._entries and self.bytes + size > self.max_bytes: + _, (_, evicted_size, _) = self._entries.popitem(last=False) + self.bytes -= evicted_size + self.evictions += 1 + self._entries[key] = (stored, size, compute_seconds) + self.bytes += size + self.peak_bytes = max(self.peak_bytes, self.bytes) + self.stored += 1 + + def clear(self) -> None: + with self._lock: + self._entries.clear() + self.bytes = 0 + + def summary(self) -> Dict[str, Any]: + return { + "hits": self.hits, + "misses": self.misses, + "stored": self.stored, + "evictions": self.evictions, + "skipped_too_large": self.skipped_too_large, + "peak_gb": round(self.peak_bytes / _GB, 3), + "budget_gb": round(self.max_bytes / _GB, 3), + "saved_seconds": round(self.saved_seconds, 1), + } + + +class TrialNodeCache: + def __init__(self, store: TrialResultCache, signatures: Dict[str, Hashable]): + self.store = store + self.signatures = signatures + self.hits = 0 + + def get(self, node_id: str, extra: Hashable = ()) -> Optional[Any]: + signature = self.signatures.get(node_id) + if signature is None: + return None + value = self.store.get((signature, extra)) + if value is not None: + self.hits += 1 + return value + + def put( + self, node_id: str, value: Any, compute_seconds: float, extra: Hashable = () + ) -> None: + signature = self.signatures.get(node_id) + if signature is not None: + self.store.put((signature, extra), value, compute_seconds) + + +def _freeze(value: Any) -> Hashable: + if isinstance(value, dict): + return tuple(sorted((k, _freeze(v)) for k, v in value.items())) + if isinstance(value, (list, tuple)): + return tuple(_freeze(v) for v in value) + if isinstance(value, float): + return round(value, 10) + if isinstance(value, type): + return f"{value.__module__}.{value.__qualname__}" + try: + hash(value) + except TypeError: + return repr(value) + return value + + +def trial_node_signatures(dag, params: Dict[str, Any]) -> Dict[str, Hashable]: + own = {} + for name, value in params.items(): + node_id, _, param = name.rpartition("-") + own.setdefault(node_id, []).append((param, _freeze(value))) + + signatures: Dict[str, Hashable] = {} + + def visit(node_id: str) -> Hashable: + if node_id in signatures: + return signatures[node_id] + node = dag.get_node_by_id(node_id) + inputs = tuple(visit(input_id) for input_id in node.inputs) + signature = (node_id, tuple(sorted(own.get(node_id, ()))), inputs) + signatures[node_id] = signature + return signature + + visit(dag.root_node_id) + return signatures + + +class timed: + def __enter__(self): + self._start = time.perf_counter() + return self + + def __exit__(self, *exc): + self.seconds = time.perf_counter() - self._start + return False diff --git a/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py b/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py index bd18a9beb52..ced501ba095 100644 --- a/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py +++ b/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py @@ -211,6 +211,9 @@ def _count_results_by_modality(self, results) -> Dict[Any, int]: return counts + def _skip_for_modality(self, modality_id) -> int: + return int(self._checkpoint_manager.skip_remaining_by_key.get(modality_id, 0)) + def resume_from_checkpoint(self): loaded = self._checkpoint_manager.resume_from_checkpoint( "eval_count_by_modality", self._count_results_by_modality @@ -233,10 +236,7 @@ def optimize_parallel(self, n_workers=None): executor.submit( self._process_modality, modality, - self._checkpoint_manager.skip_remaining_by_key.get( - modality.modality_id, 0 - ) - / len(self.tasks), + self._skip_for_modality(modality.modality_id), scheduler=None, ): modality for modality in self.modalities @@ -282,14 +282,7 @@ def _optimize_modalities(self): local_result, execution_time = self._process_modality( modality, ( - int( - round( - self._checkpoint_manager.skip_remaining_by_key.get( - modality.modality_id, 0 - ) - / len(self.tasks) - ) - ) + self._skip_for_modality(modality.modality_id) if self.resume else 0 ), @@ -488,11 +481,17 @@ def _build_execution_dags_for_modality( ) dags = self.add_aggregation_operator(self.builders[modality.modality_id], dags) - dags_with_pushdown = pushdown_aggregation(dags) + built = len(dags) if skip_remaining > 0: dags = dags[skip_remaining:] + print( + f"[resume] modality {modality.modality_id}: built {built} dag(s), " + f"skip {skip_remaining}, executing {len(dags)}", + flush=True, + ) + dags_with_pushdown = pushdown_aggregation(dags) expanded_dags_with_task_roots = self._expand_dags_with_task_roots( dags_with_pushdown ) @@ -671,6 +670,11 @@ def _dag_output_length(self, dag: RepresentationDag) -> Optional[int]: return None return int(output_shape[0]) + def _dag_output_has_time_axis(self, dag: RepresentationDag) -> bool: + stats = self._node_output_stats(dag).get(dag.root_node_id, None) + output_shape = getattr(stats, "output_shape", None) + return output_shape is not None and len(output_shape) >= 2 + def _aggregation_needed(self, dag: RepresentationDag) -> bool: return len(self._node_output_stats(dag)[dag.root_node_id].output_shape) > 1 @@ -720,6 +724,19 @@ def default_context_operators( context_operators = self._get_context_operators( rep_root.operation().output_modality_type ) + if not self._dag_output_has_time_axis(rep_dag): + for context_op in context_operators: + self.pruned.append( + { + "operation": context_op.__name__, + "input": rep_root.operation.__name__, + "reason": ( + "representation output is a rank-1 feature " + "vector with no time axis to window over" + ), + } + ) + return dags output_length = self._dag_output_length(rep_dag) for context_op in context_operators: context_operator_instance = context_op() @@ -941,7 +958,7 @@ def add_task_result(self, task_result: ResultEntry, dags: List[RepresentationDag if self._dag_index_source is not dags: self._dag_index = {dag.dag_id: dag for dag in dags} self._dag_index_source = dags - task_result.dag = self._dag_index.get(dag_id) + task_result.dag = self._dag_index.get(dag_id, task_result.dag) self.results[task_result.dag.nodes[0].modality_id][task_name].append( task_result ) diff --git a/src/main/python/systemds/scuro/modality/indexed.py b/src/main/python/systemds/scuro/modality/indexed.py index 9841b8b222d..a45eb0a0a7b 100644 --- a/src/main/python/systemds/scuro/modality/indexed.py +++ b/src/main/python/systemds/scuro/modality/indexed.py @@ -39,9 +39,16 @@ def __post_init__(self): num_instances = len(self.modality.data) elif self.modality.has_metadata(): num_instances = len(self.modality.metadata) + elif not instance_ids: + num_instances = 0 else: raise ValueError( - "IndexedModality requires materialized data or instance metadata" + "IndexedModality requires materialized data or instance metadata " + f"for {len(instance_ids)} instance ID(s); modality " + f"id={getattr(self.modality, 'modality_id', None)} " + f"type={type(self.modality).__name__} " + f"data={type(getattr(self.modality, 'data', None)).__name__} " + f"metadata_len={len(self.modality.metadata or [])}" ) if len(instance_ids) != num_instances: diff --git a/src/main/python/systemds/scuro/modality/unimodal_modality.py b/src/main/python/systemds/scuro/modality/unimodal_modality.py index e91a30f7891..a615035bce3 100644 --- a/src/main/python/systemds/scuro/modality/unimodal_modality.py +++ b/src/main/python/systemds/scuro/modality/unimodal_modality.py @@ -60,21 +60,31 @@ def copy_from_instance(self): new_instance.metadata = self.metadata.copy() return new_instance - def subset(self, indices): - if self.data_loader.chunk_size is None: - return super().subset(indices) - + def _subset_loader(self, indices): indices = list(indices) subset_loader = copy.copy(self.data_loader) subset_loader.indices = [self.data_loader.indices[i] for i in indices] subset_loader.stats = copy.copy(self.data_loader.stats) if hasattr(subset_loader.stats, "num_instances"): subset_loader.stats.num_instances = len(indices) + if getattr(subset_loader.stats, "lengths", None) is not None: + subset_loader.stats.lengths = tuple( + subset_loader.stats.lengths[i] for i in indices + ) subset_loader.reset() # Re-run the setter after replacing indices so num_chunks reflects the # subset rather than the original dataset. subset_loader.chunk_size = self.data_loader.chunk_size + return subset_loader + + def subset(self, indices): + if self.data_loader.chunk_size is None or self.has_data(): + subset_modality = super().subset(indices) + subset_modality.data_loader = self._subset_loader(indices) + subset_modality.stats = subset_modality.data_loader.stats + return subset_modality + subset_loader = self._subset_loader(indices) subset_modality = type(self)(subset_loader) subset_modality.modality_id = self.modality_id subset_modality.transform_time = self.transform_time diff --git a/src/main/python/systemds/scuro/representations/bow.py b/src/main/python/systemds/scuro/representations/bow.py index 93e72ddf8fb..ee42160ccda 100644 --- a/src/main/python/systemds/scuro/representations/bow.py +++ b/src/main/python/systemds/scuro/representations/bow.py @@ -47,6 +47,10 @@ class BoW(UnimodalRepresentation): def __init__(self, ngram_range=2, min_df=2, output_file=None, params=None): parameters = {"ngram_range": [2, 3, 5, 10], "min_df": [1, 2, 4, 8]} super().__init__("BoW", ModalityType.EMBEDDING, parameters) + if params is not None: + ngram_range = params.get("ngram_range", ngram_range) + min_df = params.get("min_df", min_df) + output_file = params.get("output_file", output_file) self.ngram_range = int(ngram_range) self.min_df = int(min_df) self.output_file = output_file diff --git a/src/main/python/systemds/scuro/representations/covarep_audio_features.py b/src/main/python/systemds/scuro/representations/covarep_audio_features.py index 973b7a99c2f..ab3c5d02afa 100644 --- a/src/main/python/systemds/scuro/representations/covarep_audio_features.py +++ b/src/main/python/systemds/scuro/representations/covarep_audio_features.py @@ -44,7 +44,9 @@ def __init__(self, hop_length=512, params=None): "hop_length": [256, 512, 1024, 2048], } # TODO super().__init__("Spectral", ModalityType.EMBEDDING, parameters, False) - self.hop_length = hop_length + if params is not None: + hop_length = params.get("hop_length", hop_length) + self.hop_length = int(hop_length) def transform(self, modality, aggregation=None): transformed_modality = TransformedModality( @@ -137,7 +139,9 @@ def __init__(self, hop_length=512, params=None): "hop_length": [256, 512, 1024, 2048], } # TODO super().__init__("ZeroCrossing", ModalityType.EMBEDDING, parameters, False) - self.hop_length = hop_length + if params is not None: + hop_length = params.get("hop_length", hop_length) + self.hop_length = int(hop_length) def transform(self, modality, aggregation=None): transformed_modality = TransformedModality( @@ -201,8 +205,11 @@ def __init__(self, frame_length=1024, hop_length=512, params=None): "hop_length": [256, 512, 1024, 2048], } # TODO super().__init__("RMSE", ModalityType.EMBEDDING, parameters, False) - self.hop_length = hop_length - self.frame_length = frame_length + if params is not None: + frame_length = params.get("frame_length", frame_length) + hop_length = params.get("hop_length", hop_length) + self.hop_length = int(hop_length) + self.frame_length = int(frame_length) def transform(self, modality, aggregation=None): transformed_modality = TransformedModality( @@ -266,7 +273,9 @@ def __init__(self, hop_length=512, params=None): "hop_length": [256, 512, 1024, 2048], } # TODO super().__init__("Pitch", ModalityType.EMBEDDING, parameters, False) - self.hop_length = hop_length + if params is not None: + hop_length = params.get("hop_length", hop_length) + self.hop_length = int(hop_length) def transform(self, modality, aggregation=None): transformed_modality = TransformedModality( diff --git a/src/main/python/systemds/scuro/representations/elmo.py b/src/main/python/systemds/scuro/representations/elmo.py index 85b04183d58..d4e8ee63a76 100644 --- a/src/main/python/systemds/scuro/representations/elmo.py +++ b/src/main/python/systemds/scuro/representations/elmo.py @@ -43,6 +43,11 @@ def __init__( output_file=None, params=None, ): + if params is not None: + model_name = params.get("model_name", model_name) + layer = params.get("layer", layer) + pooling = params.get("pooling", pooling) + output_file = params.get("output_file", output_file) self.data_type = torch.float32 self.model_name = model_name self.layer_name = layer diff --git a/src/main/python/systemds/scuro/representations/openface.py b/src/main/python/systemds/scuro/representations/openface.py index 7715e85951f..2824f501d7f 100644 --- a/src/main/python/systemds/scuro/representations/openface.py +++ b/src/main/python/systemds/scuro/representations/openface.py @@ -20,8 +20,7 @@ # ------------------------------------------------------------- import os import tempfile -from contextlib import redirect_stdout -from io import StringIO +from contextlib import contextmanager from pathlib import Path import cv2 @@ -45,6 +44,20 @@ get_device, ) + +@contextmanager +def _without_tf32(): + cudnn_tf32 = torch.backends.cudnn.allow_tf32 + matmul_tf32 = torch.backends.cuda.matmul.allow_tf32 + torch.backends.cudnn.allow_tf32 = False + torch.backends.cuda.matmul.allow_tf32 = False + try: + yield + finally: + torch.backends.cudnn.allow_tf32 = cudnn_tf32 + torch.backends.cuda.matmul.allow_tf32 = matmul_tf32 + + _retinaface_pretrain_patched = False _star_dirs_patched = False @@ -152,11 +165,16 @@ class OpenFace(UnimodalRepresentation): FEATURE_COLUMNS = MULTITASK_COLUMNS + LANDMARK_COLUMNS FEATURE_DIM = len(FEATURE_COLUMNS) + FRAME_BATCH_SIZE = 32 + FACE_MODEL_MEMORY_BYTES = 8 * 1024 * 1024 MULTITASK_MODEL_MEMORY_BYTES = 128 * 1024 * 1024 LANDMARK_MODEL_MEMORY_BYTES = 192 * 1024 * 1024 CPU_RUNTIME_OVERHEAD_BYTES = 128 * 1024 * 1024 GPU_RUNTIME_OVERHEAD_BYTES = 64 * 1024 * 1024 + LANDMARK_RUNTIME_BYTES = 2400 * 1024 * 1024 + LANDMARK_BYTES_PER_FRAME = 52 * 1024 * 1024 + MULTITASK_BYTES_PER_FRAME = 26 * 1024 * 1024 def __init__( self, @@ -282,8 +300,12 @@ def estimate_peak_memory_bytes(self, input_stats) -> dict: max_height * max_width * max_channels * np.dtype(np.float32).itemsize ) model_bytes = self.FACE_MODEL_MEMORY_BYTES + self.MULTITASK_MODEL_MEMORY_BYTES + batch_bytes = self.MULTITASK_BYTES_PER_FRAME + runtime_bytes = self.GPU_RUNTIME_OVERHEAD_BYTES if self.feature_set in ("landmarks", "all"): model_bytes += self.LANDMARK_MODEL_MEMORY_BYTES + batch_bytes += self.LANDMARK_BYTES_PER_FRAME + runtime_bytes += self.LANDMARK_RUNTIME_BYTES output_bytes = self.estimate_output_memory_bytes(input_stats) return { @@ -294,7 +316,9 @@ def estimate_peak_memory_bytes(self, input_stats) -> dict: + self.CPU_RUNTIME_OVERHEAD_BYTES ), "gpu_peak_bytes": int( - model_bytes + 2 * frame_bytes + self.GPU_RUNTIME_OVERHEAD_BYTES + model_bytes + + self.FRAME_BATCH_SIZE * (frame_bytes + batch_bytes) + + runtime_bytes ), } @@ -338,55 +362,77 @@ def transform(self, modality, aggregation=None): return transformed_modality def _extract_image_features(self, images): - return [self._extract_features(image) for image in images] + return list(self._extract_batch_features([self._to_bgr(i) for i in images])) def _extract_video_features(self, frames, owner_id): if len(frames) == 0: raise ValueError(f"Video instance {owner_id} contains no frames") - return np.stack([self._extract_features(frame) for frame in frames]) + return self._extract_batch_features([self._to_bgr(frame) for frame in frames]) def _extract_features(self, image): + return self._extract_batch_features([self._to_bgr(image)])[0] + + def _extract_batch_features(self, images_bgr): self._load_models() - image_bgr = self._to_bgr(image) - face, detections = self._face_detector.get_face(image_bgr) - if face is None or face.size == 0: - return np.zeros(self.feature_dim, dtype=np.float32) + with _without_tf32(): + return self._extract_batch_features_fp32(images_bgr) + + def _extract_batch_features_fp32(self, images_bgr): + features = np.zeros((len(images_bgr), self.feature_dim), dtype=np.float32) + faces, detections = self._detect_faces(images_bgr) + with_face = [ + index + for index, face in enumerate(faces) + if face is not None and face.size > 0 + ] + if not with_face: + return features needs_backbone = self.feature_set in ("backbone", "all") - if needs_backbone: - self._ensure_backbone_hook() - self._backbone_output = None - - emotion_output, gaze_output, action_unit_output = ( - self._multitask_predictor.predict(face) + emotion, gaze, action_units, backbone = self._predict_multitask( + [faces[index] for index in with_face], needs_backbone ) - behavioral = self._behavioral_features(gaze_output, action_unit_output) - emotion = self._checked_vector( - emotion_output, len(self.EMOTION_COLUMNS), "emotion" - ) - multitask = np.concatenate((behavioral, emotion)) - - if self.feature_set == "behavioral": - return behavioral - if self.feature_set == "multitask": - return multitask.astype(np.float32, copy=False) - - backbone = None - if needs_backbone: - backbone = self._checked_vector( - self._backbone_output, self.BACKBONE_DIM, "backbone" + landmarks = None + if self.feature_set in ("landmarks", "all"): + landmarks = self._landmark_features( + [images_bgr[index] for index in with_face], + [detections[index] for index in with_face], ) - if self.feature_set == "backbone": - return backbone - - landmarks = self._landmark_features(image_bgr, detections) - if self.feature_set == "landmarks": - return np.concatenate((multitask, landmarks)).astype(np.float32, copy=False) - detection = self._detection_features(detections, image_bgr.shape) - return np.concatenate((multitask, detection, landmarks, backbone)).astype( - np.float32, copy=False - ) + for row, index in enumerate(with_face): + behavioral = self._behavioral_features(gaze[row], action_units[row]) + if self.feature_set == "behavioral": + features[index] = behavioral + continue + multitask = np.concatenate( + ( + behavioral, + self._checked_vector( + emotion[row], len(self.EMOTION_COLUMNS), "emotion" + ), + ) + ) + if self.feature_set == "multitask": + features[index] = multitask + continue + backbone_row = None + if needs_backbone: + backbone_row = self._checked_vector( + backbone[row], self.BACKBONE_DIM, "backbone" + ) + if self.feature_set == "backbone": + features[index] = backbone_row + continue + if self.feature_set == "landmarks": + features[index] = np.concatenate((multitask, landmarks[row])) + continue + detection = self._detection_features( + detections[index], images_bgr[index].shape + ) + features[index] = np.concatenate( + (multitask, detection, landmarks[row], backbone_row) + ) + return features def _behavioral_features(self, gaze_output, action_unit_output): gaze = self._checked_vector(gaze_output, 2, "gaze") @@ -395,27 +441,133 @@ def _behavioral_features(self, gaze_output, action_unit_output): ) return np.concatenate((gaze, action_units)).astype(np.float32, copy=False) - def _landmark_features(self, image_bgr, detections): - with redirect_stdout(StringIO()): - landmarks = self._landmark_detector.detect_landmarks( - image_bgr, - detections[:1], - confidence_threshold=self.vis_threshold, + def _chunks(self, count): + for start in range(0, count, self.FRAME_BATCH_SIZE): + yield start, min(start + self.FRAME_BATCH_SIZE, count) + + def _detect_faces(self, images_bgr): + 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 + variance = detector.cfg["variance"] + faces = [None] * len(images_bgr) + detections = [None] * len(images_bgr) + priors_by_shape = {} + + by_shape = {} + for index, image in enumerate(images_bgr): + by_shape.setdefault(image.shape, []).append(index) + + for shape, indices in by_shape.items(): + height, width = shape[:2] + if shape not in priors_by_shape: + priors_by_shape[shape] = ( + PriorBox(detector.cfg, image_size=(height, width)) + .forward() + .to(detector.device) + .data + ) + priors = priors_by_shape[shape] + box_scale = torch.Tensor([width, height, width, height]).to(detector.device) + landmark_scale = torch.Tensor([width, height] * 5).to(detector.device) + + for start, stop in self._chunks(len(indices)): + chunk = indices[start:stop] + batch = torch.cat( + [detector.preprocess_image(images_bgr[i])[0] for i in chunk] + ) + 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 + ) + keep = py_cpu_nms(dets, detector.nms_threshold) + dets = np.concatenate((dets[keep], points[keep]), axis=1) + + if len(dets) == 0 or dets[0][4] < detector.vis_threshold: + continue + bbox = dets[0][:4].astype(int) + faces[index] = images_bgr[index][ + bbox[1] : bbox[3], bbox[0] : bbox[2] + ] + detections[index] = dets + return faces, detections + + def _predict_multitask(self, faces, needs_backbone): + predictor = self._multitask_predictor + if 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]] ) - if not landmarks: - return np.zeros(len(self.LANDMARK_COLUMNS), dtype=np.float32) + self._backbone_output = None + with torch.no_grad(): + emotion, gaze, action_units = predictor.model(batch) + for output, values in zip( + outputs, (emotion, gaze, action_units, self._backbone_output) + ): + if values is not None: + output.append(values.detach().cpu().numpy()) + emotion, gaze, action_units, backbone = ( + np.concatenate(output) if output else None for output in outputs + ) + return emotion, gaze, action_units, backbone - points = np.asarray(landmarks[0], dtype=np.float32) - if points.shape != (self.NUM_LANDMARKS, 2): - raise RuntimeError( - "OpenFace 3.0 returned unexpected landmark dimensions: " - f"{points.shape}" + def _landmark_features(self, images_bgr, detections): + alignment = self._landmark_detector.alignment + features = np.zeros( + (len(images_bgr), len(self.LANDMARK_COLUMNS)), dtype=np.float32 + ) + crops = [] + for index, dets in enumerate(detections): + det = dets[0] + 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), + ) ) - height, width = image_bgr.shape[:2] - points = points.copy() - points[:, 0] /= width - points[:, 1] /= height - return points.reshape(-1) + + for start, stop in self._chunks(len(crops)): + chunk = crops[start:stop] + 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) + return features @classmethod def _detection_features(cls, detections, image_shape): @@ -555,6 +707,15 @@ def _to_bgr(image): f"Expected an image tensor with 2 or 3 dimensions, got {image.shape}" ) + if image.dtype in (np.float16, np.float32) and image.size: + values = image.astype(np.float32, copy=False) + low, high = values.min(), values.max() + if np.isfinite(low) and np.isfinite(high): + if low >= 0.0 and high <= 1.0: + values = values * np.float32(255.0) + if image.dtype == np.float16: + values = values.astype(np.float16).astype(np.float32) + image = np.clip(values, 0, 255).astype(np.uint8) if np.issubdtype(image.dtype, np.floating): finite = image[np.isfinite(image)] if finite.size and finite.min() >= 0.0 and finite.max() <= 1.0: diff --git a/src/main/python/systemds/scuro/representations/spectrogram.py b/src/main/python/systemds/scuro/representations/spectrogram.py index 9d14931ef0e..e4822386d77 100644 --- a/src/main/python/systemds/scuro/representations/spectrogram.py +++ b/src/main/python/systemds/scuro/representations/spectrogram.py @@ -50,6 +50,9 @@ class Spectrogram(UnimodalRepresentation): def __init__(self, hop_length=512, n_fft=2048, params=None): parameters = {"hop_length": [256, 512, 1024, 2048], "n_fft": [1024, 2048, 4096]} super().__init__("Spectrogram", ModalityType.TIMESERIES, parameters, False) + if params is not None: + hop_length = params.get("hop_length", hop_length) + n_fft = params.get("n_fft", n_fft) self.hop_length = int(hop_length) self.n_fft = int(n_fft) self.window_size = self.n_fft 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 f57ac18f037..3f75758cccd 100644 --- a/src/main/python/systemds/scuro/representations/swin_video_transformer.py +++ b/src/main/python/systemds/scuro/representations/swin_video_transformer.py @@ -32,13 +32,15 @@ from systemds.scuro.drsearch.operator_registry import register_representation from systemds.scuro.dataloader.video_loader import VideoStats from systemds.scuro.representations.utils import ( - LengthBucketBatchSampler, OwnerAccumulator, get_sequence_lengths, - move_batch_to_device, - pin_memory_for, transformer_inference_context, ) +from systemds.scuro.representations.video_batching import ( + MIB, + FrameBudgetedVideoModel, + VideoMemoryProfile, +) from systemds.scuro.utils.torch_dataset import CustomDataset from systemds.scuro.utils.static_variables import ( @@ -48,11 +50,18 @@ @register_representation([ModalityType.VIDEO]) -class SwinVideoTransformer(UnimodalRepresentation): +class SwinVideoTransformer(FrameBudgetedVideoModel, UnimodalRepresentation): _EMBED_DIM = 768 cache_in_worker = True - - def __init__(self, layer_name="avgpool", batch_size=8, params=None): + _MEMORY_PROFILE = VideoMemoryProfile( + clip_frame_bytes=1 * MIB, + batch_frame_bytes=19 * MIB, + cpu_frame_bytes=24 * MIB, + ) + + def __init__( + self, layer_name="avgpool", batch_size=8, params=None, max_batch_frames=None + ): parameters = { "layer_name": [ "features", @@ -73,6 +82,8 @@ def __init__(self, layer_name="avgpool", batch_size=8, params=None): batch_size = int(params.get("batch_size", batch_size)) self.layer_name = layer_name 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 @@ -99,43 +110,8 @@ def estimate_output_memory_bytes(self, input_stats: VideoStats) -> int: dt = int(torch.tensor([], dtype=self.data_type).element_size()) return input_stats.num_instances * self._EMBED_DIM * dt - def estimate_peak_memory_bytes(self, input_stats: VideoStats) -> dict: - dt = int(torch.tensor([], dtype=self.data_type).element_size()) - temporal = max(input_stats.max_length, 1) - input_bytes = ( - dt - * input_stats.max_channels - * temporal - * input_stats.max_height - * input_stats.max_width - ) - output_bytes = self.estimate_output_memory_bytes(input_stats) - n = max(input_stats.num_instances, 1) - output_bytes_batch = output_bytes / n - - batch_peak_bytes = (input_bytes + self._EMBED_DIM * dt) * 2 - - 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 - - cpu_peak = ( - size_all_bytes * 2 * dt - + output_bytes_batch - + output_bytes - + input_bytes - + safety_margin_bytes - ) - gpu_peak = (size_all_bytes * dt + batch_peak_bytes) * 6 - return {"cpu_peak_bytes": int(cpu_peak), "gpu_peak_bytes": int(gpu_peak)} + def _memory_profile(self) -> VideoMemoryProfile: + return self._MEMORY_PROFILE def transform(self, modality, aggregation=None): self.model = self.model.to(self.device) @@ -160,31 +136,26 @@ def hook( dataset = CustomDataset(modality.data, self.data_type, "cpu") lengths = get_sequence_lengths(modality.data, modality.metadata) - dataloader = torch.utils.data.DataLoader( - dataset, - batch_sampler=LengthBucketBatchSampler( - lengths, self.batch_size, exact=True - ), - pin_memory=pin_memory_for(self.device), - ) accumulator = OwnerAccumulator(len(dataset), len(dataset), aggregation) + def forward(model, batch): + video_ids = batch["id"].long() + _ = model(batch["data"].permute(0, 2, 1, 3, 4)) + values = self.swin_output + if isinstance(values, tuple): + values = values[0] + if values.ndim == 2: + pooled = values + elif self.layer_name.startswith("features"): + pooled = values.mean(dim=tuple(range(1, values.ndim - 1))) + else: + pooled = values.mean(dim=tuple(range(2, values.ndim))) + # A batch run on the CPU accumulates on the representation's device. + video_ids = video_ids.to(self.device) + accumulator.update(pooled.to(self.device), video_ids, video_ids) + with transformer_inference_context(self.device): - for batch in dataloader: - batch = move_batch_to_device(batch, self.device) - video_ids = batch["id"].long() - frames = batch["data"].permute(0, 2, 1, 3, 4) - _ = self.model(frames) - values = self.swin_output - if isinstance(values, tuple): - values = values[0] - if values.ndim == 2: - pooled = values - elif self.layer_name.startswith("features"): - pooled = values.mean(dim=tuple(range(1, values.ndim - 1))) - else: - pooled = values.mean(dim=tuple(range(2, values.ndim))) - accumulator.update(pooled, video_ids, video_ids) + self._run_batches(dataset, lengths, forward) transformed_modality = TransformedModality( modality, self, self.output_modality_type diff --git a/src/main/python/systemds/scuro/representations/text_context_with_indices.py b/src/main/python/systemds/scuro/representations/text_context_with_indices.py index 6d8964c1d14..ef4067982fc 100644 --- a/src/main/python/systemds/scuro/representations/text_context_with_indices.py +++ b/src/main/python/systemds/scuro/representations/text_context_with_indices.py @@ -94,6 +94,9 @@ def __init__(self, max_words=55, overlap=0, params=None): "overlap": [0, 10, 20, 30], } super().__init__("WordCountSplit", parameters) + if params is not None: + max_words = params.get("max_words", max_words) + overlap = params.get("overlap", overlap) self.max_words = int(max_words) self.overlap = max(0, int(overlap)) diff --git a/src/main/python/systemds/scuro/representations/utils.py b/src/main/python/systemds/scuro/representations/utils.py index 04e6829d7ab..4ad52e926f9 100644 --- a/src/main/python/systemds/scuro/representations/utils.py +++ b/src/main/python/systemds/scuro/representations/utils.py @@ -65,21 +65,34 @@ def aggregate_chunk_embeddings(embeddings, aggregation): class LengthBucketBatchSampler(torch.utils.data.Sampler): - """Build deterministic batches from samples with similar sequence lengths.""" - def __init__(self, lengths, batch_size, exact=False): + def __init__(self, lengths, batch_size, exact=False, max_batch_items=None): self.lengths = [int(length) for length in lengths] self.batch_size = max(1, int(batch_size)) self.exact = exact + self.max_batch_items = ( + None if max_batch_items is None else max(1, int(max_batch_items)) + ) self._batches = self._build_batches() + def _batch_size_for(self, length): + if self.max_batch_items is None: + return self.batch_size + return max(1, min(self.batch_size, self.max_batch_items // max(length, 1))) + def _build_batches(self): indices = sorted(range(len(self.lengths)), key=lambda i: (self.lengths[i], i)) if not self.exact: - return [ - indices[start : start + self.batch_size] - for start in range(0, len(indices), self.batch_size) - ] + batches = [] + batch = [] + for index in indices: + if batch and len(batch) >= self._batch_size_for(self.lengths[index]): + batches.append(batch) + batch = [] + batch.append(index) + if batch: + batches.append(batch) + return batches batches = [] start = 0 @@ -88,9 +101,10 @@ def _build_batches(self): end = start while end < len(indices) and self.lengths[indices[end]] == length: end += 1 + batch_size = self._batch_size_for(length) batches.extend( - indices[offset : min(offset + self.batch_size, end)] - for offset in range(start, end, self.batch_size) + indices[offset : min(offset + batch_size, end)] + for offset in range(start, end, batch_size) ) start = end return batches diff --git a/src/main/python/systemds/scuro/representations/video_batching.py b/src/main/python/systemds/scuro/representations/video_batching.py new file mode 100644 index 00000000000..2b9802e956c --- /dev/null +++ b/src/main/python/systemds/scuro/representations/video_batching.py @@ -0,0 +1,301 @@ +# ------------------------------------------------------------- +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# +# ------------------------------------------------------------- +import copy +from dataclasses import dataclass +from typing import List, Optional, Sequence + +import numpy as np +import torch + +from systemds.scuro.representations.utils import ( + LengthBucketBatchSampler, + move_batch_to_device, + pin_memory_for, +) + +MIB = 1024 * 1024 + + +@dataclass(frozen=True) +class VideoMemoryProfile: + clip_frame_bytes: float + batch_frame_bytes: float + clip_bytes: float = 0.0 + cpu_frame_bytes: float = 0.0 + + def gpu_bytes(self, clip_length: int, batch_frames: int, clips: int) -> float: + return ( + self.clip_frame_bytes * clip_length + + self.batch_frame_bytes * batch_frames + + self.clip_bytes * clips + ) + + +class FrameBudgetedVideoModel: + _BATCH_ACTIVATION_BUDGET_BYTES = 10 * 1024 * MIB + _ALLOCATOR_OVERHEAD = 1.65 + _GPU_WORKSPACE_BYTES = 2048 * MIB + _SAFETY_MARGIN_BYTES = 100 * MIB + _GPU_MEMORY_FRACTION = 0.9 + _MODEL_INPUT_FRAME_BYTES = 3 * 224 * 224 * 4 + + _max_batch_frames: Optional[int] = None + + def _memory_profile(self) -> VideoMemoryProfile: + raise NotImplementedError + + def _model_lengths(self, lengths: Sequence[int]) -> List[int]: + """Frames per clip as the model receives them.""" + return [max(int(length), 1) for length in lengths] + + @property + def max_batch_frames(self) -> int: + if self._max_batch_frames is not None: + return self._max_batch_frames + return max( + 1, + int( + self._BATCH_ACTIVATION_BUDGET_BYTES + // self._memory_profile().batch_frame_bytes + ), + ) + + @max_batch_frames.setter + 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: + 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: + torch.cuda.empty_cache() + self._last_batch_shape = shape + + def _longest_clip(self, input_stats) -> int: + max_length = int(getattr(input_stats, "max_length", 1) or 1) + return self._model_lengths([max_length])[0] + + def gpu_memory_limit(self) -> Optional[int]: + if not torch.cuda.is_available(): + return None + total = min( + torch.cuda.get_device_properties(i).total_memory + for i in range(torch.cuda.device_count()) + ) + return int(total * self._GPU_MEMORY_FRACTION) + + def _batch_gpu_bytes(self, clip_length: int, clips: int) -> float: + activations = self._memory_profile().gpu_bytes( + clip_length, clip_length * clips, clips + ) + return ( + self._model_bytes() + + activations * self._ALLOCATOR_OVERHEAD + + self._GPU_WORKSPACE_BYTES + ) + + def _clips_fitting_gpu(self, clip_length: int, clips: int, limit) -> int: + if limit is None: + return clips + while clips > 0 and self._batch_gpu_bytes(clip_length, clips) > limit: + clips -= 1 + return clips + + def _plan_batches(self, lengths, batch_size=None, limit=None): + lengths = self._model_lengths(lengths) + sampler = LengthBucketBatchSampler( + lengths, + self.batch_size if batch_size is None else batch_size, + exact=True, + max_batch_items=self.max_batch_frames, + ) + on_gpu, on_cpu = [], [] + for batch in sampler: + fitting = self._clips_fitting_gpu(lengths[batch[0]], len(batch), limit) + if fitting == 0: + on_cpu.extend((tuple([position]), False) for position in batch) + continue + on_gpu.extend( + (tuple(batch[start : start + fitting]), True) + for start in range(0, len(batch), fitting) + ) + return on_gpu + on_cpu + + def _device_limit(self): + device = getattr(self, "device", None) + if device is None or torch.device(device).type != "cuda": + return None + return self.gpu_memory_limit() + + def _cpu_model(self): + if torch.device(self.device).type == "cpu": + return self.model + return copy.deepcopy(self.model).to("cpu") + + def _run_batches(self, dataset, lengths, forward, collate_fn=None) -> None: + plan = self._plan_batches(lengths, limit=self._device_limit()) + on_device = [list(batch) for batch, gpu in plan if gpu] + on_cpu = [list(batch) for batch, gpu in plan if not gpu] + loader_args = {} if collate_fn is None else {"collate_fn": collate_fn} + failed = [] + if on_device: + loader = torch.utils.data.DataLoader( + dataset, + batch_sampler=on_device, + pin_memory=pin_memory_for(self.device), + **loader_args, + ) + for batch_positions, batch in zip(on_device, loader): + out_of_memory = False + try: + device_batch = move_batch_to_device(batch, self.device) + self._release_cache_on_new_shape(device_batch["data"]) + forward(self.model, device_batch) + except torch.cuda.OutOfMemoryError: + out_of_memory = True + if out_of_memory: + device_batch = None + torch.cuda.empty_cache() + self._last_batch_shape = None + failed.append(batch_positions) + on_cpu = failed + on_cpu + if on_cpu: + model = self._cpu_model() + loader = torch.utils.data.DataLoader( + dataset, batch_sampler=on_cpu, **loader_args + ) + for batch in loader: + forward(model, batch) + + def _peak_batch_frames(self, input_stats) -> int: + lengths = getattr(input_stats, "lengths", None) + if lengths: + model_lengths = self._model_lengths(lengths) + return max( + model_lengths[batch[0]] * len(batch) + for batch, _ in self._plan_batches(lengths) + ) + longest = self._longest_clip(input_stats) + return max(longest, min(self.batch_size * longest, self.max_batch_frames)) + + def _planned_peaks(self, input_stats, limit): + lengths = getattr(input_stats, "lengths", None) + if lengths: + model_lengths = self._model_lengths(lengths) + gpu, cpu_frames = 0.0, 0 + for batch, on_gpu in self._plan_batches(lengths, limit=limit): + length = model_lengths[batch[0]] + if on_gpu: + gpu = max(gpu, self._batch_gpu_bytes(length, len(batch))) + else: + cpu_frames = max(cpu_frames, length) + return gpu, cpu_frames + longest = self._longest_clip(input_stats) + gpu = ( + self._model_bytes() + + self._memory_profile().gpu_bytes( + longest, + self._peak_batch_frames(input_stats), + min(self.batch_size, self.max_batch_frames), + ) + * self._ALLOCATOR_OVERHEAD + + self._GPU_WORKSPACE_BYTES + ) + cpu_frames = 0 + if limit is not None: + gpu = min(gpu, limit) + if self._clips_fitting_gpu(longest, 1, limit) == 0: + cpu_frames = longest + return gpu, cpu_frames + + def _model_bytes(self) -> int: + return sum( + t.nelement() * t.element_size() + for t in (*self.model.parameters(), *self.model.buffers()) + ) + + @staticmethod + def _raw_frame_bytes(input_stats) -> int: + raw_dtype = getattr(input_stats, "dtype", None) + try: + itemsize = np.dtype(np.float32 if raw_dtype is None else raw_dtype).itemsize + except TypeError: + itemsize = np.dtype(np.float32).itemsize + return ( + max(int(getattr(input_stats, "max_height", 224) or 224), 1) + * max(int(getattr(input_stats, "max_width", 224) or 224), 1) + * max(int(getattr(input_stats, "max_channels", 3) or 3), 1) + * itemsize + ) + + def estimate_peak_memory_bytes(self, input_stats) -> dict: + frames = self._peak_batch_frames(input_stats) + model_bytes = self._model_bytes() + gpu_peak, cpu_frames = self._planned_peaks(input_stats, self.gpu_memory_limit()) + cpu_peak = ( + model_bytes + + frames + * (self._raw_frame_bytes(input_stats) + 3 * self._MODEL_INPUT_FRAME_BYTES) + + self.estimate_output_memory_bytes(input_stats) + + self._SAFETY_MARGIN_BYTES + ) + if cpu_frames: + cpu_peak += ( + model_bytes + cpu_frames * self._memory_profile().cpu_frame_bytes + ) + return {"cpu_peak_bytes": int(cpu_peak), "gpu_peak_bytes": int(gpu_peak)} + + def _distinct_batch_sizes(self, values, lengths, groups, limit): + seen = set() + kept = [] + for value in values: + key = tuple( + tuple( + (tuple(group[i] for i in batch), on_gpu) + for batch, on_gpu in self._plan_batches( + [lengths[i] for i in group], int(value), limit + ) + ) + for group in groups + ) + if key not in seen: + seen.add(key) + kept.append(value) + return kept + + def filter_parameter_domain(self, name, values, input_stats): + lengths = getattr(input_stats, "lengths", None) + if name != "batch_size" or not lengths: + return values + everything = list(range(len(lengths))) + groups_per_grouping = [[everything]] + per_chunk = int(getattr(input_stats, "num_instances", 0) or 0) + if 0 < per_chunk < len(lengths): + groups_per_grouping.append( + [ + everything[start : start + per_chunk] + for start in range(0, len(lengths), per_chunk) + ] + ) + groups = [group for grouping in groups_per_grouping for group in grouping] + return self._distinct_batch_sizes( + values, lengths, groups, self.gpu_memory_limit() + ) diff --git a/src/main/python/systemds/scuro/representations/x3d.py b/src/main/python/systemds/scuro/representations/x3d.py index fae60d1f1cd..755a52b10d8 100644 --- a/src/main/python/systemds/scuro/representations/x3d.py +++ b/src/main/python/systemds/scuro/representations/x3d.py @@ -34,13 +34,15 @@ from systemds.scuro.representations.representation import RepresentationStats from systemds.scuro.representations.unimodal import UnimodalRepresentation from systemds.scuro.representations.utils import ( - LengthBucketBatchSampler, get_sequence_lengths, inference_context, - move_batch_to_device, - pin_memory_for, save_embeddings, ) +from systemds.scuro.representations.video_batching import ( + MIB, + FrameBudgetedVideoModel, + VideoMemoryProfile, +) from systemds.scuro.utils.static_variables import ( get_device, get_device_for_model, @@ -54,8 +56,23 @@ def forward(self, input_: torch.Tensor) -> torch.Tensor: @register_representation([ModalityType.VIDEO]) -class X3D(UnimodalRepresentation): +class X3D(FrameBudgetedVideoModel, UnimodalRepresentation): cache_in_worker = True + _MEMORY_PROFILES = { + "r3d": VideoMemoryProfile( + clip_frame_bytes=42 * MIB, + batch_frame_bytes=6.5 * MIB, + clip_bytes=26 * MIB, + cpu_frame_bytes=16 * MIB, + ), + "s3d": VideoMemoryProfile( + clip_frame_bytes=5.5 * MIB, + batch_frame_bytes=3 * MIB, + clip_bytes=20 * MIB, + cpu_frame_bytes=8 * MIB, + ), + } + _MIN_FRAMES = 14 def __init__( self, @@ -64,6 +81,7 @@ def __init__( output_file=None, batch_size=8, params=None, + max_batch_frames=None, ): self.data_type = torch.float32 if params is not None: @@ -77,6 +95,7 @@ def __init__( self.output_file = output_file 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 @@ -107,45 +126,11 @@ def estimate_output_memory_bytes(self, input_stats: VideoStats) -> int: embedding_dim = 400 * math.floor((max(input_stats.max_length, 14) - 5) / 8) return input_stats.num_instances * embedding_dim * self.data_type.itemsize - def estimate_peak_memory_bytes(self, input_stats: VideoStats) -> dict: - temporal = max(input_stats.max_length, 14) - input_bytes = ( - self.batch_size - * self.data_type.itemsize - * input_stats.max_channels - * temporal - * input_stats.max_height - * input_stats.max_width - ) - output_bytes = self.estimate_output_memory_bytes(input_stats) - n = max(input_stats.num_instances, 1) - output_bytes_batch = output_bytes / n * self.batch_size - - batch_peak_bytes = ( - input_bytes + self.batch_size * 512 * self.data_type.itemsize - ) * 2 + def _memory_profile(self) -> VideoMemoryProfile: + return self._MEMORY_PROFILES[self.model_name] - 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 - - cpu_peak = ( - size_all_bytes * 2 * self.data_type.itemsize - + output_bytes_batch - + output_bytes - + input_bytes - + safety_margin_bytes - ) - gpu_peak = (size_all_bytes * self.data_type.itemsize + batch_peak_bytes) * 6 - return {"cpu_peak_bytes": int(cpu_peak), "gpu_peak_bytes": int(gpu_peak)} + def _model_lengths(self, lengths): + return [max(int(length), self._MIN_FRAMES) for length in lengths] @property def model_name(self): @@ -204,7 +189,9 @@ def _get_parameters(self, high_level=True): @staticmethod def _collate_videos(samples): video_ids = torch.tensor([sample["id"] for sample in samples]) - target_length = max(14, max(sample["data"].shape[0] for sample in samples)) + target_length = max( + X3D._MIN_FRAMES, max(sample["data"].shape[0] for sample in samples) + ) videos = [] for sample in samples: frames = sample["data"] @@ -239,34 +226,25 @@ def hook( break dataset = CustomDataset(modality.data, self.data_type, "cpu") - lengths = [ - max(length, 14) - for length in get_sequence_lengths(modality.data, modality.metadata) - ] - dataloader = torch.utils.data.DataLoader( - dataset, - batch_sampler=LengthBucketBatchSampler( - lengths, self.batch_size, exact=True - ), - collate_fn=self._collate_videos, - pin_memory=pin_memory_for(self.device), - ) + 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)) + 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 + with inference_context(self.device): - for batch in dataloader: - batch = move_batch_to_device(batch, self.device) - video_ids = batch["id"].long() - frames = batch["data"].permute(0, 2, 1, 3, 4) - _ = self.model(frames) - values = self.activation - if isinstance(values, tuple): - values = values[0] - if values.ndim > 2: - values = torch.nn.functional.adaptive_avg_pool2d(values, (1, 1)) - vectors = torch.flatten(values, 1).detach().float().cpu().numpy() - for video_id, vector in zip(video_ids.cpu().tolist(), vectors): - embeddings[video_id] = vector + self._run_batches( + dataset, lengths, forward, collate_fn=self._collate_videos + ) if self.output_file is not None: save_embeddings(embeddings, self.output_file) @@ -279,9 +257,15 @@ def hook( return transformed_modality -class I3D(UnimodalRepresentation): +class I3D(FrameBudgetedVideoModel, UnimodalRepresentation): _EMBEDDING_DIM = 400 cache_in_worker = True + _MEMORY_PROFILE = VideoMemoryProfile( + clip_frame_bytes=16 * MIB, + batch_frame_bytes=4 * MIB, + clip_bytes=36 * MIB, + cpu_frame_bytes=13 * MIB, + ) def __init__( self, @@ -290,6 +274,7 @@ def __init__( output_file=None, batch_size=8, params=None, + max_batch_frames=None, ): if params is not None: layer = params.get("layer_name", layer) @@ -306,6 +291,7 @@ def __init__( self.output_file = output_file self.layer_name = layer self.batch_size = batch_size + self.max_batch_frames = max_batch_frames self.data_type = torch.float32 self._gpu_id = self.device.index self._activation_hook = None @@ -336,6 +322,9 @@ def get_output_stats(self, input_stats) -> RepresentationStats: def estimate_output_memory_bytes(self, input_stats: VideoStats) -> int: return input_stats.num_instances * self._EMBEDDING_DIM * self.data_type.itemsize + def _memory_profile(self) -> VideoMemoryProfile: + return self._MEMORY_PROFILE + def _get_parameters(self, high_level=True): parameters = { "batch_size": [1, 2, 4, 8, 16, 32], @@ -379,29 +368,24 @@ def hook( break dataset = CustomDataset(modality.data, self.data_type, "cpu") - dataloader = torch.utils.data.DataLoader( - dataset, - batch_sampler=LengthBucketBatchSampler( - get_sequence_lengths(modality.data, modality.metadata), - self.batch_size, - exact=True, - ), - pin_memory=pin_memory_for(self.device), - ) embeddings = [None] * len(dataset) + def forward(model, batch): + video_ids = batch["id"].long() + _ = model(batch["data"].permute(0, 2, 1, 3, 4)) + values = self.features + if isinstance(values, tuple): + values = values[0] + vectors = torch.flatten(values, 1).detach().float().cpu().numpy() + for video_id, vector in zip(video_ids.cpu().tolist(), vectors): + embeddings[video_id] = vector + with inference_context(self.device): - for batch in dataloader: - batch = move_batch_to_device(batch, self.device) - video_ids = batch["id"].long() - frames = batch["data"].permute(0, 2, 1, 3, 4) - _ = self.model(frames) - values = self.features - if isinstance(values, tuple): - values = values[0] - vectors = torch.flatten(values, 1).detach().float().cpu().numpy() - for video_id, vector in zip(video_ids.cpu().tolist(), vectors): - embeddings[video_id] = vector + self._run_batches( + dataset, + get_sequence_lengths(modality.data, modality.metadata), + forward, + ) if self.output_file is not None: save_embeddings(embeddings, self.output_file) diff --git a/src/main/python/systemds/scuro/utils/memory_utility.py b/src/main/python/systemds/scuro/utils/memory_utility.py index ece738d050e..a28f14b5004 100644 --- a/src/main/python/systemds/scuro/utils/memory_utility.py +++ b/src/main/python/systemds/scuro/utils/memory_utility.py @@ -38,10 +38,22 @@ _CPU_MEMORY_BUDGET_GB = os.environ.get("SCURO_CPU_MEMORY_BUDGET_GB") -def cpu_memory_budget_bytes() -> float: +def cpu_memory_budget_bytes(modality_type: Optional[str] = None) -> float: + """CPU memory budget for one executor run. + + SCURO_CPU_MEMORY_BUDGET_GB_ (e.g. ..._AUDIO) overrides + SCURO_CPU_MEMORY_BUDGET_GB for executors over that modality type, so a + modality whose memory estimates are far below actual use can be throttled + without slowing down every other modality. + """ budget = float(psutil.virtual_memory().available) * _CPU_MEMORY_BUDGET_FRACTION - if _CPU_MEMORY_BUDGET_GB: - budget = min(budget, float(_CPU_MEMORY_BUDGET_GB) * 1024**3) + limit = _CPU_MEMORY_BUDGET_GB + if modality_type: + limit = os.environ.get( + f"SCURO_CPU_MEMORY_BUDGET_GB_{str(modality_type).upper()}", limit + ) + if limit: + budget = min(budget, float(limit) * 1024**3) return budget From d7d62f166c495c428f2cc8a49a44a2e1354d8ca5 Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Mon, 12 Oct 2026 08:32:10 +0200 Subject: [PATCH 15/15] fix merge issues --- .../scuro/drsearch/hyperparameter_tuner.py | 7 +- .../systemds/scuro/drsearch/node_executor.py | 2 +- .../scuro/drsearch/operator_registry.py | 13 ++- .../scuro/drsearch/representation_dag.py | 93 ++++++------------- .../scuro/modality/unimodal_modality.py | 4 +- .../representations/swin_video_transformer.py | 1 + .../systemds/scuro/representations/x3d.py | 1 + .../tests/scuro/test_operator_registry.py | 48 +++++----- 8 files changed, 78 insertions(+), 91 deletions(-) diff --git a/src/main/python/systemds/scuro/drsearch/hyperparameter_tuner.py b/src/main/python/systemds/scuro/drsearch/hyperparameter_tuner.py index aa806c26aa4..bfb62bcc890 100644 --- a/src/main/python/systemds/scuro/drsearch/hyperparameter_tuner.py +++ b/src/main/python/systemds/scuro/drsearch/hyperparameter_tuner.py @@ -1110,7 +1110,12 @@ def _evaluate_dag_once( dag_copy, modalities, task, node_cache=node_cache ).evaluate() return params, evaluation.scores - modified_modality = dag_copy.execute(modalities, task, node_cache=node_cache) + modified_modality = dag_copy.execute( + modalities, + task, + node_cache=node_cache, + execution_cache=self.execution_cache, + ) return params, task.run(modified_modality.data) def _evaluate_dag_trial( diff --git a/src/main/python/systemds/scuro/drsearch/node_executor.py b/src/main/python/systemds/scuro/drsearch/node_executor.py index 3219b37f279..a498c2752e3 100644 --- a/src/main/python/systemds/scuro/drsearch/node_executor.py +++ b/src/main/python/systemds/scuro/drsearch/node_executor.py @@ -286,7 +286,7 @@ def _run_task(): return { "scores": result[0], "task_time": result[1], - "task_timing": task.last_run_timing, + "task_timing": getattr(task, "last_run_timing", {}), "memory": measurement, "peak_bytes": measurement.increment_bytes if measurement else -1, "gpu_peak_bytes": gpu_peak_bytes, diff --git a/src/main/python/systemds/scuro/drsearch/operator_registry.py b/src/main/python/systemds/scuro/drsearch/operator_registry.py index c3780590434..f8a83b28e59 100644 --- a/src/main/python/systemds/scuro/drsearch/operator_registry.py +++ b/src/main/python/systemds/scuro/drsearch/operator_registry.py @@ -27,13 +27,22 @@ _BUILTIN_MODULES = { ModalityType.AUDIO: ( "covarep_audio_features", + "image_bind", "mel_spectrogram", "mfcc", "spectrogram", "timeseries_representations", "wav2vec", ), - ModalityType.TEXT: ("bert", "bow", "clip", "glove", "tfidf", "word2vec"), + ModalityType.TEXT: ( + "bert", + "bow", + "clip", + "glove", + "image_bind", + "tfidf", + "word2vec", + ), ModalityType.IMAGE: ("clip", "color_histogram", "openface", "resnet", "vgg"), ModalityType.VIDEO: ( "openface", @@ -56,7 +65,7 @@ ModalityType.AUDIO: ("window_aggregation",), ModalityType.TEXT: ("text_context", "text_context_with_indices"), ModalityType.VIDEO: ("window_aggregation",), - ModalityType.TIMESERIES: ("window_aggregation",), + ModalityType.TIMESERIES: ("physiological_window", "window_aggregation"), ModalityType.PHYSIOLOGICAL: ("physiological_window", "window_aggregation"), } _FUSION_MODULES = ( diff --git a/src/main/python/systemds/scuro/drsearch/representation_dag.py b/src/main/python/systemds/scuro/drsearch/representation_dag.py index 2eecd3d14d0..d84d092bbaf 100644 --- a/src/main/python/systemds/scuro/drsearch/representation_dag.py +++ b/src/main/python/systemds/scuro/drsearch/representation_dag.py @@ -390,72 +390,35 @@ def compute_node(node_id: str, task) -> TransformedModality: input_mods = [execute_node(input_id, task) for input_id in node.inputs] is_unimodal = len(input_mods) == 1 - if external_cache is not None and is_unimodal: - cached = external_cache.get(node_id) - if cached is not None: - result = cached - else: - 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 len(input_mods) == 1: - # It's a unimodal operation - if isinstance(node_operation, Context): - result = input_mods[0].context(node_operation) - elif isinstance(node_operation, DimensionalityReduction): - result = input_mods[0].dimensionality_reduction( - node_operation + 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 len(input_mods) == 1: + # It's a unimodal operation + if isinstance(node_operation, Context): + result = input_mods[0].context(node_operation) + elif isinstance(node_operation, DimensionalityReduction): + result = input_mods[0].dimensionality_reduction(node_operation) + elif isinstance(node_operation, AggregatedRepresentation): + result = node_operation.transform(input_mods[0]) + elif isinstance(node_operation, UnimodalRepresentation): + if rep_cache is not None: + result = rep_cache[node_operation.name] + else: + agg = pushdown_aggregation_for_node(node.parameters) + result = input_mods[0].apply_representation( + node_operation, aggregation=agg ) - elif isinstance(node_operation, AggregatedRepresentation): - result = node_operation.transform(input_mods[0]) - elif isinstance(node_operation, UnimodalRepresentation): - if rep_cache is not None: - result = rep_cache[node_operation.name] - else: - agg = pushdown_aggregation_for_node(node.parameters) - result = input_mods[0].apply_representation( - node_operation, aggregation=agg - ) - else: - # It's a fusion operation - fusion_op = node_operation - if getattr(fusion_op, "needs_training", False): - raise RuntimeError( - "Fitted representations require SplitAwareDAGEvaluator" - ) - result = input_mods[0].combine(input_mods[1:], fusion_op) - else: - - 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 len(input_mods) == 1: - # It's a unimodal operation - if isinstance(node_operation, Context): - result = input_mods[0].context(node_operation) - elif isinstance(node_operation, DimensionalityReduction): - result = input_mods[0].dimensionality_reduction( - node_operation - ) - elif isinstance(node_operation, AggregatedRepresentation): - result = node_operation.transform(input_mods[0]) - elif isinstance(node_operation, UnimodalRepresentation): - if rep_cache is not None: - result = rep_cache[node_operation.name] - else: - agg = pushdown_aggregation_for_node(node.parameters) - result = input_mods[0].apply_representation( - node_operation, aggregation=agg - ) - else: - # It's a fusion operation - fusion_op = node_operation - if getattr(fusion_op, "needs_training", False): - raise RuntimeError( - "Fitted representations require SplitAwareDAGEvaluator" - ) - result = input_mods[0].combine(input_mods[1:], fusion_op) + else: + # It's a fusion operation + fusion_op = node_operation + if getattr(fusion_op, "needs_training", False): + raise RuntimeError( + "Fitted representations require SplitAwareDAGEvaluator" + ) + result = input_mods[0].combine(input_mods[1:], fusion_op) + return result if execution_cache is not None and is_unimodal: result = execution_cache.get_or_compute_representation( diff --git a/src/main/python/systemds/scuro/modality/unimodal_modality.py b/src/main/python/systemds/scuro/modality/unimodal_modality.py index a615035bce3..7a40cd59360 100644 --- a/src/main/python/systemds/scuro/modality/unimodal_modality.py +++ b/src/main/python/systemds/scuro/modality/unimodal_modality.py @@ -78,7 +78,9 @@ def _subset_loader(self, indices): return subset_loader def subset(self, indices): - if self.data_loader.chunk_size is None or self.has_data(): + if self.data_loader.chunk_size is None or ( + self.has_data() and len(self.data) == len(self.data_loader.indices) + ): subset_modality = super().subset(indices) subset_modality.data_loader = self._subset_loader(indices) subset_modality.stats = subset_modality.data_loader.stats 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 2ea010651ab..5bf4a6d0add 100644 --- a/src/main/python/systemds/scuro/representations/swin_video_transformer.py +++ b/src/main/python/systemds/scuro/representations/swin_video_transformer.py @@ -53,6 +53,7 @@ @register_representation([ModalityType.VIDEO]) +@register_expensive_representation([ModalityType.VIDEO]) class SwinVideoTransformer(FrameBudgetedVideoModel, UnimodalRepresentation): _EMBED_DIM = 768 cache_in_worker = True diff --git a/src/main/python/systemds/scuro/representations/x3d.py b/src/main/python/systemds/scuro/representations/x3d.py index eb4cd4d08f0..d3f6b50e16b 100644 --- a/src/main/python/systemds/scuro/representations/x3d.py +++ b/src/main/python/systemds/scuro/representations/x3d.py @@ -59,6 +59,7 @@ def forward(self, input_: torch.Tensor) -> torch.Tensor: @register_representation([ModalityType.VIDEO]) +@register_expensive_representation([ModalityType.VIDEO]) class X3D(FrameBudgetedVideoModel, UnimodalRepresentation): cache_in_worker = True _MEMORY_PROFILES = { diff --git a/src/main/python/tests/scuro/test_operator_registry.py b/src/main/python/tests/scuro/test_operator_registry.py index 99bfeb2f2a4..d60788ea3f1 100644 --- a/src/main/python/tests/scuro/test_operator_registry.py +++ b/src/main/python/tests/scuro/test_operator_registry.py @@ -89,30 +89,36 @@ class TestOperatorRegistry(unittest.TestCase): def test_audio_representations_in_registry(self): registry = Registry() - assert registry.get_representations(ModalityType.AUDIO) == [ - MelSpectrogram, - MFCC, - ImageBind, - Spectrogram, - Wav2Vec, - Spectral, - ZeroCrossing, - RMSE, - Pitch, - ] + self.assertCountEqual( + registry.get_representations(ModalityType.AUDIO), + [ + MelSpectrogram, + MFCC, + ImageBind, + Spectrogram, + Wav2Vec, + Spectral, + ZeroCrossing, + RMSE, + Pitch, + ], + ) def test_video_representations_in_registry(self): registry = Registry() - assert registry.get_representations(ModalityType.VIDEO) == [ - ResNet, - OpenFace, - ImageBind, - SwinVideoTransformer, - X3D, - ColorHistogram, - VGG19, - CLIPVisual, - ] + self.assertCountEqual( + registry.get_representations(ModalityType.VIDEO), + [ + ResNet, + OpenFace, + ImageBind, + SwinVideoTransformer, + X3D, + ColorHistogram, + VGG19, + CLIPVisual, + ], + ) def test_timeseries_representations_in_registry(self): registry = Registry()