From 2651d489998e35bc4d0a5df031a39fc712a9de39 Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Tue, 1 Sep 2026 11:13:49 +0200 Subject: [PATCH 01/17] 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/17] 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/17] 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/17] 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/17] 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/17] 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/17] 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/17] 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/17] 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/17] 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/17] 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/17] 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/17] 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/17] 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 40893aeecbc1be4828c450621491c151c75a19e3 Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Mon, 12 Oct 2026 11:09:20 +0200 Subject: [PATCH 15/17] add algebra rules --- .../python/systemds/scuro/drsearch/algebra.py | 577 ++++++++++++++++++ .../scuro/drsearch/hyperparameter_tuner.py | 123 +++- .../scuro/drsearch/modality_shared_memory.py | 25 +- .../scuro/drsearch/multimodal_ga_optimizer.py | 233 ++++++- .../scuro/drsearch/multimodal_optimizer.py | 77 ++- .../systemds/scuro/drsearch/node_executor.py | 37 +- .../systemds/scuro/drsearch/node_scheduler.py | 69 ++- .../scuro/drsearch/representation_dag.py | 76 ++- .../scuro/drsearch/split_aware_evaluator.py | 22 +- .../scuro/drsearch/unimodal_optimizer.py | 19 +- .../systemds/scuro/modality/transformed.py | 7 +- .../python/systemds/scuro/models/model.py | 1 + .../scuro/representations/aggregate.py | 12 +- .../systemds/scuro/representations/average.py | 4 +- .../systemds/scuro/representations/clip.py | 3 +- .../scuro/representations/concatenation.py | 8 + .../systemds/scuro/representations/fusion.py | 12 +- .../systemds/scuro/representations/glove.py | 4 +- .../scuro/representations/hadamard.py | 4 +- .../systemds/scuro/representations/lstm.py | 2 +- .../systemds/scuro/representations/max.py | 15 +- .../scuro/representations/openface.py | 157 +++-- .../systemds/scuro/representations/sum.py | 19 +- .../representations/swin_video_transformer.py | 3 +- .../systemds/scuro/representations/vgg.py | 3 +- .../representations/window_aggregation.py | 9 +- .../scuro/representations/zero_pad.py | 48 ++ .../tests/scuro/test_algebra_canonical.py | 269 ++++++++ .../python/tests/scuro/test_algebra_cse.py | 82 +++ .../tests/scuro/test_algebra_exhaustive.py | 76 +++ .../tests/scuro/test_algebra_explicit_prep.py | 218 +++++++ .../tests/scuro/test_algebra_ga_dedup.py | 65 ++ .../python/tests/scuro/test_algebra_laws.py | 136 +++++ .../scuro/test_algebra_order_invariant.py | 56 ++ .../tests/scuro/test_algebra_plan_counts.py | 71 +++ .../tests/scuro/test_algebra_signatures.py | 86 +++ .../python/tests/scuro/test_algebra_stats.py | 52 ++ .../scuro/test_algebra_subtree_sharing.py | 115 ++++ .../tests/scuro/test_algebra_typecheck.py | 121 ++++ .../python/tests/scuro/test_data_loaders.py | 34 ++ .../scuro/test_multimodal_ga_optimizer.py | 74 ++- .../tests/scuro/test_unimodal_optimizer.py | 184 ++++++ 42 files changed, 3049 insertions(+), 159 deletions(-) create mode 100644 src/main/python/systemds/scuro/drsearch/algebra.py create mode 100644 src/main/python/systemds/scuro/representations/zero_pad.py create mode 100644 src/main/python/tests/scuro/test_algebra_canonical.py create mode 100644 src/main/python/tests/scuro/test_algebra_cse.py create mode 100644 src/main/python/tests/scuro/test_algebra_exhaustive.py create mode 100644 src/main/python/tests/scuro/test_algebra_explicit_prep.py create mode 100644 src/main/python/tests/scuro/test_algebra_ga_dedup.py create mode 100644 src/main/python/tests/scuro/test_algebra_laws.py create mode 100644 src/main/python/tests/scuro/test_algebra_order_invariant.py create mode 100644 src/main/python/tests/scuro/test_algebra_plan_counts.py create mode 100644 src/main/python/tests/scuro/test_algebra_signatures.py create mode 100644 src/main/python/tests/scuro/test_algebra_stats.py create mode 100644 src/main/python/tests/scuro/test_algebra_subtree_sharing.py create mode 100644 src/main/python/tests/scuro/test_algebra_typecheck.py diff --git a/src/main/python/systemds/scuro/drsearch/algebra.py b/src/main/python/systemds/scuro/drsearch/algebra.py new file mode 100644 index 00000000000..880e7afe7ea --- /dev/null +++ b/src/main/python/systemds/scuro/drsearch/algebra.py @@ -0,0 +1,577 @@ +# ------------------------------------------------------------- +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# +# ------------------------------------------------------------- +import hashlib +import itertools +import json +from collections import Counter +from dataclasses import dataclass +from fractions import Fraction +from functools import lru_cache +from math import comb, factorial, prod +from typing import NamedTuple, Optional + +import numpy as np + +from systemds.scuro.drsearch.operator_registry import Registry, load_builtin_operators +from systemds.scuro.modality.type import ModalityType +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.fusion import Fusion +from systemds.scuro.representations.representation import ( + FitRequirement, + RepresentationStats, +) +from systemds.scuro.representations.zero_pad import ZeroPad + +ALL_LAWS = frozenset( + { + "comm", + "assoc", + "avg_flatten", + "concat_perm", + "agg_pushdown", + "implicit_agg", + "implicit_pad", + } +) + +EXACT_LAWS = frozenset({"agg_pushdown", "implicit_agg", "implicit_pad"}) + + +@dataclass(frozen=True) +class AlgebraConfig: + """Switches for the Beta-algebra rules; the defaults are the legacy search.""" + + laws: frozenset = frozenset() + feature_order_invariant: Optional[bool] = None + cse_law_aware: bool = False + ga_canonical_dedup: bool = False + exhaustive_dedup_permutations: bool = False + exhaustive_canonical: bool = False + typecheck: str = "off" # off | warn | strict + explicit_fusion_prep: bool = False + explicit_padding: str = "off" # off | pad | reject + typed_chaining: bool = False + share_fused_subtrees: bool = False + + @classmethod + def full(cls) -> "AlgebraConfig": + return cls( + laws=ALL_LAWS, + feature_order_invariant=True, + cse_law_aware=True, + ga_canonical_dedup=True, + exhaustive_dedup_permutations=True, + exhaustive_canonical=True, + typecheck="warn", + explicit_fusion_prep=True, + explicit_padding="pad", + share_fused_subtrees=True, + ) + + +class OperatorSignature(NamedTuple): + kind: str # encode | context | aggregate | reduce | fuse | task + input_types: frozenset + output_type: Optional[ModalityType] # None: the input's type + arity: object # 1, or "n" for fusion + input_rank: str # any | seq | windows | vector + output_rank: str + + +ANY_TYPE = frozenset(m for m in ModalityType if isinstance(m.value, int)) + + +def operator_signature(op_cls, parameters=None) -> OperatorSignature: + """`op_cls` None is a task node.""" + if op_cls is None: + return OperatorSignature("task", ANY_TYPE, None, 1, "any", "any") + if issubclass(op_cls, AggregatedRepresentation): + return OperatorSignature("aggregate", ANY_TYPE, None, 1, "seq", "vector") + if issubclass(op_cls, ZeroPad): + return OperatorSignature("reduce", ANY_TYPE, None, 1, "vector", "vector") + if issubclass(op_cls, Fusion): + keep_axis = (parameters or {}).get("preserve_leading_axis") + seq_in = op_cls.preserves_leading_axis if keep_axis is None else keep_axis + ranks = ["seq" if seq else "vector" for seq in (seq_in, keep_axis)] + return OperatorSignature("fuse", ANY_TYPE, _output_type(op_cls), "n", *ranks) + if issubclass(op_cls, Context): + kind, lists, ranks = "context", Registry._context_operators, ("seq", "windows") + elif issubclass(op_cls, DimensionalityReduction): + kind, lists = "reduce", Registry._dimensionality_reduction_operators + ranks = ("vector", "vector") + else: + kind, lists, ranks = "encode", Registry._representations, ("any", "any") + load_builtin_operators(ANY_TYPE) + input_types = frozenset(m for m, ops in lists.items() if op_cls in ops) + return OperatorSignature(kind, input_types, _output_type(op_cls), 1, *ranks) + + +@lru_cache(maxsize=None) +def _needs_alignment(op_cls) -> bool: + return op_cls().needs_alignment + + +@lru_cache(maxsize=None) +def _output_type(op_cls): + return getattr(op_cls, "OUTPUT_MODALITY_TYPE", None) or getattr( + op_cls(), "output_modality_type", None + ) + + +class Violation(NamedTuple): + severity: str # error | implicit + kind: str # input_type | arity | hidden_aggregation | zero_padding + node_id: str + operator: str + detail: str + + +def typecheck(dag, leaf_types, leaf_stats=None) -> list: + nodes = {node.node_id: node for node in dag.nodes} + types, violations = {}, [] + stats = output_stats(nodes.__getitem__, leaf_stats) + + def key(node): + return node.modality_id, node.representation_index + + def flag(severity, kind, node, detail): + name = node.operation.__name__ + violations.append(Violation(severity, kind, node.node_id, name, detail)) + + def check_fusion(node): + operands = [stats(i) for i in node.inputs] + if None in operands: + return + op = node.operation(params=node.parameters) + if not op.preserves_leading_axis: + ranks = [len(s.output_shape) for s in operands] + if max(ranks) > 1: + flag("implicit", "hidden_aggregation", node, f"operand ranks {ranks}") + if op.needs_alignment: + widths = [s.output_shape[-1] for s in op._fusion_input_stats(operands)] + if len(set(widths)) > 1: + flag("implicit", "zero_padding", node, f"operand widths {widths}") + + def out_type(nid): + if nid in types: + return types[nid] + node = nodes[nid] + if not node.inputs: + types[nid] = leaf_types.get(key(node)) + return types[nid] + in_types = [out_type(i) for i in node.inputs] + if node.operation is None: + types[nid] = in_types[0] + return types[nid] + sig = operator_signature(node.operation, node.parameters) + if (sig.arity == "n") != (len(node.inputs) > 1): + flag("error", "arity", node, f"{len(node.inputs)} input(s)") + for t in in_types: + if t is not None and t not in sig.input_types: + flag("error", "input_type", node, f"{t.name} input") + if sig.kind == "fuse" and leaf_stats is not None: + check_fusion(node) + types[nid] = sig.output_type or in_types[0] + return types[nid] + + out_type(dag.root_node_id) + return violations + + +def output_stats(get_node, leaf_stats): + from systemds.scuro.drsearch.representation_dag import ( + pushdown_aggregation_for_node, + ) + + memo = {} + + def stats(nid): + if nid not in memo: + node = get_node(nid) + memo[nid] = None + if not node.inputs: + key = node.modality_id, node.representation_index + memo[nid] = (leaf_stats or {}).get(key) + elif node.operation is not None: + inputs = [stats(i) for i in node.inputs] + if None not in inputs: + try: + out = node.operation(params=node.parameters).get_output_stats( + inputs if len(inputs) > 1 else inputs[0] + ) + agg = pushdown_aggregation_for_node(node.parameters) + memo[nid] = agg.get_output_stats(out) if agg else out + except Exception: + pass + return memo[nid] + + return stats + + +def sample_stats(modality) -> Optional[RepresentationStats]: + 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 explicit_prep_on(algebra) -> bool: + return algebra.explicit_fusion_prep or algebra.explicit_padding == "pad" + + +def explicit_fusion_operands( + builder, operation, parameters, inputs, leaf_stats, algebra +) -> list: + if not issubclass(operation, Fusion): + return inputs + op = operation(params=parameters) + stats = output_stats(builder.get_node, leaf_stats) + if algebra.explicit_fusion_prep and not op.preserves_leading_axis: + agg = AggregatedRepresentation(target_dimensions=1, aggregate_leading=True) + inputs = [ + ( + builder.create_operation_node( + AggregatedRepresentation, [i], agg.get_current_parameters() + ) + if stats(i) is not None and len(stats(i).output_shape) > 1 + else i + ) + for i in inputs + ] + shapes = [stats(i) for i in inputs] + if ( + algebra.explicit_padding == "pad" + and op.needs_alignment + and None not in shapes + and all(len(s.output_shape) == 1 for s in shapes) + ): + width = max(s.output_shape[0] for s in shapes) + inputs = [ + ( + builder.create_operation_node( + ZeroPad, + [i], + ZeroPad(width, s.output_shape[0]).get_current_parameters(), + ) + if s.output_shape[0] < width + else i + ) + for i, s in zip(inputs, shapes) + ] + return inputs + + +def count_explicit_prep(counter, dag) -> None: + """Counts the explicit operand preparation in a plan about to be evaluated.""" + nodes = {n.node_id: n for n in dag.filter_connected_nodes(dag.nodes)} + padded = False + for node in nodes.values(): + if not (node.operation and issubclass(node.operation, Fusion)): + continue + for operand in (nodes[i] for i in node.inputs): + if operand.operation is ZeroPad: + padded = True + ratio = operand.parameters["width"] / operand.parameters["input_width"] + counter["explicit_paddings"] += 1 + counter["explicit_padding_ratio_sum"] += ratio + counter["explicit_padding_ratio_max"] = max( + counter["explicit_padding_ratio_max"], ratio + ) + operand = nodes[operand.inputs[0]] + if operand.operation is AggregatedRepresentation: + counter["explicit_aggregations"] += 1 + counter["explicit_prep_plans"] += 1 + counter["explicit_padding_plans"] += padded + + +def count_violations(counter, violations) -> list: + """Adds per-kind and per-operator counts; returns the errors.""" + counter["typecheck_plans"] += 1 + for v in violations: + counter[f"typecheck_{v.kind}"] += 1 + counter[f"typecheck_{v.kind}:{v.operator}"] += 1 + return [v for v in violations if v.severity == "error"] + + +def order_invariant(config: AlgebraConfig, tasks) -> bool: + if config.feature_order_invariant is not None: + return config.feature_order_invariant + return bool(tasks) and all( + getattr(task.model, "feature_order_invariant", None) is True for task in tasks + ) + + +def canonical_signature(dag, laws, invariant, node_id=None) -> tuple: + # representation_dag imports this module. + from systemds.scuro.drsearch.representation_dag import _spec_safe + + def view(nid): + node = dag.get_node_by_id(nid) + op = node.operation + params = {k: _spec_safe(v) for k, v in (node.parameters or {}).items()} + name = op.__name__ if op is not None else None + return ( + name, + op, + params, + node.inputs, + node.modality_id, + node.representation_index, + ) + + return _signature(view, node_id or dag.root_node_id, laws, invariant) + + +def canonical_signature_from_spec(spec: dict, laws, invariant) -> tuple: + nodes = {node["id"]: node for node in spec["nodes"]} + classes = _fusion_classes() + + def view(nid): + n = nodes[nid] + return ( + n["op"], + classes.get(n["op"]), + n["params"], + n["inputs"], + n["modality_id"], + n["representation_index"], + ) + + return _signature(view, spec["root"], laws, invariant) + + +def canonical_id(dag, laws, invariant, node_id=None) -> str: + """`dag` is a RepresentationDag or a to_spec() dict.""" + if isinstance(dag, dict): + sig = canonical_signature_from_spec(dag, laws, invariant) + else: + sig = canonical_signature(dag, laws, invariant, node_id) + blob = json.dumps(sig, sort_keys=True, separators=(",", ":")) + return hashlib.sha1(blob.encode()).hexdigest()[:16] + + +def _fusion_classes() -> dict: + classes = {cls.__name__: cls for cls in Registry().get_fusion_operators()} + stack = [Fusion] + while stack: + cls = stack.pop() + classes.setdefault(cls.__name__, cls) + stack.extend(cls.__subclasses__()) + return classes + + +def _params_key(params) -> str: + return json.dumps(params, sort_keys=True, separators=(",", ":")) + + +def _signature(view, root, laws, invariant) -> tuple: + # The aggregation Fusion._prepare_modalities applies to every operand of a + # fusion that does not preserve the leading axis. + implicit_agg = ( + "op", + AggregatedRepresentation.__name__, + _params_key( + AggregatedRepresentation( + target_dimensions=1, aggregate_leading=True + ).get_current_parameters() + ), + ) + memo = {} + + def sig(nid, laws): + if (nid, laws) not in memo: + memo[nid, laws] = compute(nid, laws) + return memo[nid, laws] + + def compute(nid, laws): + name, op, params, inputs, modality_id, representation_index = view(nid) + if not inputs: + return ("leaf", modality_id, representation_index) + + children = [sig(i, laws & EXACT_LAWS if _fitted(op) else laws) for i in inputs] + params = dict(params) + pushed = ( + params.pop("_pushdown_aggregation", None) + if "agg_pushdown" in laws + else None + ) + head = ("op", name, _params_key(params)) + fusion = isinstance(op, type) and issubclass(op, Fusion) + op_laws = op.algebraic_laws(params) & laws if fusion else frozenset() + + if ( + fusion + and "implicit_agg" in laws + and not params.get("preserve_leading_axis", op.preserves_leading_axis) + ): + children = [c[3][0] if c[:3] == implicit_agg else c for c in children] + if fusion and "implicit_pad" in laws and _needs_alignment(op): + children = [ + c[3][0] if c[:2] == ("op", ZeroPad.__name__) else c for c in children + ] + if "assoc" in op_laws: + children = [g for c in children for g in (c[3] if c[:3] == head else (c,))] + if "avg_flatten" in op_laws: + weights = Counter() + for c in children: + for w, g in c[3] if c[:3] == head else (("1", c),): + weights[g] += Fraction(w) / len(children) + children = sorted(((str(w), g) for g, w in weights.items()), key=repr) + elif "comm" in op_laws or ("concat_perm" in op_laws and invariant): + children = sorted(children, key=repr) + + out = head + (tuple(children),) + if pushed is not None: + out = ("op", AggregatedRepresentation.__name__, _params_key(pushed), (out,)) + return out + + return sig(root, frozenset(laws)) + + +def count_fusion_plans(n, fusion_ops, laws, invariant) -> dict: + heads = _fusion_heads(fusion_ops, laws) + f = len(heads) + reorders = sum(_reorders(op_laws, invariant) for *_, op_laws in heads) + ordered = factorial(2 * n - 2) // factorial(n - 1) * f ** (n - 1) + if n == 1: + canonical = 1 + elif reorders == f and all("assoc" in op_laws for *_, op_laws in heads): + canonical = f * _flattened_trees(n, f) + elif n <= 4: + canonical = _brute_force_canonical(n, fusion_ops, laws, invariant) + else: + canonical = None + return { + "ordered": ordered, + "legacy_exhaustive": factorial(n) * ordered, + "commutative": _commutative_plans(n, heads, invariant), + "canonical": canonical, + } + + +def count_search_space( + ks, fusion_ops, laws, invariant, min_modalities=2, max_modalities=None +) -> dict: + """`ks[i]` is the number of representations of modality i.""" + totals = dict.fromkeys( + ("ordered", "legacy_exhaustive", "commutative", "canonical"), 0 + ) + for r in range(min_modalities, (max_modalities or len(ks)) + 1): + leaves = sum(prod(c) for c in itertools.combinations(ks, r)) + for key, plans in count_fusion_plans(r, fusion_ops, laws, invariant).items(): + if plans is None or totals[key] is None: + totals[key] = None + else: + totals[key] += leaves * plans + return totals + + +def _fitted(op) -> bool: + return getattr(op, "fit_requirement", None) not in (None, FitRequirement.NONE) + + +def _commutative_plans(n, heads, invariant) -> int: + plans = {(1, True): 1, (1, False): 1} + for m in range(2, n + 1): + for free in (True, False): + plans[m, free] = 0 + for *_, op, _, op_laws in heads: + below = free and not _fitted(op) + splits = sum( + comb(m, k) * plans[k, below] * plans[m - k, below] + for k in range(1, m) + ) + plans[m, free] += splits // ( + 2 if free and _reorders(op_laws, invariant) else 1 + ) + return plans[n, True] + + +def _reorders(op_laws, invariant) -> bool: + return "comm" in op_laws or ("concat_perm" in op_laws and invariant) + + +def _fusion_heads(fusion_ops, laws) -> list: + from systemds.scuro.drsearch.representation_dag import _spec_safe + + heads = [] + for op in fusion_ops: + params = {k: _spec_safe(v) for k, v in op().get_current_parameters().items()} + heads.append((op.__name__, op, params, op.algebraic_laws(params) & laws)) + return heads + + +def _flattened_trees(n, f) -> int: + tree, block, partitions = [0, 1], [0, 1], [1, 1] + for m in range(2, n + 1): + tree.append( + sum(comb(m - 1, j - 1) * block[j] * partitions[m - j] for j in range(1, m)) + ) + block.append((f - 1) * tree[m]) + partitions.append(tree[m] + block[m]) + return tree[n] + + +def _ordered_trees(leaves): + if len(leaves) == 1: + yield leaves[0] + return + for split in range(1, len(leaves)): + for left in itertools.combinations(leaves, split): + right = tuple(i for i in leaves if i not in left) + for l_tree in _ordered_trees(left): + for r_tree in _ordered_trees(right): + yield (l_tree, r_tree) + + +def _brute_force_canonical(n, fusion_ops, laws, invariant) -> int: + heads = _fusion_heads(fusion_ops, laws) + signatures = set() + for tree in _ordered_trees(tuple(range(n))): + for choice in itertools.product(heads, repeat=n - 1): + nodes, ops = {}, iter(choice) + + def add(t): + if isinstance(t, int): + nodes[t] = (None, None, {}, (), t, 0) + return t + inputs = (add(t[0]), add(t[1])) + name, op, params, _ = next(ops) + nid = f"n{len(nodes)}" + nodes[nid] = (name, op, params, inputs, None, None) + return nid + + root = add(tree) + signatures.add(_signature(nodes.__getitem__, root, laws, invariant)) + return len(signatures) diff --git a/src/main/python/systemds/scuro/drsearch/hyperparameter_tuner.py b/src/main/python/systemds/scuro/drsearch/hyperparameter_tuner.py index bfb62bcc890..7805ac66b3f 100644 --- a/src/main/python/systemds/scuro/drsearch/hyperparameter_tuner.py +++ b/src/main/python/systemds/scuro/drsearch/hyperparameter_tuner.py @@ -32,7 +32,10 @@ import random from systemds.scuro.drsearch.representation_dag import ( RepresentationNode, + get_modality_by_id_and_instance_id, ) +from systemds.scuro.representations.representation import infer_stats_from_data +from systemds.scuro.utils.memory_utility import cpu_memory_budget_bytes from systemds.scuro.modality.modality import Modality from systemds.scuro.drsearch.split_aware_evaluator import ( SplitAwareDAGEvaluator, @@ -151,6 +154,65 @@ def _is_window_operation(op: Any) -> bool: return False +def _is_leading_axis_concatenation(node: RepresentationNode) -> bool: + from systemds.scuro.representations.concatenation import Concatenation + + return ( + inspect.isclass(node.operation) + and issubclass(node.operation, Concatenation) + and bool((node.parameters or {}).get("preserve_leading_axis")) + ) + + +def _tied_window_params(dag) -> Dict[str, str]: + parent = {} + + def find(node_id): + while parent[node_id] != node_id: + node_id = parent[node_id] + return node_id + + def windows_under(node_id): + node = dag.get_node_by_id(node_id) + if _is_window_operation(node.operation): + return [node_id] + if _is_leading_axis_concatenation(node): + return [w for input_id in node.inputs for w in windows_under(input_id)] + return [] + + for node in dag.nodes: + if not _is_leading_axis_concatenation(node): + continue + windows = windows_under(node.node_id) + for window_id in windows: + parent.setdefault(window_id, window_id) + for window_id in windows[1:]: + parent[find(window_id)] = find(windows[0]) + + tied = {} + canonical = {} + for window_id in parent: + name = dag.get_node_by_id(window_id).operation.granularity_parameter + if name is None: + continue + key = (find(window_id), name) + canonical.setdefault(key, f"{window_id}-{name}") + if canonical[key] != f"{window_id}-{name}": + tied[f"{window_id}-{name}"] = canonical[key] + return tied + + +def _expand_tied_params( + config: Dict[str, Any], tied_params: Dict[str, str] +) -> Dict[str, Any]: + return { + **config, + **{ + alias: config[name] for alias, name in tied_params.items() if name in config + }, + } + + def _materialize_node_params( node: RepresentationNode, flat_params: Dict[str, Any] ) -> Dict[str, Any]: @@ -581,6 +643,8 @@ def visit_node(node_id): for node_id in hyperparams } param_specs = self._build_param_specs(hyperparams, node_parameters) + tied_params = _tied_window_params(dag) + param_specs = [s for s in param_specs if s["name"] not in tied_params] discrete_size = self._estimate_discrete_search_size(param_specs) n_calls = min(discrete_size, max_evals) if max_evals else discrete_size all_results = self._search_best_configs( @@ -593,6 +657,7 @@ def visit_node(node_id): budget=n_calls, initial_config=None, rep_name=rep_name, + tied_params=tied_params, ) all_results.append(baseline) @@ -924,6 +989,7 @@ def _search_best_configs( budget: int, initial_config: Optional[Dict[str, Any]], rep_name: str = "", + tied_params: Optional[Dict[str, str]] = None, ) -> List[Tuple[Dict[str, Any], Any]]: import optuna from optuna.trial import TrialState @@ -1002,12 +1068,26 @@ def _search_best_configs( trial_results: List[Tuple[Dict[str, Any], Any]] = [] def objective(trial: optuna.Trial) -> float: - config = self._suggest_config_from_specs(trial, param_specs) + config = _expand_tied_params( + self._suggest_config_from_specs(trial, param_specs), tied_params or {} + ) key = self._config_key(config) if key in seen: trial.set_user_attr("duplicate", True) raise optuna.TrialPruned() + estimate = self._estimate_trial_peak_bytes( + dag, config, modality_ids, task, modalities_override + ) + budget_bytes = cpu_memory_budget_bytes() + if estimate is not None and estimate > budget_bytes: + self.logger.warning( + f"{rep_name}: skipping trial estimated at {estimate / 1024**3:.1f} GB " + f"(budget {budget_bytes / 1024**3:.1f} GB): {config}" + ) + trial.set_user_attr("over_memory_budget", True) + raise optuna.TrialPruned() + params, scores = self.evaluate_dag_config( dag, config, @@ -1072,6 +1152,47 @@ def _get_cached_modalities_for_task(self, task, modality_ids): ) return cached_modalities + def _estimate_trial_peak_bytes( + self, dag, params, modality_ids, task, modalities_override=None + ) -> Optional[int]: + modalities = ( + modalities_override + if modalities_override is not None + else self.get_modalities_by_id(modality_ids) + ) + stats = {} + + def visit(node_id): + if node_id in stats: + return stats[node_id] + node = dag.get_node_by_id(node_id) + if not node.inputs: + modality = get_modality_by_id_and_instance_id( + modalities, int(node.modality_id), node.representation_index + ) + stats[node_id] = ( + modality.get_stats() + if hasattr(modality, "get_stats") + else infer_stats_from_data(modality.data) + ) + return stats[node_id] + inputs = [visit(input_id) for input_id in node.inputs] + if any(s is None for s in inputs): + stats[node_id] = None + return None + inputs = inputs[0] if len(inputs) == 1 else inputs + operation = node.operation(params=_apply_trial_params_to_node(node, params)) + peaks.append(operation.estimate_peak_memory_bytes(inputs)["cpu_peak_bytes"]) + stats[node_id] = operation.get_output_stats(inputs) + return stats[node_id] + + peaks = [] + output_stats = visit(dag.root_node_id) + if output_stats is None: + return None + peaks.append(task.estimate_peak_memory_bytes(output_stats)["cpu_peak_bytes"]) + return int(max(peaks)) + def _trial_node_cache(self, dag, params): store = getattr(self, "_trial_cache", None) if store is None: 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 b7472a35a6d..1f78ff67421 100644 --- a/src/main/python/systemds/scuro/drsearch/modality_shared_memory.py +++ b/src/main/python/systemds/scuro/drsearch/modality_shared_memory.py @@ -18,11 +18,28 @@ # under the License. # # ------------------------------------------------------------- +import os +import shutil from typing import Any, List, Tuple import numpy as np from multiprocessing import shared_memory SHARED_MEMORY_MIN_BYTES = 1 * 1024 * 1024 +_SHM_DIR = "/dev/shm" + + +def _shm_reserve_bytes() -> int: + return int(float(os.environ.get("SCURO_SHM_RESERVE_GB", "8")) * 1024**3) + + +def _should_offload(data_nbytes: int) -> bool: + if data_nbytes < SHARED_MEMORY_MIN_BYTES: + return False + try: + free = shutil.disk_usage(_SHM_DIR).free + except OSError: + return True + return data_nbytes + _shm_reserve_bytes() <= free def unlink_shm(name: str) -> None: @@ -324,7 +341,7 @@ def _is_nested_shared_memory_candidate(data: Any) -> bool: def add_shared_memory_candidate(data: Any, resident_bytes: int = 0) -> bool: if _is_shared_memory_candidate(data): data_nbytes = _estimate_ndarray_list_nbytes(data) - if data_nbytes >= SHARED_MEMORY_MIN_BYTES: + if _should_offload(data_nbytes): dtype = data[0].dtype offsets = [] total_elems = 0 @@ -351,7 +368,7 @@ def add_shared_memory_candidate(data: Any, resident_bytes: int = 0) -> bool: elif _is_shared_ndarray_candidate(data): arr = data data_nbytes = int(arr.nbytes) - if data_nbytes >= SHARED_MEMORY_MIN_BYTES: + if _should_offload(data_nbytes): shm = shared_memory.SharedMemory(create=True, size=data_nbytes) shm_arr = np.ndarray(arr.shape, dtype=arr.dtype, buffer=shm.buf) np.copyto(shm_arr, arr, casting="no") @@ -371,7 +388,7 @@ def add_shared_memory_candidate(data: Any, resident_bytes: int = 0) -> bool: leaves.append(arr) group_bounds.append((lo, len(leaves))) data_nbytes = _estimate_ndarray_list_nbytes(leaves) - if data_nbytes >= SHARED_MEMORY_MIN_BYTES: + if _should_offload(data_nbytes): dtype = leaves[0].dtype offsets = [] total_elems = 0 @@ -400,7 +417,7 @@ def add_shared_memory_candidate(data: Any, resident_bytes: int = 0) -> bool: elif _is_string_list_shared_memory_candidate(data): encoded = [s.encode("utf-8") for s in data] data_nbytes = int(sum(len(b) for b in encoded)) - if data_nbytes >= SHARED_MEMORY_MIN_BYTES: + if _should_offload(data_nbytes): shm = shared_memory.SharedMemory(create=True, size=data_nbytes) mv = shm.buf cursor = 0 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 79a04ca262c..0c182338101 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 +from collections import Counter import gc import os import pickle @@ -42,6 +43,18 @@ unlink_shm, ) +from systemds.scuro.drsearch.algebra import ( + AlgebraConfig, + canonical_id, + canonical_signature, + count_explicit_prep, + count_violations, + explicit_fusion_operands, + explicit_prep_on, + order_invariant, + sample_stats, + typecheck, +) from systemds.scuro.drsearch.operator_registry import Registry from systemds.scuro.drsearch.representation_dag import ( RepresentationDAGBuilder, @@ -53,6 +66,7 @@ SplitAwareDAGEvaluator, dag_requires_fitting, ) +from systemds.scuro.drsearch.trial_result_cache import TrialNodeCache, TrialResultCache from systemds.scuro.drsearch.worker_pool import ( PersistentWorkerPool, create_mp_context, @@ -62,7 +76,7 @@ AggregatedRepresentation, ) from systemds.scuro.representations.fusion import FusionInputTooLargeError -from systemds.scuro.representations.representation import RepresentationStats +from systemds.scuro.utils.memory_utility import gpu_memory_info from systemds.scuro.utils.schema_helpers import get_shape Tree = Union[int, Tuple["Tree", "Tree"]] @@ -175,6 +189,7 @@ class DagGenome: _TIMING_OBJECTIVES = {"runtime", "task_time", "representation_time"} +_MAX_ERROR_SAMPLES = 5 ObjectiveSpec = Tuple[str, str] # (name, "max" | "min") @@ -187,6 +202,14 @@ def _objective_value( return val_score[name] +def _task_uses_gpu(task: Task) -> bool: + uses_gpu = getattr(task.model, "uses_gpu", None) + if uses_gpu is not None: + return bool(uses_gpu) + device = getattr(task.model, "device", None) + return str(getattr(device, "type", device) or "").startswith("cuda") + + def _failure_fitness(objective_specs: List[ObjectiveSpec]) -> Tuple[float, ...]: return tuple( float("-inf") if direction == "max" else float("inf") @@ -205,15 +228,45 @@ def _release_cuda_memory() -> None: pass +_SUBTREE_CACHE: Optional[TrialResultCache] = None + + +def _subtree_cache() -> TrialResultCache: + global _SUBTREE_CACHE + if _SUBTREE_CACHE is None: + gb = float(os.environ.get("SCURO_GA_SUBTREE_CACHE_GB", "2")) + _SUBTREE_CACHE = TrialResultCache(int(gb * 1024**3)) + return _SUBTREE_CACHE + + +def _subtree_node_cache(dag: RepresentationDag, task: Task) -> TrialNodeCache: + signatures = { + node.node_id: ( + task.model.name, + canonical_signature(dag, frozenset(), False, node.node_id), + ) + for node in dag.nodes + if node.inputs + } + return TrialNodeCache(_subtree_cache(), signatures) + + def _evaluate_genome_body( dag: RepresentationDag, task: Task, modalities: List[Any], objective_specs: List[ObjectiveSpec], + share_subtrees: bool = False, ) -> Tuple[Optional[Tuple[float, ...]], Optional[Dict[str, Any]]]: start = time.time() if dag_requires_fitting(dag): - evaluation = SplitAwareDAGEvaluator(dag, modalities, task).evaluate() + node_cache = _subtree_node_cache(dag, task) if share_subtrees else None + if node_cache is not None: + store = node_cache.store + before = store.hits, store.misses, store.saved_seconds + evaluation = SplitAwareDAGEvaluator( + dag, modalities, task, node_cache=node_cache, share_fitted=share_subtrees + ).evaluate() scores = evaluation.scores total = evaluation.representation_time + evaluation.task_time val_score = scores[1].average_scores @@ -238,6 +291,13 @@ def _evaluate_genome_body( "task_timing": getattr(task, "last_run_timing", {}), **timing, } + if node_cache is not None: + payload["subtree_cache"] = { + "hits": store.hits - before[0], + "misses": store.misses - before[1], + "saved_seconds": store.saved_seconds - before[2], + "bytes": store.bytes, + } return fitness, payload fused = dag.execute(modalities, task, enable_cache=False) @@ -281,11 +341,11 @@ def _evaluate_genome_body( def _dispatch_genome_evaluation(payload, _gpu_id): - dag, task, modalities, objective_specs = payload + dag, task, modalities, objective_specs, *share = payload try: try: fitness, result = _evaluate_genome_body( - dag, task, modalities, objective_specs + dag, task, modalities, objective_specs, *share ) if fitness is None: fitness = _failure_fitness(objective_specs) @@ -310,6 +370,9 @@ def __init__(self, values, fitness_type): class MultimodalDeapOptimizer: + # For instances made without __init__, as regret_analysis/exhaustive_enumerator.py does. + algebra = AlgebraConfig() + def __init__( self, modalities: List[Any], @@ -336,6 +399,7 @@ def __init__( allow_repeated_modalities: bool = False, threads_per_worker: Optional[int] = None, max_feature_dim: Optional[int] = DEFAULT_MAX_FUSED_FEATURE_DIM, + algebra: Optional[AlgebraConfig] = None, ): self.modalities = modalities self.tasks = tasks @@ -343,6 +407,13 @@ def __init__( 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.algebra = algebra or AlgebraConfig() + self.algebra_stats = Counter() + self._order_invariant = order_invariant(self.algebra, tasks) + self._canonical_signatures: Dict[Tuple, Tuple] = {} + self._seen_legacy_signatures: set = set() + # (task, canonical signature) -> [first evaluation runtime, evaluations avoided] + self._canonical_classes: Dict[Tuple, List] = {} self.min_modalities = max(1, min_modalities) requested_max = max_modalities or len(modalities) @@ -399,6 +470,7 @@ def __init__( self.optimization_results: Dict[str, List[FusionSearchResult]] = {} self.evaluation_errors: Dict[str, int] = {} + self.evaluation_error_samples: Dict[str, List[str]] = {} self.hall_of_fame_size = max(1, hall_of_fame_size) self.hall_of_fame: Dict[str, List[FusionSearchResult]] = {} @@ -659,11 +731,18 @@ def _start_parallel_runtime(self, task_name: str) -> None: shared_modality._data = wrapped shm_names.append(shm_name) shared_modalities.append(shared_modality) + # Workers are spawned with CUDA_VISIBLE_DEVICES="" unless they get + # a GPU lane, so without these a CUDA task model cannot run at all. worker_pool = PersistentWorkerPool( self.max_workers, _WORKER_DISPATCH, ctx=create_mp_context(), threads_per_worker=self.threads_per_worker, + gpu_devices=[info["index"] for info in gpu_memory_info()], + gpu_slots_per_device=int( + os.environ.get("SCURO_GPU_SLOTS_PER_DEVICE", "1") + ), + gpu_demand_fraction=1.0, ) except Exception: for shm_name in shm_names: @@ -684,6 +763,17 @@ def _shutdown_parallel_runtime(self) -> None: unlink_shm(shm_name) self._parallel_shm_names = [] + def _idle_lane(self, uses_gpu: bool): + pool = self._worker_pool + if uses_gpu: + for gpu_id in pool.gpu_worker_devices: + if pool.has_idle_worker_for(gpu_id): + return gpu_id + return None if pool.has_idle_cpu_worker else False + if pool.has_idle_worker_for(None, allow_gpu_worker_for_cpu=True): + return None + return False + def _evaluate_individuals_parallel( self, individuals: List[Any], task: Task ) -> None: @@ -711,20 +801,25 @@ def _evaluate_individuals_parallel( ind.fitness.values = self._reject_oversized(task_name, genome, reason) continue dag = self._genome_to_dag(genome) + if explicit_prep_on(self.algebra): + count_explicit_prep(self.algebra_stats, dag) pending_followers[sig] = [] pending_work.append((ind, dag, sig)) + uses_gpu = _task_uses_gpu(task) try: while pending_work or jobs: - while ( - pending_work - and self._worker_pool.has_idle_worker - and len(jobs) < self.batch_size - ): + while pending_work and len(jobs) < self.batch_size: + lane = self._idle_lane(uses_gpu) + if lane is False: + break ind, dag, sig = pending_work.pop(0) job_id = self._worker_pool.submit( "genome", - (dag, task, self._parallel_modalities, self.objective_specs), + (dag, task, self._parallel_modalities, self.objective_specs) + + self._share_args(), + gpu_id=lane, + allow_gpu_worker_for_cpu=not uses_gpu, ) jobs[job_id] = (ind, dag, sig) @@ -771,10 +866,22 @@ def _record_evaluation( self.evaluation_errors[task_name] = ( self.evaluation_errors.get(task_name, 0) + 1 ) - if self.debug and error is not None: + samples = self.evaluation_error_samples.setdefault(task_name, []) + if error is not None and len(samples) < _MAX_ERROR_SAMPLES: + spec = dag.to_spec() if dag is not None else None + samples.append(f"{spec}\n{error}") + if error is not None and (self.debug or len(samples) == 1): print(f"[GA] genome evaluation failed for task={task_name}:\n{error}") return + cache_stats = payload.get("subtree_cache") + if cache_stats: + for key in ("hits", "misses", "saved_seconds"): + self.algebra_stats[f"subtree_cache_{key}"] += cache_stats[key] + self.algebra_stats["subtree_cache_max_bytes"] = max( + self.algebra_stats["subtree_cache_max_bytes"], cache_stats["bytes"] + ) + result = FusionSearchResult( dag=dag, train_score=payload["train_score"], @@ -798,6 +905,13 @@ def _record_evaluation( ) self.optimization_results.setdefault(task_name, []).append(result) self._update_hall_of_fame(task_name, result) + if self.algebra.ga_canonical_dedup: + cls = self._canonical_classes.setdefault( + (task_name, self._canonical_signature(dag)), [None, 0] + ) + if cls[0] is None: + cls[0] = result.runtime + self.algebra_stats["ga_seconds_avoided"] += cls[0] * cls[1] self._eval_counter += 1 def _dominates(self, a: Tuple[float, ...], b: Tuple[float, ...]) -> bool: @@ -962,18 +1076,35 @@ def _genome_to_dag(self, genome: DagGenome) -> RepresentationDag: builder.create_leaf_node(mod_id, repr_idx) for mod_id, repr_idx in genome.leaves ] + leaf_stats = ( + { + leaf: self._leaf_stats(self._current_task_name, leaf) + for leaf in genome.leaves + } + if self is not None + and explicit_prep_on(self.algebra) + and self._current_task_name is not None + else None + ) def build(subtree: Tree, path: str) -> str: if isinstance(subtree, int): return leaf_ids[subtree] left, right = subtree - left_id = build(left, path + "L") - right_id = build(right, path + "R") + inputs = [build(left, path + "L"), build(right, path + "R")] op_cls = genome.fusion_ops[path] op = op_cls() - return builder.create_operation_node( - op.__class__, [left_id, right_id], op.get_current_parameters() - ) + params = op.get_current_parameters() + if leaf_stats is not None: + inputs = explicit_fusion_operands( + builder, + op.__class__, + params, + inputs, + leaf_stats, + self.algebra, + ) + return builder.create_operation_node(op.__class__, inputs, params) return builder.build(build(genome.tree, "")) @@ -981,11 +1112,41 @@ def _genome_signature(self, g: DagGenome) -> Tuple: def norm(t: Tree): return t if isinstance(t, int) else (norm(t[0]), norm(t[1])) - return ( + legacy = ( tuple(g.leaves), norm(g.tree), tuple(sorted((p, c.__name__) for p, c in g.fusion_ops.items())), ) + if not self.algebra.ga_canonical_dedup: + return legacy + canon = self._canonical_signatures.get(legacy) + if canon is None: + canon = self._canonical_signature(self._genome_to_dag(g)) + self._canonical_signatures[legacy] = canon + self._count_signature_request(legacy, canon) + return canon + + def _canonical_signature(self, dag: RepresentationDag) -> Tuple: + return ( + "canon", + canonical_id(dag, self.algebra.laws, self._order_invariant), + ) + + def _count_signature_request(self, legacy: Tuple, canon: Tuple) -> None: + task_name = self._current_task_name + self.algebra_stats["ga_signature_requests"] += 1 + if (task_name, legacy) in self._seen_legacy_signatures: + return + self._seen_legacy_signatures.add((task_name, legacy)) + self.algebra_stats["ga_unique_legacy"] += 1 + cls = self._canonical_classes.get((task_name, canon)) + if cls is None: + self._canonical_classes[(task_name, canon)] = [None, 0] + self.algebra_stats["ga_unique_canonical"] += 1 + return + cls[1] += 1 + self.algebra_stats["ga_evaluations_avoided"] += 1 + self.algebra_stats["ga_seconds_avoided"] += cls[0] or 0.0 def _leaf_stats(self, task_name: str, leaf: Tuple[Any, int]): """Shape of one candidate leaf representation, or None if unknown.""" @@ -996,18 +1157,7 @@ def _leaf_stats(self, task_name: str, leaf: Tuple[Any, int]): ] 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) + return sample_stats(modality) def _predicted_feature_dim( self, task_name: str, genome: DagGenome @@ -1071,7 +1221,24 @@ def walk(subtree: Tree, path: str): except Exception: return None + def _type_error(self, task_name: str, genome: DagGenome) -> Optional[str]: + reps = self.k_best_representations[task_name] + types = { + (mid, idx): getattr(reps[mid][idx], "modality_type", None) + for mid, idx in genome.leaves + } + stats = {leaf: self._leaf_stats(task_name, leaf) for leaf in genome.leaves} + violations = typecheck(self._genome_to_dag(genome), types, stats) + errors = count_violations(self.algebra_stats, violations) + if errors and self.algebra.typecheck == "strict": + return f"type_error: {errors}" + return None + def _oversized_reason(self, task_name: str, genome: DagGenome) -> Optional[str]: + if self.algebra.typecheck != "off" and task_name is not None: + reason = self._type_error(task_name, genome) + if reason is not None: + return reason reason = self._operator_input_violation(task_name, genome) if reason is not None: return f"{reason}; candidate rejected before execution" @@ -1112,13 +1279,15 @@ def _evaluate_genome(self, genome: DagGenome, task: Task) -> Tuple[float, ...]: return self._reject_oversized(task_name, genome, reason) dag = self._genome_to_dag(genome) + if explicit_prep_on(self.algebra): + count_explicit_prep(self.algebra_stats, dag) modalities = list( chain.from_iterable(self.k_best_representations[task_name].values()) ) try: fitness, payload = _evaluate_genome_body( - dag, task, modalities, self.objective_specs + dag, task, modalities, self.objective_specs, *self._share_args() ) error = None if fitness is None: @@ -1135,6 +1304,10 @@ def _evaluate_genome(self, genome: DagGenome, task: Task) -> Tuple[float, ...]: cache[sig] = fitness return fitness + def _share_args(self) -> Tuple: + # Legacy passes no extra argument, so 4-argument overrides keep working. + return (True,) if self.algebra.share_fused_subtrees else () + def _crossover_genomes( self, g1: DagGenome, g2: DagGenome ) -> Tuple[DagGenome, DagGenome]: diff --git a/src/main/python/systemds/scuro/drsearch/multimodal_optimizer.py b/src/main/python/systemds/scuro/drsearch/multimodal_optimizer.py index 54cddfd142b..22476bce60c 100644 --- a/src/main/python/systemds/scuro/drsearch/multimodal_optimizer.py +++ b/src/main/python/systemds/scuro/drsearch/multimodal_optimizer.py @@ -37,12 +37,23 @@ from systemds.scuro.representations.aggregated_representation import ( AggregatedRepresentation, ) +from systemds.scuro.drsearch.algebra import ( + AlgebraConfig, + canonical_id, + count_explicit_prep, + explicit_fusion_operands, + explicit_prep_on, + order_invariant, + sample_stats, +) from systemds.scuro.drsearch.operator_registry import Registry from systemds.scuro.utils.schema_helpers import get_shape from concurrent.futures import ProcessPoolExecutor, wait, FIRST_COMPLETED +import math import pickle import copy +from collections import Counter import time import traceback from itertools import chain @@ -128,6 +139,7 @@ def __init__( checkpoint_every: int = None, checkpoint_dir: Optional[str] = None, resume: bool = False, + algebra: Optional[AlgebraConfig] = None, ): self.modalities = modalities self.tasks = tasks @@ -135,6 +147,9 @@ def __init__( self.debug = debug self.min_modalities = max(2, min_modalities) self.max_modalities = max_modalities or len(modalities) + self.algebra = algebra or AlgebraConfig() + self.algebra_stats = Counter() + self._order_invariant = order_invariant(self.algebra, tasks) self.operator_registry = Registry() self.fusion_operators = self.operator_registry.get_fusion_operators() @@ -217,7 +232,7 @@ def optimize_parallel( break for dag in self._generate_fusion_dags( - modality_subset, repr_combo + modality_subset, repr_combo, task.model.name ): if max_combinations and evaluated_count >= max_combinations: stop_generation = True @@ -353,9 +368,20 @@ def _generate_representation_combinations( } def _generate_fusion_dags( - self, modality_subset: List[str], representation_combo: Dict[str, int] + self, + modality_subset: List[str], + representation_combo: Dict[str, int], + task_name: str = None, ) -> Generator[RepresentationDag, None, None]: leaf_infos = [(m, representation_combo[m]) for m in modality_subset] + leaf_stats = ( + { + (m, i): sample_stats(self.k_best_representations[task_name][m][i]) + for m, i in leaf_infos + } + if explicit_prep_on(self.algebra) and task_name is not None + else None + ) def gen_trees(indices: List[int]): if len(indices) == 1: @@ -393,18 +419,35 @@ def build_variants( for fusion_op_class in self.fusion_operators: new_builder = copy.deepcopy(right_builder) fusion_op = fusion_op_class() + params = fusion_op.get_current_parameters() + inputs = [left_root, right_root] + if leaf_stats is not None: + inputs = explicit_fusion_operands( + new_builder, + fusion_op.__class__, + params, + inputs, + leaf_stats, + self.algebra, + ) fusion_id = new_builder.create_operation_node( - fusion_op.__class__, - [left_root, right_root], - fusion_op.get_current_parameters(), + fusion_op.__class__, inputs, params ) variants.append((new_builder, fusion_id)) return variants n = len(leaf_infos) + dedup_permutations = self.algebra.exhaustive_dedup_permutations + # gen_trees already yields every leaf order, so permuting the leaves repeats each tree n! times. + leaf_orders = ( + [leaf_infos] + if dedup_permutations + else itertools.permutations(leaf_infos, n) + ) + seen_canonical = set() - for permuted_leaf_infos in itertools.permutations(leaf_infos, n): + for permuted_leaf_infos in leaf_orders: base_builder = RepresentationDAGBuilder() leaf_id_map = {} for idx, (modality_id, repr_idx) in enumerate(permuted_leaf_infos): @@ -416,12 +459,28 @@ def build_variants( for tree in gen_trees(indices): variants = build_variants(tree, base_builder, leaf_id_map) for builder_variant, root_id in variants: + self.algebra_stats["exhaustive_generated"] += ( + math.factorial(n) if dedup_permutations else 1 + ) + if dedup_permutations: + self.algebra_stats["exhaustive_after_permutation_dedup"] += 1 try: - yield builder_variant.build(root_id) + dag = builder_variant.build(root_id) except ValueError: if self.debug: print(f"[DEBUG] Skipping invalid DAG for root {root_id}") continue + if self.algebra.exhaustive_canonical: + cid = canonical_id( + dag, self.algebra.laws, self._order_invariant + ) + if cid in seen_canonical: + continue + seen_canonical.add(cid) + self.algebra_stats["exhaustive_after_canonical"] += 1 + if leaf_stats is not None: + count_explicit_prep(self.algebra_stats, dag) + yield dag def _evaluate_dag(self, dag: RepresentationDag, task: Task) -> "OptimizationResult": start_time = time.time() @@ -525,7 +584,9 @@ def optimize( modality_subset, task.model.name ): - for dag in self._generate_fusion_dags(modality_subset, repr_combo): + for dag in self._generate_fusion_dags( + modality_subset, repr_combo, task.model.name + ): if max_combinations and evaluated_count >= max_combinations: break if skip_remaining > 0: diff --git a/src/main/python/systemds/scuro/drsearch/node_executor.py b/src/main/python/systemds/scuro/drsearch/node_executor.py index a498c2752e3..4be1b669af0 100644 --- a/src/main/python/systemds/scuro/drsearch/node_executor.py +++ b/src/main/python/systemds/scuro/drsearch/node_executor.py @@ -163,6 +163,10 @@ def _offload_to_shared_memory(result: Any): return shm_name, resident_bytes, int(shm_bytes or 0), actual_stats +def _measure_memory() -> bool: + return DEBUG or os.environ.get("SCURO_MEASURE_MEMORY", "0") == "1" + + def _execute_node_worker(node, input_mods: List[Any], gpu_id: Optional[int]): start_time = time.perf_counter() if gpu_id is not None: @@ -205,7 +209,7 @@ def _run_node_op(): gpu_peak_bytes = -1 measurement = None - if DEBUG: + if _measure_memory(): 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, cleanup_after=False), @@ -271,7 +275,7 @@ def _run_task(): gpu_peak_bytes = -1 measurement = None - if DEBUG: + if _measure_memory(): result, measurement = measure_memory_during( lambda: _run_gpu_op(_run_task, gpu_id), input_resident_bytes=estimate_modality_bytes(modality), @@ -480,6 +484,7 @@ def __init__( self._cache_misses = 0 self._nodes_executed = 0 self._peak_cache_bytes = 0 + self._leaf_cached_bytes = self.scheduler.memory_stats["cpu_cached"] self._eval_counter = 0 self._nodes_executed = 0 self._search_start = ( @@ -741,6 +746,12 @@ def _process_result(self, jr) -> None: finally: if retained: self._result_cache.release_shm_names(retained) + self._sync_cache_accounting() + + def _sync_cache_accounting(self) -> None: + cache_bytes = self._result_cache.get_memory_total_memory_usage() + self._peak_cache_bytes = max(self._peak_cache_bytes, cache_bytes) + self.scheduler.set_cached_cpu_bytes(self._leaf_cached_bytes + cache_bytes) def _handle_job_failure(self, unit: Union[_NodeUnit, _BatchUnit], jr): reason = jr.error or "unknown worker failure" @@ -860,8 +871,8 @@ def _handle_modality_result( operation_name, actual_stats, measurement=measurement, + estimated_stats=estimated_stats, ) - before_bytes = self._result_cache.get_memory_total_memory_usage() self._manage_result_cache( node_id, transformed_modality, @@ -869,9 +880,7 @@ def _handle_modality_result( resident_bytes=resident_bytes, 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._sync_cache_accounting() self.scheduler.complete_node(node_id) def _manage_result_cache( @@ -906,9 +915,12 @@ def _checkpoint_memory_usage( operation_name: str, actual_stats: Optional[RepresentationStats], measurement: Optional[MemoryMeasurement] = None, + estimated_stats: Optional[RepresentationStats] = None, ): if self.enable_checkpointing: self._memory_usage_checkpoint.increment(node_id) + if estimated_stats is None: + estimated_stats = self.scheduler.node_stats.get(node_id) if measurement is not None: peak_bytes = measurement.footprint_bytes @@ -962,6 +974,12 @@ def _checkpoint_memory_usage( "rss_delta_bytes": measurement.rss_delta_bytes if measurement else -1, "num_instances": getattr(actual_stats, "num_instances", None), "output_shape": getattr(actual_stats, "output_shape", None), + "estimated_output_shape": getattr( + estimated_stats, "output_shape", None + ), + "estimated_num_instances": getattr( + estimated_stats, "num_instances", None + ), "dtype": str(getattr(actual_stats, "dtype", None)), "container": str(getattr(actual_stats, "container", None)), } @@ -1100,6 +1118,13 @@ def _reuse_report(self) -> Dict[str, Any]: "cpu_fallback_nodes": len( getattr(self.scheduler, "cpu_fallback_nodes", []) or [] ), + "planned_task_nodes": len(getattr(self.scheduler, "roots", ()) or ()), + "completed_task_nodes": self._eval_counter, + "exclusive_memory_nodes": len( + getattr(self.scheduler, "exclusive_memory_nodes", {}) or {} + ), + "deadlock": bool(getattr(self.scheduler, "deadlock", False)), + "deadlock_reason": getattr(self.scheduler, "deadlock_reason", None), } ) 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 78ef5b953fa..b69c50f43cc 100644 --- a/src/main/python/systemds/scuro/drsearch/node_scheduler.py +++ b/src/main/python/systemds/scuro/drsearch/node_scheduler.py @@ -44,6 +44,10 @@ def _exclusive_over_budget_enabled() -> bool: return os.environ.get("SCURO_EXCLUSIVE_OVER_BUDGET", "0") == "1" +def _cache_drain_fraction() -> float: + return float(os.environ.get("SCURO_CACHE_DRAIN_FRACTION", "0.5")) + + class MemoryAwareNodeScheduler: def __init__( @@ -141,16 +145,45 @@ def can_start_now(self, node_id: str) -> bool: def update_cpu_memory_in_use(self, delta_bytes: int): self.memory_stats["cpu_cached"] += delta_bytes + def set_cached_cpu_bytes(self, cached_bytes: int) -> None: + self.memory_stats["cpu_cached"] = cached_bytes + + def _cache_draining(self) -> bool: + budget = self.memory_budget["cpu"] + if _exclusive_over_budget_enabled() and self._total_cpu_in_use() >= budget: + return True + return self.memory_stats["cpu_cached"] >= _cache_drain_fraction() * budget + + def _release_bytes(self, node_id: str) -> int: + release_bytes = 0 + for parent_id in self.parents.get(node_id, set()): + if ( + parent_id not in self.leaves + and self.remaining_children.get(parent_id, 0) == 1 + ): + release_bytes += self.node_resources[parent_id][0] + return release_bytes + + def _grows_cache(self, node_id: str) -> bool: + adds_bytes = ( + self.node_resources[node_id][0] if self.children.get(node_id) else 0 + ) + return adds_bytes > self._release_bytes(node_id) + def get_runnable(self) -> List[RepresentationNode]: runnable_nodes = self._get_runnable_nodes() admitted_bytes = self._pending_admitted_cpu_bytes() pending_gpu = {gpu_id: 0 for gpu_id in self.gpu_devices} pending_slots = {gpu_id: 0 for gpu_id in self.gpu_devices} + draining = self._cache_draining() + hold_growing_nodes = False for node in runnable_nodes: if node in self._ready_set: continue + if hold_growing_nodes and self._grows_cache(node): + continue ok, gpu_id = self._check_memory_constraints( node, admitted_bytes, pending_gpu, pending_slots ) @@ -163,6 +196,13 @@ def get_runnable(self) -> List[RepresentationNode]: self._candidates.discard(node) self.ready_nodes.append(node) self._ready_set.add(node) + elif ( + draining + and not self._grows_cache(node) + and self.node_resources[node][0] + > self.memory_budget["cpu"] - self._total_cpu_in_use() - admitted_bytes + ): + hold_growing_nodes = True chunked_leaf_by_device = defaultdict(list) for node in self.ready_nodes: if isinstance(node, list): @@ -189,28 +229,17 @@ 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"] - ) + # Upward rank always puts a producer ahead of its consumers, so without + # a counterweight every large intermediate is materialized before the + # tasks that would release it. Once the cache holds a large share of + # the budget (or, with exclusive admission, the budget is exceeded), + # nodes whose completion does not grow the cache (task roots, or last + # children releasing at least their own output) run first. + draining_first = self._cache_draining() def _score(node_id: str): - release_bytes = 0 - for parent_id in self.parents.get(node_id, set()): - if ( - parent_id not in self.leaves - 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 + release_bytes = self._release_bytes(node_id) + grows_cache = draining_first and self._grows_cache(node_id) return ( grows_cache, -self.node_priorities.get(node_id, 0.0), diff --git a/src/main/python/systemds/scuro/drsearch/representation_dag.py b/src/main/python/systemds/scuro/drsearch/representation_dag.py index d84d092bbaf..efc4692dcdf 100644 --- a/src/main/python/systemds/scuro/drsearch/representation_dag.py +++ b/src/main/python/systemds/scuro/drsearch/representation_dag.py @@ -36,8 +36,20 @@ DimensionalityReduction, ) from systemds.scuro.utils.identifier import get_op_id, get_node_id +from systemds.scuro.drsearch.algebra import ( + ALL_LAWS, + AlgebraConfig, + canonical_id, + count_explicit_prep, + count_violations, + explicit_fusion_operands, + explicit_prep_on, + order_invariant, + typecheck, +) +from systemds.scuro.representations.fusion import Fusion import time -from collections import OrderedDict, defaultdict, deque +from collections import Counter, OrderedDict, defaultdict, deque def pushdown_aggregation_for_node( @@ -156,6 +168,9 @@ def pipeline_id(self) -> str: blob = json.dumps(self.to_spec(), sort_keys=True, separators=(",", ":")) return hashlib.sha1(blob.encode()).hexdigest()[:16] + def canonical_id(self, laws=ALL_LAWS, invariant=False) -> str: + return canonical_id(self, laws, invariant) + def _topological_order(self) -> List[str]: node_map = {node.node_id: node for node in self.nodes} visited = [] @@ -670,7 +685,13 @@ def pushdown_aggregation(dag_group: List[RepresentationDag]) -> List[Representat class CSEAwareDAGBuilder: - def __init__(self): + def __init__(self, algebra: Optional[AlgebraConfig] = None, tasks=None, leaves=()): + self.algebra = algebra or AlgebraConfig() + self.leaves = leaves + self.algebra_stats = Counter() + self._order_invariant = order_invariant(self.algebra, tasks) + self._seen_ordered = set() + self._seen_legacy = set() self.global_nodes: List[RepresentationNode] = [] self.signature_to_node: Dict[Hashable, str] = {} self.node_to_signature: Dict[str, Hashable] = {} @@ -685,8 +706,28 @@ def _compute_node_signature( input_sigs = tuple(sorted(ip, key=repr)) if inputs else () op_cls = operation().name params_items = tuple(sorted((parameters or {}).items())) + if self.algebra.cse_law_aware and not self._may_reorder(operation, parameters): + input_sigs = tuple(ip) return ("op", op_cls, params_items, input_sigs) + def _may_reorder(self, operation: Any, parameters: Dict[str, Any]) -> bool: + if not issubclass(operation, Fusion): + return False + laws = operation.algebraic_laws(parameters or {}) + return "comm" in laws or ("concat_perm" in laws and self._order_invariant) + + def _count_law_merges(self, signature: Hashable, inputs: List[str]): + ip = tuple(self.node_to_signature[inp] for inp in inputs) + ordered = signature[:3] + (ip,) + legacy = signature[:3] + (tuple(sorted(ip, key=repr)),) + if signature in self.signature_to_node: + if ordered not in self._seen_ordered: + self.algebra_stats["cse_merges_by_law"] += 1 + elif legacy in self._seen_legacy: + self.algebra_stats["cse_unsafe_merges_prevented"] += 1 + self._seen_ordered.add(ordered) + self._seen_legacy.add(legacy) + def _compute_leaf_signature( self, modality_id: str, representation_index: int = -1 ) -> Hashable: @@ -705,6 +746,8 @@ def _get_or_create_node( signature = self._compute_leaf_signature(modality_id, representation_index) else: signature = self._compute_node_signature(operation, inputs, parameters) + if self.algebra.cse_law_aware: + self._count_law_merges(signature, inputs) try: if signature in self.signature_to_node: @@ -751,6 +794,15 @@ def create_leaf_node( def create_operation_node( self, operation: Any, inputs: List[str], parameters: Dict[str, Any] = None ): + if explicit_prep_on(self.algebra): + inputs = explicit_fusion_operands( + self, + operation, + parameters, + inputs, + self._leaf_stats(), + self.algebra, + ) return self._get_or_create_node( operation=operation, inputs=inputs, parameters=parameters, is_leaf=False ) @@ -769,11 +821,31 @@ def build(self, root_node_id: str, dag_id: int = None) -> RepresentationDag: self.dag_counter += 1 if not dag.validate(): raise ValueError("Invalid DAG construction") + if self.algebra.typecheck != "off": + self._typecheck(dag) + if explicit_prep_on(self.algebra): + count_explicit_prep(self.algebra_stats, dag) if dag_id is None: self._dag_by_root[root_node_id] = dag return dag + def _leaf_stats(self) -> Dict[Hashable, Any]: + return {(m.modality_id, -1): getattr(m, "stats", None) for m in self.leaves} + + def _typecheck(self, dag: RepresentationDag) -> None: + violations = typecheck( + dag, + { + (m.modality_id, -1): getattr(m, "modality_type", None) + for m in self.leaves + }, + self._leaf_stats(), + ) + errors = count_violations(self.algebra_stats, violations) + if errors and self.algebra.typecheck == "strict": + raise ValueError(f"Ill-typed plan: {errors}") + def get_node(self, node_id: str) -> Optional[RepresentationNode]: for node in self.global_nodes: if node.node_id == node_id: 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 97de3da3a58..05b70d03af6 100644 --- a/src/main/python/systemds/scuro/drsearch/split_aware_evaluator.py +++ b/src/main/python/systemds/scuro/drsearch/split_aware_evaluator.py @@ -63,9 +63,17 @@ def dag_requires_fitting(dag: RepresentationDag) -> bool: class SplitAwareDAGEvaluator: - def __init__(self, dag: RepresentationDag, modalities, task, node_cache=None): + def __init__( + self, + dag: RepresentationDag, + modalities, + task, + node_cache=None, + share_fitted=False, + ): self.dag = dag self.node_cache = node_cache + self.share_fitted = share_fitted self.modalities = ( list(modalities) if isinstance(modalities, (list, tuple)) else [modalities] ) @@ -231,10 +239,12 @@ 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) + reusable = self.node_cache is not None and ( + self.share_fitted + or ( + 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) @@ -254,7 +264,7 @@ def execute_node(node_id): if requirement is not FitRequirement.NONE: fitted.add(node_id) - reusable = False + reusable = reusable and self.share_fitted labels = None fit_input = input_partitions[0] fit_ids = ( diff --git a/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py b/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py index ced501ba095..8b717e5ced3 100644 --- a/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py +++ b/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py @@ -19,6 +19,7 @@ # # ------------------------------------------------------------- import copy +from collections import Counter import math import pickle import time @@ -44,6 +45,7 @@ AggregatedRepresentation, ) from systemds.scuro.modality.modality import Modality +from systemds.scuro.drsearch.algebra import AlgebraConfig from systemds.scuro.drsearch.operator_registry import Registry from systemds.scuro.representations.window_aggregation import WindowAggregation from systemds.scuro.utils.checkpointing import CheckpointManager @@ -71,10 +73,13 @@ def __init__( max_num_workers: int = -1, enable_checkpointing: bool = False, window_combination_chains: int = 1, + algebra: Optional[AlgebraConfig] = None, ): self._node_stats: Dict[str, Any] = {} self.pruned = [] self.window_combination_chains = window_combination_chains + self.algebra = algebra or AlgebraConfig() + self.algebra_stats = Counter() self.enable_checkpointing = enable_checkpointing self.modalities = modalities self.tasks = tasks @@ -96,7 +101,8 @@ def __init__( ) # TODO: check if we should make this a local variable (might keep unnecessary memory usage) self.builders = { - modality.modality_id: CSEAwareDAGBuilder() for modality in modalities + modality.modality_id: CSEAwareDAGBuilder(self.algebra, tasks, [modality]) + for modality in modalities } self.debug = debug @@ -392,6 +398,17 @@ def _process_modality(self, modality, skip_remaining: int = 0, scheduler=None): exec_out = node_executor.run() task_results.extend(exec_out["task_results"]) statistics = exec_out["statistics"] + reuse = statistics.get("reuse", {}) + planned = reuse.get("planned_task_nodes", 0) + completed = reuse.get("completed_task_nodes", planned) + if reuse.get("deadlock") or completed < planned: + print( + f"[unimodal] WARNING: modality {modality.modality_id} " + f"evaluated {completed} of {planned} planned task nodes; " + f"deadlock={reuse.get('deadlock')} " + f"reason={reuse.get('deadlock_reason')}", + flush=True, + ) if fitted_dags: fitted_task_dags = self._expand_dags_with_task_roots(fitted_dags) diff --git a/src/main/python/systemds/scuro/modality/transformed.py b/src/main/python/systemds/scuro/modality/transformed.py index 62cc03eec04..87cc3e91434 100644 --- a/src/main/python/systemds/scuro/modality/transformed.py +++ b/src/main/python/systemds/scuro/modality/transformed.py @@ -185,10 +185,13 @@ def combine(self, other: Union[Modality, List[Modality]], fusion_method): fused_modality = TransformedModality( self, fusion_method, ModalityType.EMBEDDING ) + modalities = self.create_modality_list(other) start_time = time.time() - fused_modality.data = fusion_method.transform(self.create_modality_list(other)) + fused_modality.data = fusion_method.transform(modalities) end_time = time.time() - fused_modality.transform_time = end_time - start_time + fused_modality.transform_time += (end_time - start_time) + sum( + getattr(m, "transform_time", 0) or 0 for m in modalities[1:] + ) return fused_modality def create_modality_list(self, other: Union[Modality, List[Modality]]): diff --git a/src/main/python/systemds/scuro/models/model.py b/src/main/python/systemds/scuro/models/model.py index fcd4150cab2..965058eb62c 100644 --- a/src/main/python/systemds/scuro/models/model.py +++ b/src/main/python/systemds/scuro/models/model.py @@ -22,6 +22,7 @@ class Model: uses_gpu = None + feature_order_invariant = None def __init__(self, name: str): """ diff --git a/src/main/python/systemds/scuro/representations/aggregate.py b/src/main/python/systemds/scuro/representations/aggregate.py index 366bb9f3ce9..2c54913445c 100644 --- a/src/main/python/systemds/scuro/representations/aggregate.py +++ b/src/main/python/systemds/scuro/representations/aggregate.py @@ -51,18 +51,22 @@ def _mode_agg(data, aggregate_dim=0): if arr.ndim == 1: values, counts = np.unique(arr, return_counts=True) return values[np.argmax(counts)] - moved = np.moveaxis(arr, axis, -1) - flat = moved.reshape(-1, moved.shape[-1]) + axes = tuple( + a % arr.ndim for a in (axis if isinstance(axis, tuple) else (axis,)) + ) + moved = np.moveaxis(arr, axes, tuple(range(arr.ndim - len(axes), arr.ndim))) + kept = moved.shape[: arr.ndim - len(axes)] + flat = moved.reshape(int(np.prod(kept)), -1) modes = np.empty(flat.shape[0], dtype=arr.dtype) for i in range(flat.shape[0]): values, counts = np.unique(flat[i], return_counts=True) modes[i] = values[np.argmax(counts)] - return modes.reshape(moved.shape[:-1]) + return modes.reshape(kept) @staticmethod def _normalize_axis(aggregate_dim): if isinstance(aggregate_dim, tuple): - return aggregate_dim[0] if aggregate_dim else 0 + return aggregate_dim if aggregate_dim else 0 return aggregate_dim _aggregation_function = { diff --git a/src/main/python/systemds/scuro/representations/average.py b/src/main/python/systemds/scuro/representations/average.py index 06c435d6580..ade8b03eae0 100644 --- a/src/main/python/systemds/scuro/representations/average.py +++ b/src/main/python/systemds/scuro/representations/average.py @@ -32,14 +32,14 @@ @register_fusion_operator() class Average(Fusion): + _laws = frozenset({"comm", "avg_flatten"}) + def __init__(self, params=None): """ Combines modalities using averaging """ super().__init__("Average") self.needs_alignment = True - self.associative = True - self.commutative = True def execute(self, modalities: List[Modality]): data = np.array(modalities[0].data, dtype=np.float64) diff --git a/src/main/python/systemds/scuro/representations/clip.py b/src/main/python/systemds/scuro/representations/clip.py index a64354b1ace..55f841e42e8 100644 --- a/src/main/python/systemds/scuro/representations/clip.py +++ b/src/main/python/systemds/scuro/representations/clip.py @@ -63,12 +63,13 @@ @register_representation([ModalityType.VIDEO, ModalityType.IMAGE]) @register_expensive_representation([ModalityType.VIDEO, ModalityType.IMAGE]) class CLIPVisual(UnimodalRepresentation): + OUTPUT_MODALITY_TYPE = ModalityType.EMBEDDING supports_aggregation_pushdown = True cache_in_worker = True def __init__(self, output_file=None, batch_size=32, layer_name="", params=None): parameters = self._get_parameters() - super().__init__("CLIPVisual", ModalityType.EMBEDDING, parameters) + super().__init__("CLIPVisual", self.OUTPUT_MODALITY_TYPE, parameters) self.params = params self._activation_hook = None self.model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32") diff --git a/src/main/python/systemds/scuro/representations/concatenation.py b/src/main/python/systemds/scuro/representations/concatenation.py index cee30382963..f599aba0142 100644 --- a/src/main/python/systemds/scuro/representations/concatenation.py +++ b/src/main/python/systemds/scuro/representations/concatenation.py @@ -34,6 +34,14 @@ @register_fusion_operator() class Concatenation(Fusion): + _laws = frozenset({"assoc", "concat_perm"}) + + @classmethod + def algebraic_laws(cls, parameters: dict) -> frozenset: + if parameters.get("preserve_leading_axis"): + return frozenset({"assoc"}) + return cls._laws + def __init__(self, params=None, preserve_leading_axis=False): super().__init__("Concatenation") if params is not None: diff --git a/src/main/python/systemds/scuro/representations/fusion.py b/src/main/python/systemds/scuro/representations/fusion.py index 5cbe128cae0..95bdec2c1a9 100644 --- a/src/main/python/systemds/scuro/representations/fusion.py +++ b/src/main/python/systemds/scuro/representations/fusion.py @@ -55,22 +55,28 @@ class FusionInputTooLargeError(ValueError): class Fusion(Representation): + _laws = frozenset() + preserves_leading_axis = False + def __init__(self, name, parameters=None): """ Parent class for different multimodal fusion types :param name: Name of the fusion type """ super().__init__(name, parameters) - self.associative = False - self.commutative = False + self.associative = "assoc" in self._laws + self.commutative = "comm" in self._laws self.needs_alignment = False self.needs_training = False self.needs_instance_alignment = False - 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 + @classmethod + def algebraic_laws(cls, parameters: dict) -> frozenset: + return cls._laws + @staticmethod def _as_stats_list(input_stats) -> List[RepresentationStats]: if isinstance(input_stats, RepresentationStats): diff --git a/src/main/python/systemds/scuro/representations/glove.py b/src/main/python/systemds/scuro/representations/glove.py index 9513760fb9f..8e9cf0f8a72 100644 --- a/src/main/python/systemds/scuro/representations/glove.py +++ b/src/main/python/systemds/scuro/representations/glove.py @@ -49,8 +49,10 @@ def load_glove_embeddings(file_path): @register_representation(ModalityType.TEXT) @register_expensive_representation(ModalityType.TEXT) class GloVe(UnimodalRepresentation): + OUTPUT_MODALITY_TYPE = ModalityType.EMBEDDING + def __init__(self, output_file=None, params=None): - super().__init__("GloVe", ModalityType.TEXT) + super().__init__("GloVe", self.OUTPUT_MODALITY_TYPE) file_path = hf_hub_download( repo_id="stanfordnlp/glove", filename="glove.6B.zip" ) diff --git a/src/main/python/systemds/scuro/representations/hadamard.py b/src/main/python/systemds/scuro/representations/hadamard.py index 0697a9f67ea..a88d7d68634 100644 --- a/src/main/python/systemds/scuro/representations/hadamard.py +++ b/src/main/python/systemds/scuro/representations/hadamard.py @@ -33,14 +33,14 @@ @register_fusion_operator() class Hadamard(Fusion): + _laws = frozenset({"comm", "assoc"}) + def __init__(self, params=None): """ Combines modalities using elementwise multiply (Hadamard product) """ super().__init__("Hadamard") self.needs_alignment = True # zero padding falsifies the result - self.commutative = True - self.associative = True def execute(self, modalities: List[Modality], train_indices=None): fused_data = np.prod([m.data for m in modalities], axis=0) diff --git a/src/main/python/systemds/scuro/representations/lstm.py b/src/main/python/systemds/scuro/representations/lstm.py index 97b318e8bad..46e0842c0c1 100644 --- a/src/main/python/systemds/scuro/representations/lstm.py +++ b/src/main/python/systemds/scuro/representations/lstm.py @@ -42,6 +42,7 @@ @register_fusion_operator() class LSTM(Fusion): fit_requirement = FitRequirement.SUPERVISED + preserves_leading_axis = True def __init__( self, @@ -89,7 +90,6 @@ def __init__( self.device = get_device() self.needs_training = True self.needs_alignment = False - self.preserves_leading_axis = True self.model = None self.input_dim = None self.num_classes = None diff --git a/src/main/python/systemds/scuro/representations/max.py b/src/main/python/systemds/scuro/representations/max.py index 2e0df4c464e..f221a0c882c 100644 --- a/src/main/python/systemds/scuro/representations/max.py +++ b/src/main/python/systemds/scuro/representations/max.py @@ -31,6 +31,8 @@ @register_fusion_operator() class RowMax(Fusion): + _laws = frozenset({"comm", "assoc"}) + def __init__(self, params=None): """ Combines modalities by computing the outer product of a modality combination and @@ -38,8 +40,6 @@ def __init__(self, params=None): """ super().__init__("RowMax") self.needs_alignment = True - self.associative = True - self.commutative = True def execute( self, @@ -64,3 +64,14 @@ def get_output_stats(self, input_stats_list) -> RepresentationStats: 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) + + def estimate_peak_memory_bytes(self, input_stats) -> dict: + stats_list = self._as_stats_list(input_stats) + input_bytes = sum(self._stats_bytes(s) for s in stats_list) + output_bytes = self._stats_bytes(self.get_output_stats(input_stats)) + + raw_bytes = self._raw_input_bytes(input_stats) + cpu_peak = ( + int((raw_bytes + 2 * input_bytes + output_bytes) * 1.15) + 8 * 1024 * 1024 + ) + return {"cpu_peak_bytes": cpu_peak, "gpu_peak_bytes": 0} diff --git a/src/main/python/systemds/scuro/representations/openface.py b/src/main/python/systemds/scuro/representations/openface.py index 60599e806c3..b2d82cc6e91 100644 --- a/src/main/python/systemds/scuro/representations/openface.py +++ b/src/main/python/systemds/scuro/representations/openface.py @@ -37,6 +37,7 @@ from systemds.scuro.representations.representation import ( CONTAINER_LIST, RepresentationStats, + stats_itemsize, ) from systemds.scuro.representations.unimodal import UnimodalRepresentation from systemds.scuro.representations.utils import get_sequence_lengths @@ -170,6 +171,9 @@ class OpenFace(UnimodalRepresentation): FEATURE_DIM = len(FEATURE_COLUMNS) FRAME_BATCH_SIZE = 32 + MAX_SIDE = 1024 + SIZE_WINDOW_CHUNKS = 8 + DETECTOR_MEAN_BGR = (104, 117, 123) FACE_MODEL_MEMORY_BYTES = 8 * 1024 * 1024 MULTITASK_MODEL_MEMORY_BYTES = 128 * 1024 * 1024 @@ -177,6 +181,7 @@ class OpenFace(UnimodalRepresentation): CPU_RUNTIME_OVERHEAD_BYTES = 128 * 1024 * 1024 GPU_RUNTIME_OVERHEAD_BYTES = 64 * 1024 * 1024 LANDMARK_RUNTIME_BYTES = 2400 * 1024 * 1024 + DETECTOR_BYTES_PER_PIXEL = 144 LANDMARK_BYTES_PER_FRAME = 52 * 1024 * 1024 MULTITASK_BYTES_PER_FRAME = 26 * 1024 * 1024 @@ -186,6 +191,7 @@ def __init__( confidence_threshold=0.02, nms_threshold=0.4, vis_threshold=0.5, + max_side=MAX_SIDE, params=None, ): if params is not None: @@ -195,6 +201,7 @@ def __init__( ) nms_threshold = params.get("nms_threshold", nms_threshold) vis_threshold = params.get("vis_threshold", vis_threshold) + max_side = params.get("max_side", max_side) parameters = { "feature_set": list(self.FEATURE_SETS), @@ -207,6 +214,7 @@ def __init__( self.confidence_threshold = float(confidence_threshold) self.nms_threshold = float(nms_threshold) self.vis_threshold = float(vis_threshold) + self.max_side = int(max_side) self.params = params self.data_type = np.float32 self._gpu_id = None @@ -300,9 +308,28 @@ def estimate_peak_memory_bytes(self, input_stats) -> dict: max_height = int(getattr(input_stats, "max_height", 224)) max_width = int(getattr(input_stats, "max_width", 224)) max_channels = int(getattr(input_stats, "max_channels", 3)) - frame_bytes = ( - max_height * max_width * max_channels * np.dtype(np.float32).itemsize + float_size = np.dtype(np.float32).itemsize + + if isinstance(input_stats, VideoStats): + video_bytes = ( + int(input_stats.max_length) + * max_height + * max_width + * max_channels + * stats_itemsize(input_stats) + ) + sample_bytes = 2 * video_bytes + 2 * max_height * max_width * 3 * float_size + else: + sample_bytes = 2 * max_height * max_width * max_channels + + canvas_pixels = min(max_height, self.max_side) * min(max_width, self.max_side) + window_bytes = ( + self.FRAME_BATCH_SIZE * self.SIZE_WINDOW_CHUNKS * canvas_pixels * 3 ) + chunk_bytes = window_bytes + ( + self.FRAME_BATCH_SIZE * canvas_pixels * 3 * 2 * float_size + ) + 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 @@ -315,13 +342,15 @@ def estimate_peak_memory_bytes(self, input_stats) -> dict: return { "cpu_peak_bytes": int( output_bytes - + frame_bytes + + sample_bytes + + chunk_bytes + model_bytes + self.CPU_RUNTIME_OVERHEAD_BYTES ), "gpu_peak_bytes": int( model_bytes - + self.FRAME_BATCH_SIZE * (frame_bytes + batch_bytes) + + self.FRAME_BATCH_SIZE + * (canvas_pixels * self.DETECTOR_BYTES_PER_PIXEL + batch_bytes) + runtime_bytes ), } @@ -366,15 +395,54 @@ def transform(self, modality, aggregation=None): return transformed_modality def _extract_image_features(self, images): - return list(self._extract_batch_features([self._to_bgr(i) for i in images])) + return list(self._extract_streamed_features(images)) def _extract_video_features(self, frames, owner_id): if len(frames) == 0: raise ValueError(f"Video instance {owner_id} contains no frames") - return self._extract_batch_features([self._to_bgr(frame) for frame in frames]) + return self._extract_streamed_features(frames) def _extract_features(self, image): - return self._extract_batch_features([self._to_bgr(image)])[0] + return self._extract_batch_features([self._prepare(image)])[0] + + def _extract_streamed_features(self, samples): + features = np.zeros((len(samples), self.feature_dim), dtype=np.float32) + window = self.FRAME_BATCH_SIZE * self.SIZE_WINDOW_CHUNKS + for window_start in range(0, len(samples), window): + indices = range(window_start, min(window_start + window, len(samples))) + prepared = {index: self._prepare(samples[index]) for index in indices} + order = sorted(indices, key=lambda index: prepared[index].shape[:2]) + for start, stop in self._chunks(len(order)): + chunk = order[start:stop] + features[chunk] = self._extract_batch_features( + [prepared[index] for index in chunk] + ) + return features + + def _prepare(self, image): + image = self._to_bgr(image) + height, width = image.shape[:2] + longer = max(height, width) + if longer <= self.max_side: + return image + scale = self.max_side / longer + size = (max(1, round(width * scale)), max(1, round(height * scale))) + return cv2.resize(image, size, interpolation=cv2.INTER_AREA) + + def _pad_to(self, image, height, width): + pad_bottom = height - image.shape[0] + pad_right = width - image.shape[1] + if not pad_bottom and not pad_right: + return image + return cv2.copyMakeBorder( + image, + 0, + pad_bottom, + 0, + pad_right, + cv2.BORDER_CONSTANT, + value=self.DETECTOR_MEAN_BGR, + ) def _extract_batch_features(self, images_bgr): self._load_models() @@ -460,53 +528,52 @@ def _detect_faces(self, 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] = ( + for start, stop in self._chunks(len(images_bgr)): + chunk = range(start, stop) + height = max(images_bgr[i].shape[0] for i in chunk) + width = max(images_bgr[i].shape[1] for i in chunk) + if (height, width) not in priors_by_shape: + priors_by_shape[(height, width)] = ( PriorBox(detector.cfg, image_size=(height, width)) .forward() .to(detector.device) .data ) - priors = priors_by_shape[shape] + priors = priors_by_shape[(height, width)] box_scale = torch.Tensor([width, height, width, height]).to(detector.device) landmark_scale = torch.Tensor([width, height] * 5).to(detector.device) - 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] + batch = torch.cat( + [ + detector.preprocess_image( + self._pad_to(images_bgr[i], height, width) + )[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 ) - 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 + 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): diff --git a/src/main/python/systemds/scuro/representations/sum.py b/src/main/python/systemds/scuro/representations/sum.py index 58cac2e7210..0b9ce2e399e 100644 --- a/src/main/python/systemds/scuro/representations/sum.py +++ b/src/main/python/systemds/scuro/representations/sum.py @@ -33,6 +33,8 @@ @register_fusion_operator() class Sum(Fusion): + _laws = frozenset({"comm", "assoc"}) + def __init__(self, params=None): """ Combines modalities using colum-wise sum @@ -41,16 +43,13 @@ def __init__(self, params=None): self.needs_alignment = True def execute(self, modalities: List[Modality]): - data = np.array( - modalities[0].data, - dtype=modalities[0].metadata[0]["data_layout"]["type"], - ) - - for m in range(1, len(modalities)): - data += np.asarray( - modalities[m].data, - dtype=modalities[m].metadata[0]["data_layout"]["type"], - ) + arrays = [ + np.asarray(m.data, dtype=m.metadata[0]["data_layout"]["type"]) + for m in modalities + ] + data = np.array(arrays[0], dtype=np.result_type(*arrays)) + for array in arrays[1:]: + data += array return data def get_output_stats(self, input_stats_list) -> RepresentationStats: 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 5bf4a6d0add..2172e6af4ab 100644 --- a/src/main/python/systemds/scuro/representations/swin_video_transformer.py +++ b/src/main/python/systemds/scuro/representations/swin_video_transformer.py @@ -55,6 +55,7 @@ @register_representation([ModalityType.VIDEO]) @register_expensive_representation([ModalityType.VIDEO]) class SwinVideoTransformer(FrameBudgetedVideoModel, UnimodalRepresentation): + OUTPUT_MODALITY_TYPE = ModalityType.EMBEDDING _EMBED_DIM = 768 cache_in_worker = True _MEMORY_PROFILE = VideoMemoryProfile( @@ -80,7 +81,7 @@ def __init__( # "batch_size": [1, 2, 4, 8, 16, 32], } self.data_type = torch.float32 - super().__init__("SwinVideoTransformer", ModalityType.EMBEDDING, parameters) + super().__init__("SwinVideoTransformer", self.OUTPUT_MODALITY_TYPE, parameters) if params is not None: layer_name = params.get("layer_name", layer_name) batch_size = int(params.get("batch_size", batch_size)) diff --git a/src/main/python/systemds/scuro/representations/vgg.py b/src/main/python/systemds/scuro/representations/vgg.py index 6f82937952a..214b56a59b0 100644 --- a/src/main/python/systemds/scuro/representations/vgg.py +++ b/src/main/python/systemds/scuro/representations/vgg.py @@ -56,6 +56,7 @@ def forward(self, input_: torch.Tensor) -> torch.Tensor: @register_representation([ModalityType.IMAGE, ModalityType.VIDEO]) @register_expensive_representation([ModalityType.IMAGE, ModalityType.VIDEO]) class VGG19(UnimodalRepresentation): + OUTPUT_MODALITY_TYPE = ModalityType.EMBEDDING supports_aggregation_pushdown = True cache_in_worker = True @@ -71,7 +72,7 @@ def __init__( self.model = models.vgg19(weights=models.VGG19_Weights.DEFAULT) self.model = self.model.to(self.device) parameters = self._get_parameters() - super().__init__("VGG19", ModalityType.EMBEDDING, parameters) + super().__init__("VGG19", self.OUTPUT_MODALITY_TYPE, parameters) self.params = params if params is not None: batch_size = int((params or {}).get("batch_size", batch_size)) diff --git a/src/main/python/systemds/scuro/representations/window_aggregation.py b/src/main/python/systemds/scuro/representations/window_aggregation.py index e8189e94312..60b5b467ff3 100644 --- a/src/main/python/systemds/scuro/representations/window_aggregation.py +++ b/src/main/python/systemds/scuro/representations/window_aggregation.py @@ -693,10 +693,11 @@ def _feature_shape(self, in_shape): if isinstance(self.aggregation_function, Aggregation): return in_shape[1:] length = in_shape[0] if in_shape else 0 - approx_window_size = ( - max(1, int(length / self._effective_num_windows(length))) if in_shape else 1 - ) - return self._per_window_feature_shape(approx_window_size) + # Window sizes grow geometrically and _pad_stack pads every window's + # features to the largest one, so size them from the largest window + # (about 4x the mean), not the mean. + max_window_size = int(self._window_sizes(length).max()) if in_shape else 1 + return self._per_window_feature_shape(max_window_size) def get_output_stats(self, input_stats: RepresentationStats) -> tuple: in_shape = tuple(int(s) for s in input_stats.output_shape) diff --git a/src/main/python/systemds/scuro/representations/zero_pad.py b/src/main/python/systemds/scuro/representations/zero_pad.py new file mode 100644 index 00000000000..505545e35a3 --- /dev/null +++ b/src/main/python/systemds/scuro/representations/zero_pad.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 systemds.scuro.modality.transformed import TransformedModality +from systemds.scuro.representations.fusion import Fusion +from systemds.scuro.representations.representation import derive_stats +from systemds.scuro.representations.unimodal import UnimodalRepresentation + + +class ZeroPad(UnimodalRepresentation): + def __init__(self, width=None, input_width=None, params=None): + super().__init__("ZeroPad", None) + params = params or {} + self.width = params.get("width", width) + self.input_width = params.get("input_width", input_width) + + def transform(self, modality, aggregation=None): + padded = TransformedModality( + modality, self, self_contained=modality.self_contained + ) + padded.data = modality.data + Fusion._normalize_for_fusion(padded) + padded.pad(max_len=self.width) + return padded + + def get_output_stats(self, input_stats): + shape = (*input_stats.output_shape[:-1], self.width) + return derive_stats(input_stats, output_shape=shape) + + def get_current_parameters(self): + return {"width": self.width, "input_width": self.input_width} diff --git a/src/main/python/tests/scuro/test_algebra_canonical.py b/src/main/python/tests/scuro/test_algebra_canonical.py new file mode 100644 index 00000000000..f1db174cb4a --- /dev/null +++ b/src/main/python/tests/scuro/test_algebra_canonical.py @@ -0,0 +1,269 @@ +# ------------------------------------------------------------- +# +# 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 json + +import pytest + +from systemds.scuro.drsearch.algebra import ALL_LAWS, canonical_id +from systemds.scuro.drsearch.representation_dag import RepresentationDAGBuilder +from systemds.scuro.representations.aggregated_representation import ( + AggregatedRepresentation, +) +from systemds.scuro.representations.average import Average +from systemds.scuro.representations.bow import BoW +from systemds.scuro.representations.concatenation import Concatenation +from systemds.scuro.representations.hadamard import Hadamard +from systemds.scuro.representations.lstm import LSTM +from systemds.scuro.representations.max import RowMax +from systemds.scuro.representations.sum import Sum +from systemds.scuro.representations.zero_pad import ZeroPad + +IMPLICIT = AggregatedRepresentation( + target_dimensions=1, aggregate_leading=True +).get_current_parameters() +FEATURE_AGG = AggregatedRepresentation(target_dimensions=1).get_current_parameters() +NONE = frozenset() + + +def S(*c, op=Sum): + return (op, {}, *c) + + +def A(*c): + return (Average, {}, *c) + + +def C(*c): + return (Concatenation, {"preserve_leading_axis": False}, *c) + + +def W(*c): + return (Concatenation, {"preserve_leading_axis": True}, *c) + + +def L(*c): + return (LSTM, {}, *c) + + +def G(c, params=IMPLICIT): + return (AggregatedRepresentation, params, c) + + +def P(c): + return (ZeroPad, ZeroPad(8, 3).get_current_parameters(), c) + + +def pushed(c, params): + return (BoW, {"_pushdown_aggregation": params}, c) + + +# (name, left, right, laws, invariant, equal) +CASES = [ + ("leaf", "a", "a", ALL_LAWS, False, True), + ("leaf_index", "a", "b", ALL_LAWS, False, False), + ("params", (BoW, {"x": 1}, "a"), (BoW, {"x": 2}, "a"), ALL_LAWS, False, False), + *[ + ( + f"comm_{op.__name__}", + S("a", "b", op=op), + S("b", "a", op=op), + ALL_LAWS, + False, + True, + ) + for op in (Sum, RowMax, Hadamard) + ], + ("comm_off", S("a", "b"), S("b", "a"), NONE, False, False), + ("assoc", S(S("a", "b"), "c"), S("a", S("b", "c")), ALL_LAWS, False, True), + ("assoc_nary", S(S("a", "b"), "c"), S("a", "b", "c"), ALL_LAWS, False, True), + ("assoc_off", S(S("a", "b"), "c"), S("a", "b", "c"), {"comm"}, False, False), + ( + "assoc_other_op", + S(S("a", "b", op=RowMax), "c"), + S("a", "b", "c"), + ALL_LAWS, + False, + False, + ), + ("concat_order_kept", C("a", "b"), C("b", "a"), ALL_LAWS, False, False), + ("concat_perm", C("a", "b"), C("b", "a"), ALL_LAWS, True, True), + ("concat_assoc", C(C("a", "b"), "c"), C("a", C("b", "c")), ALL_LAWS, False, True), + ("window_concat_never_sorted", W("a", "b"), W("b", "a"), ALL_LAWS, True, False), + ( + "window_concat_assoc", + W(W("a", "b"), "c"), + W("a", "b", "c"), + ALL_LAWS, + True, + True, + ), + ("window_concat_not_concat", W("a", "b"), C("a", "b"), ALL_LAWS, True, False), + ("lstm_never_sorted", L("a", "b"), L("b", "a"), ALL_LAWS, True, False), + ( + "no_comm_below_lstm", + L(S("a", "b"), "c"), + L(S("b", "a"), "c"), + ALL_LAWS, + True, + False, + ), + ( + "no_perm_below_lstm", + L(C("a", "b"), "c"), + L(C("b", "a"), "c"), + ALL_LAWS, + True, + False, + ), + ( + "pad_below_lstm", + L(S(P("a"), "b"), "c"), + L(S("a", "b"), "c"), + ALL_LAWS, + False, + True, + ), + ( + "comm_above_lstm", + S(L("a", "b"), "c"), + S("c", L("a", "b")), + ALL_LAWS, + False, + True, + ), + ("avg_flatten", A(A("a", "b"), "c"), A("c", A("b", "a")), ALL_LAWS, False, True), + ("avg_not_assoc", A(A("a", "b"), "c"), A("a", A("b", "c")), ALL_LAWS, False, False), + ( + "avg_weights_add", + A(A("a", "b"), A("a", "c")), + A("a", A("b", "c")), + ALL_LAWS, + False, + True, + ), + ( + "avg_flatten_off", + A(A("a", "b"), A("a", "c")), + A("a", A("b", "c")), + {"comm"}, + False, + False, + ), + ( + "agg_pushdown", + pushed("a", FEATURE_AGG), + G((BoW, {}, "a"), FEATURE_AGG), + ALL_LAWS, + False, + True, + ), + ( + "agg_pushdown_off", + pushed("a", FEATURE_AGG), + G((BoW, {}, "a"), FEATURE_AGG), + NONE, + False, + False, + ), + ("implicit_agg", S(G("a"), "b"), S("a", "b"), ALL_LAWS, False, True), + ( + "implicit_agg_pushed", + S(pushed("a", IMPLICIT), "b"), + S((BoW, {}, "a"), "b"), + ALL_LAWS, + False, + True, + ), + ( + "implicit_agg_off", + S(G("a"), "b"), + S("a", "b"), + ALL_LAWS - {"implicit_agg"}, + False, + False, + ), + ( + "implicit_agg_other_params", + S(G("a", FEATURE_AGG), "b"), + S("a", "b"), + ALL_LAWS, + False, + False, + ), + ("implicit_agg_window_concat", W(G("a"), "b"), W("a", "b"), ALL_LAWS, False, False), + ("implicit_agg_lstm", L(G("a"), "b"), L("a", "b"), ALL_LAWS, False, False), + ("implicit_agg_unary", (BoW, {}, G("a")), (BoW, {}, "a"), ALL_LAWS, False, False), + ("implicit_pad", S(P("a"), "b"), S("a", "b"), ALL_LAWS, False, True), + ("implicit_pad_avg", A(P("a"), "b"), A("a", "b"), ALL_LAWS, False, True), + ( + "implicit_pad_off", + S(P("a"), "b"), + S("a", "b"), + ALL_LAWS - {"implicit_pad"}, + False, + False, + ), + ("implicit_pad_concat", C(P("a"), "b"), C("a", "b"), ALL_LAWS, True, False), +] + + +def build(expr): + builder = RepresentationDAGBuilder() + leaves = {} + + def make(e): + if isinstance(e, str): + if e not in leaves: + i = "abc".index(e) + leaves[e] = builder.create_leaf_node(f"m{i}", i) + return leaves[e] + op, params, *children = e + return builder.create_operation_node(op, [make(c) for c in children], params) + + return builder.build(make(expr)) + + +@pytest.mark.parametrize( + "left,right,laws,invariant,equal", [c[1:] for c in CASES], ids=[c[0] for c in CASES] +) +def test_canonical_pairs(left, right, laws, invariant, equal): + laws = frozenset(laws) + left, right = build(left), build(right) + ids = [canonical_id(dag, laws, invariant) for dag in (left, right)] + assert (ids[0] == ids[1]) == equal + for dag, dag_id in zip((left, right), ids): + spec = json.loads(json.dumps(dag.to_spec())) + assert canonical_id(spec, laws, invariant) == dag_id + + +def test_dag_method_defaults(): + dag = build(C("a", S("c", "b"))) + assert dag.canonical_id() == canonical_id(dag, ALL_LAWS, False) + assert dag.canonical_id(laws=NONE) != dag.canonical_id() + + +def test_sub_plan_signature(): + dag = build(C("a", S("b", "c"))) + sub = dag.get_node_by_id(dag.get_node_by_id(dag.root_node_id).inputs[1]) + assert canonical_id(dag, ALL_LAWS, False, sub.node_id) == canonical_id( + build(S("c", "b")), ALL_LAWS, False + ) diff --git a/src/main/python/tests/scuro/test_algebra_cse.py b/src/main/python/tests/scuro/test_algebra_cse.py new file mode 100644 index 00000000000..df86a55bd21 --- /dev/null +++ b/src/main/python/tests/scuro/test_algebra_cse.py @@ -0,0 +1,82 @@ +# ------------------------------------------------------------- +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# +# ------------------------------------------------------------- + +import unittest +from types import SimpleNamespace + +from systemds.scuro.drsearch.algebra import AlgebraConfig +from systemds.scuro.drsearch.unimodal_optimizer import UnimodalOptimizer +from systemds.scuro.representations.concatenation import Concatenation +from systemds.scuro.representations.sum import Sum + +LEGACY = AlgebraConfig() +LAW_AWARE = AlgebraConfig(cse_law_aware=True) +FULL = AlgebraConfig.full() + + +def _builder(algebra): + optimizer = UnimodalOptimizer( + [SimpleNamespace(modality_id=0)], + [SimpleNamespace(model=SimpleNamespace(name="task"), expected_dim=1)], + debug=False, + window_combination_chains=2, + algebra=algebra, + ) + builder = optimizer.builders[0] + return optimizer, builder, [builder.create_leaf_node(0, i) for i in range(2)] + + +def _nodes(builder, op): + return [n for n in builder.global_nodes if n.operation is op] + + +class TestLawAwareCSE(unittest.TestCase): + def test_window_concatenation_rotations(self): + for algebra, expected in ((LEGACY, 1), (LAW_AWARE, 2), (FULL, 2)): + with self.subTest(algebra=algebra): + optimizer, builder, windows = _builder(algebra) + optimizer.combine_windowed_representations(builder, windows) + self.assertEqual(len(_nodes(builder, Concatenation)), expected) + prevented = builder.algebra_stats["cse_unsafe_merges_prevented"] + self.assertEqual(prevented, int(algebra.cse_law_aware)) + + def test_commutative_fusion_merges_in_both_modes(self): + cases = [ + (Sum, {}, LEGACY, 1), + (Sum, {}, LAW_AWARE, 1), + (Concatenation, {"preserve_leading_axis": False}, LEGACY, 1), + (Concatenation, {"preserve_leading_axis": False}, LAW_AWARE, 2), + (Concatenation, {"preserve_leading_axis": False}, FULL, 1), + ] + for op, params, algebra, expected in cases: + with self.subTest(op=op.__name__, algebra=algebra): + _, builder, (a, b) = _builder(algebra) + builder.create_operation_node(op, [a, b], params) + builder.create_operation_node(op, [b, a], params) + self.assertEqual(len(_nodes(builder, op)), expected) + merged_by_law = builder.algebra_stats["cse_merges_by_law"] + self.assertEqual( + merged_by_law, int(algebra.cse_law_aware and expected == 1) + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/main/python/tests/scuro/test_algebra_exhaustive.py b/src/main/python/tests/scuro/test_algebra_exhaustive.py new file mode 100644 index 00000000000..249720ed8c2 --- /dev/null +++ b/src/main/python/tests/scuro/test_algebra_exhaustive.py @@ -0,0 +1,76 @@ +# ------------------------------------------------------------- +# +# 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 tempfile +import unittest + +from systemds.scuro.drsearch.algebra import ALL_LAWS, AlgebraConfig, canonical_id +from systemds.scuro.drsearch.multimodal_optimizer import MultimodalOptimizer +from systemds.scuro.representations.average import Average +from systemds.scuro.representations.concatenation import Concatenation +from systemds.scuro.representations.sum import Sum +from tests.scuro.test_multimodal_ga_optimizer import ( + _FakeModality, + _FakeUnimodalResults, + _make_task, +) + +PERMUTATIONS = AlgebraConfig(exhaustive_dedup_permutations=True) +CANONICAL = AlgebraConfig( + laws=ALL_LAWS, + feature_order_invariant=True, + exhaustive_dedup_permutations=True, + exhaustive_canonical=True, +) + + +def _fusion_dags(algebra, n=3): + ids = [f"m{i}" for i in range(n)] + with tempfile.TemporaryDirectory() as tmp: + optimizer = MultimodalOptimizer( + [_FakeModality(m) for m in ids], + _FakeUnimodalResults({m: [object()] for m in ids}), + [_make_task()], + debug=False, + checkpoint_dir=tmp, + algebra=algebra, + ) + optimizer.fusion_operators = [Concatenation, Sum, Average] + dags = list(optimizer._generate_fusion_dags(ids, {m: 0 for m in ids})) + return dags, optimizer.algebra_stats + + +class TestExhaustiveDedup(unittest.TestCase): + def test_counts(self): + legacy, _ = _fusion_dags(AlgebraConfig()) + ordered, stats = _fusion_dags(PERMUTATIONS) + self.assertEqual((len(legacy), len(ordered)), (648, 108)) + self.assertEqual(stats["exhaustive_generated"], 648) + self.assertEqual(stats["exhaustive_after_permutation_dedup"], 108) + + def test_canonical_matches_brute_force(self): + ordered, _ = _fusion_dags(PERMUTATIONS) + ids = {canonical_id(d, ALL_LAWS, True) for d in ordered} + canonical, stats = _fusion_dags(CANONICAL) + self.assertEqual( + sorted(canonical_id(d, ALL_LAWS, True) for d in canonical), sorted(ids) + ) + self.assertEqual(stats["exhaustive_after_canonical"], len(ids)) diff --git a/src/main/python/tests/scuro/test_algebra_explicit_prep.py b/src/main/python/tests/scuro/test_algebra_explicit_prep.py new file mode 100644 index 00000000000..3f57605a813 --- /dev/null +++ b/src/main/python/tests/scuro/test_algebra_explicit_prep.py @@ -0,0 +1,218 @@ +# ------------------------------------------------------------- +# +# 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 tempfile +from collections import Counter +from types import SimpleNamespace + +import numpy as np +import pytest + +from systemds.scuro.drsearch.algebra import ( + ALL_LAWS, + AlgebraConfig, + canonical_id, + count_explicit_prep, + sample_stats, + typecheck, +) +from systemds.scuro.drsearch.operator_registry import Registry +from systemds.scuro.drsearch.multimodal_ga_optimizer import ( + DagGenome, + MultimodalDeapOptimizer, +) +from systemds.scuro.drsearch.multimodal_optimizer import MultimodalOptimizer +from systemds.scuro.drsearch.representation_dag import CSEAwareDAGBuilder +from systemds.scuro.modality.transformed import TransformedModality +from systemds.scuro.modality.type import ModalityType +from systemds.scuro.representations.aggregated_representation import ( + AggregatedRepresentation, +) +from systemds.scuro.representations.average import Average +from systemds.scuro.representations.concatenation import Concatenation +from systemds.scuro.representations.hadamard import Hadamard +from systemds.scuro.representations.representation import RepresentationStats +from systemds.scuro.representations.sum import Sum +from systemds.scuro.representations.zero_pad import ZeroPad +from tests.scuro.data_generator import ModalityRandomDataGenerator +from tests.scuro.test_multimodal_ga_optimizer import ( + _FakeModality, + _FakeUnimodalResults, + _make_task, +) + +AGG = AlgebraConfig(explicit_fusion_prep=True) +PAD = AlgebraConfig(explicit_padding="pad") +BOTH = AlgebraConfig(explicit_fusion_prep=True, explicit_padding="pad") +FUSIONS = [Concatenation, Sum, Average, Hadamard] +ALIGNED = {Sum, Average, Hadamard} + + +def _reps(): + """Modality 0 is a (4, 8) sequence; 1 and 2 are vectors of width 3 and 8.""" + gen = ModalityRandomDataGenerator() + data, _ = gen.create_2d_modality(6, (4, 8)) + seq = TransformedModality(gen, "test") + seq.data = data + reps = {0: seq} + for mid, width in ((1, 3), (2, 8)): + reps[mid] = gen.create1DModality(6, width, ModalityType.TIMESERIES) + for mid, rep in reps.items(): + rep.modality_id = mid + return {mid: [rep] for mid, rep in reps.items()} + + +def _expected(algebra, op, leaves): + out = [] + if algebra.explicit_fusion_prep and 0 in leaves: + out.append(("AggregatedRepresentation", 0)) + # Padding needs every operand to be a vector first. + if algebra.explicit_padding == "pad" and op in ALIGNED: + if algebra.explicit_fusion_prep or 0 not in leaves: + out.append(("ZeroPad", 1)) + return out + + +def _prep_nodes(dag): + return sorted( + (n.operation.__name__, dag.get_node_by_id(n.inputs[0]).modality_id) + for n in dag.filter_connected_nodes(dag.nodes) + if n.operation in (AggregatedRepresentation, ZeroPad) + ) + + +def _fused(dag, reps): + out = dag.execute([m for ms in reps.values() for m in ms], enable_cache=False) + return np.asarray(out[dag.root_node_id] if isinstance(out, dict) else out.data) + + +def _assert_equivalent(legacy, explicit, reps, expected): + assert _prep_nodes(legacy) == [] + assert _prep_nodes(explicit) == expected + assert canonical_id(legacy, ALL_LAWS, False) == canonical_id( + explicit, ALL_LAWS, False + ) + np.testing.assert_array_equal(_fused(legacy, reps), _fused(explicit, reps)) + + +def _ga_dag(algebra, reps, genome): + optimizer = MultimodalDeapOptimizer( + [_FakeModality(m) for m in reps], + _FakeUnimodalResults(reps), + [_make_task()], + debug=False, + algebra=algebra, + ) + optimizer._current_task_name = "task0" + return optimizer._genome_to_dag(genome) + + +@pytest.mark.parametrize("algebra", [AGG, PAD, BOTH]) +@pytest.mark.parametrize("op", FUSIONS) +@pytest.mark.parametrize("leaves", [(0, 1), (1, 2)]) +def test_ga_genome_to_dag(algebra, op, leaves): + reps = _reps() + genome = DagGenome([(m, 0) for m in leaves], (0, 1), {"": op}) + _assert_equivalent( + _ga_dag(AlgebraConfig(), reps, genome), + _ga_dag(algebra, reps, genome), + reps, + _expected(algebra, op, leaves), + ) + + +@pytest.mark.parametrize("algebra", [AGG, BOTH]) +def test_exhaustive_enumerator(algebra): + reps = _reps() + runs = [] + for explicit in (AlgebraConfig(), algebra): + with tempfile.TemporaryDirectory() as tmp: + optimizer = MultimodalOptimizer( + [_FakeModality(m) for m in reps], + _FakeUnimodalResults(reps), + [_make_task()], + debug=False, + checkpoint_dir=tmp, + algebra=AlgebraConfig( + exhaustive_dedup_permutations=True, + explicit_fusion_prep=explicit.explicit_fusion_prep, + explicit_padding=explicit.explicit_padding, + ), + ) + optimizer.fusion_operators = FUSIONS + runs.append( + list(optimizer._generate_fusion_dags([0, 1], {0: 0, 1: 0}, "task0")) + ) + assert len(runs[0]) == len(runs[1]) == 2 * len(FUSIONS) + for legacy, explicit, op in zip(*runs, FUSIONS * 2): + _assert_equivalent(legacy, explicit, reps, _expected(algebra, op, (0, 1))) + assert optimizer.algebra_stats["explicit_prep_plans"] == 2 * len(FUSIONS) + + +def test_counts_and_typechecks_padded_plan(): + reps = _reps() + dag = _ga_dag(BOTH, reps, DagGenome([(0, 0), (1, 0)], (0, 1), {"": Sum})) + stats = Counter() + count_explicit_prep(stats, dag) + assert stats == { + "explicit_prep_plans": 1, + "explicit_padding_plans": 1, + "explicit_aggregations": 1, + "explicit_paddings": 1, + "explicit_padding_ratio_sum": 8 / 3, + "explicit_padding_ratio_max": 8 / 3, + } + leaves = {(m, 0): r[0] for m, r in reps.items()} + types = {k: r.modality_type for k, r in leaves.items()} + assert typecheck(dag, types, {k: sample_stats(r) for k, r in leaves.items()}) == [] + + +def test_zero_pad_is_not_registered(): + registry = Registry() + lists = [ + registry._representations, + registry._context_operators, + registry._dimensionality_reduction_operators, + ] + assert all(ZeroPad not in ops for lst in lists for ops in lst.values()) + + +def test_unbound_genome_to_dag_and_bare_checker(): + reps = _reps() + genome = DagGenome([(0, 0), (1, 0)], (0, 1), {"": Sum}) + assert _prep_nodes(MultimodalDeapOptimizer._genome_to_dag(None, genome)) == [] + checker = object.__new__(MultimodalDeapOptimizer) + checker.max_feature_dim = None + checker.k_best_representations = {"task0": reps} + assert checker._oversized_reason("task0", genome) is None + + +@pytest.mark.parametrize( + "params, wrapped", [({}, True), ({"preserve_leading_axis": True}, False)] +) +def test_cse_builder(params, wrapped): + leaf = SimpleNamespace(modality_id="t", stats=RepresentationStats(5, (4, 8))) + builder = CSEAwareDAGBuilder(BOTH, leaves=[leaf]) + leaf_id = builder.create_leaf_node("t") + root = builder.create_operation_node(Concatenation, [leaf_id, leaf_id], params) + builder.build(root) + assert (builder.get_node(root).inputs != [leaf_id, leaf_id]) == wrapped + assert builder.algebra_stats["explicit_aggregations"] == 2 * wrapped diff --git a/src/main/python/tests/scuro/test_algebra_ga_dedup.py b/src/main/python/tests/scuro/test_algebra_ga_dedup.py new file mode 100644 index 00000000000..8bc457dd530 --- /dev/null +++ b/src/main/python/tests/scuro/test_algebra_ga_dedup.py @@ -0,0 +1,65 @@ +# ------------------------------------------------------------- +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# +# ------------------------------------------------------------- + +import unittest +from unittest.mock import patch + +from systemds.scuro.drsearch.algebra import AlgebraConfig +from systemds.scuro.drsearch.multimodal_ga_optimizer import DagGenome +from systemds.scuro.representations.sum import Sum +from tests.scuro.test_multimodal_ga_optimizer import MODULE, _make_optimizer + + +def _body(_dag, _task, _modalities, _objective_specs): + return (0.5,), { + "train_score": {}, + "val_score": {"accuracy": 0.5}, + "test_score": {}, + "runtime": 2.0, + "task_time": 1.0, + "representation_time": 1.0, + } + + +class TestCanonicalGADedup(unittest.TestCase): + def test_commuted_sum_evaluated_once_when_on(self): + leaves = [("m0", 0), ("m1", 0)] + genomes = [DagGenome(leaves, tree, {"": Sum}) for tree in ((0, 1), (1, 0))] + for algebra, evaluations in ( + (AlgebraConfig(), 2), + (AlgebraConfig(laws=frozenset({"comm"}), ga_canonical_dedup=True), 1), + ): + with self.subTest(algebra=algebra): + optimizer, _, task = _make_optimizer( + n_modalities=2, max_workers=1, algebra=algebra + ) + optimizer._current_task_name = task.model.name + population = [optimizer._make_individual(g) for g in genomes] + with patch(f"{MODULE}._evaluate_genome_body", side_effect=_body): + optimizer._evaluate_population(population, task) + results = optimizer.optimization_results["task0"] + self.assertEqual(len(results), evaluations) + avoided = 2 - evaluations + stats = optimizer.algebra_stats + self.assertEqual(stats["ga_evaluations_avoided"], avoided) + self.assertEqual(stats["ga_seconds_avoided"], 2.0 * avoided) + self.assertEqual(stats["ga_unique_legacy"], 2 * avoided) + self.assertEqual(stats["ga_unique_canonical"], avoided) diff --git a/src/main/python/tests/scuro/test_algebra_laws.py b/src/main/python/tests/scuro/test_algebra_laws.py new file mode 100644 index 00000000000..9ef8f7d2c61 --- /dev/null +++ b/src/main/python/tests/scuro/test_algebra_laws.py @@ -0,0 +1,136 @@ +# ------------------------------------------------------------- +# +# 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 numpy as np +import pytest + +from systemds.scuro.drsearch.operator_registry import Registry +from systemds.scuro.modality.modality import Modality +from systemds.scuro.modality.transformed import TransformedModality +from systemds.scuro.modality.type import ModalityType +from systemds.scuro.representations.concatenation import Concatenation + +OPERATORS = [(op, {}) for op in Registry().get_fusion_operators()] + [ + (Concatenation, {"preserve_leading_axis": True}) +] +DECLARED = [ + (op, params, law) + for op, params in OPERATORS + for law in sorted(op.algebraic_laws(params)) +] +WIDTHS = {"equal": (6, 6, 6), "unequal": (6, 3, 9)} +DTYPES = { + "float32": (np.float32, np.float32, np.float32), + "float32/64": (np.float32, np.float64, np.float32), + "int/float": (np.int64, np.float32, np.float64), +} + + +def _operands(widths, dtypes, rank): + rng = np.random.default_rng(0) + shape = (5,) + (4,) * (rank - 1) + return [ + (rng.standard_normal(shape + (w,)) * 5).astype(dtype) + for w, dtype in zip(widths, dtypes) + ] + + +def _fuse(op, params, *operands): + modalities = [] + for data in operands: + metadata = [ModalityType.EMBEDDING.create_metadata(d) for d in data] + modality = TransformedModality( + Modality(ModalityType.EMBEDDING, 0, metadata), "operand" + ) + modality.data = data + modalities.append(modality) + return np.asarray(op(params).transform(modalities)) + + +def _assert_same(x, y, exact): + assert x.shape == y.shape + if exact: + np.testing.assert_array_equal(x, y) + else: + np.testing.assert_allclose(x, y, rtol=1e-5, atol=1e-5) + + +def _prepared(x, width): + # What _prepare_modalities hands to execute: hidden mean aggregation, then zero-padding. + x = x.mean(axis=1) if x.ndim == 3 else x + return np.pad(x, ((0, 0), (0, width - x.shape[-1]))) + + +def _comm(f, a, b, c, exact): + _assert_same(f(a, b), f(b, a), exact) + + +def _assoc(f, a, b, c, exact): + left = f(f(a, b), c) + _assert_same(left, f(a, f(b, c)), exact) + _assert_same(left, f(a, b, c), exact) + + +def _avg_flatten(f, a, b, c, exact): + width = max(x.shape[-1] for x in (a, b, c)) + expected = sum(w * _prepared(x, width) for w, x in ((0.25, a), (0.25, b), (0.5, c))) + _assert_same(f(f(a, b), c), expected, exact) + + +def _concat_perm(f, a, b, c, exact): + _assert_same(f(a, b), np.roll(f(b, a), -b.shape[-1], axis=-1), exact=True) + + +LAW_CHECKS = { + "comm": _comm, + "assoc": _assoc, + "avg_flatten": _avg_flatten, + "concat_perm": _concat_perm, +} + + +def test_every_declared_law_is_checked(): + assert {law for _, _, law in DECLARED} <= LAW_CHECKS.keys() + + +@pytest.mark.parametrize("rank", [1, 2]) +@pytest.mark.parametrize("dtypes", DTYPES.values(), ids=DTYPES.keys()) +@pytest.mark.parametrize("widths", WIDTHS.values(), ids=WIDTHS.keys()) +@pytest.mark.parametrize( + "op,params,law", + DECLARED, + ids=[f"{op.__name__}{params or ''}-{law}" for op, params, law in DECLARED], +) +def test_declared_law_holds(op, params, law, widths, dtypes, rank): + LAW_CHECKS[law]( + lambda *xs: _fuse(op, params, *xs), + *_operands(widths, dtypes, rank), + exact=op is Concatenation, + ) + + +def test_window_axis_concatenation_is_order_sensitive(): + params = {"preserve_leading_axis": True} + a, b, _ = _operands(WIDTHS["equal"], DTYPES["float32"], rank=2) + assert not np.array_equal( + _fuse(Concatenation, params, a, b), _fuse(Concatenation, params, b, a) + ) + assert not {"comm", "concat_perm"} & Concatenation.algebraic_laws(params) diff --git a/src/main/python/tests/scuro/test_algebra_order_invariant.py b/src/main/python/tests/scuro/test_algebra_order_invariant.py new file mode 100644 index 00000000000..1f6c1264300 --- /dev/null +++ b/src/main/python/tests/scuro/test_algebra_order_invariant.py @@ -0,0 +1,56 @@ +# ------------------------------------------------------------- +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# +# ------------------------------------------------------------- + +import unittest +from types import SimpleNamespace + +from systemds.scuro.drsearch.algebra import AlgebraConfig, order_invariant +from systemds.scuro.models.model import Model + + +def _tasks(*flags): + return [ + SimpleNamespace(model=SimpleNamespace(feature_order_invariant=f)) for f in flags + ] + + +class TestOrderInvariant(unittest.TestCase): + def test_model_default_is_unknown(self): + self.assertIsNone(Model("m").feature_order_invariant) + + def test_cases(self): + cases = [ + (None, _tasks(True, True), True), + (None, _tasks(True, None), False), + (None, _tasks(True, False), False), + (None, [], False), + (None, [SimpleNamespace(model=object())], False), + (True, _tasks(False), True), + (False, _tasks(True), False), + ] + for override, tasks, expected in cases: + with self.subTest(override=override, tasks=tasks): + config = AlgebraConfig(feature_order_invariant=override) + self.assertIs(order_invariant(config, tasks), expected) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/main/python/tests/scuro/test_algebra_plan_counts.py b/src/main/python/tests/scuro/test_algebra_plan_counts.py new file mode 100644 index 00000000000..8861d89b766 --- /dev/null +++ b/src/main/python/tests/scuro/test_algebra_plan_counts.py @@ -0,0 +1,71 @@ +# ------------------------------------------------------------- +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# +# ------------------------------------------------------------- + +import unittest + +from systemds.scuro.drsearch.algebra import ( + ALL_LAWS, + _brute_force_canonical, + count_fusion_plans, + count_search_space, +) +from systemds.scuro.representations.average import Average +from systemds.scuro.representations.concatenation import Concatenation +from systemds.scuro.representations.hadamard import Hadamard +from systemds.scuro.representations.lstm import LSTM +from systemds.scuro.representations.max import RowMax +from systemds.scuro.representations.sum import Sum + +COMM = frozenset({"comm", "concat_perm"}) +COMM_ASSOC = [Sum, Hadamard, RowMax] +MIXED = [Sum, Average, Concatenation, LSTM] + + +class TestPlanCounts(unittest.TestCase): + def test_closed_forms_match_brute_force(self): + for ops in (COMM_ASSOC, MIXED): + for n in (2, 3, 4): + with self.subTest(ops=[op.__name__ for op in ops], n=n): + counts = count_fusion_plans(n, ops, ALL_LAWS, True) + for key, laws in ( + ("ordered", frozenset()), + ("commutative", COMM), + ("canonical", ALL_LAWS), + ): + brute = _brute_force_canonical(n, ops, laws, True) + self.assertEqual(counts[key], brute, key) + + def test_known_values(self): + counts = count_fusion_plans(3, [Concatenation, Sum, Average], ALL_LAWS, True) + self.assertEqual((counts["legacy_exhaustive"], counts["ordered"]), (648, 108)) + # comm alone divides the ordered trees by 2^(n-1) + self.assertEqual( + count_fusion_plans(4, COMM_ASSOC, COMM, True)["commutative"], 3240 // 8 + ) + self.assertIsNone(count_fusion_plans(5, MIXED, ALL_LAWS, True)["canonical"]) + + def test_search_space_sums_over_subsets(self): + per_size = {n: count_fusion_plans(n, MIXED, ALL_LAWS, False) for n in (2, 3)} + total = count_search_space([2, 2, 3], MIXED, ALL_LAWS, False) + for key, value in total.items(): + self.assertEqual( + value, (4 + 6 + 6) * per_size[2][key] + 12 * per_size[3][key], key + ) diff --git a/src/main/python/tests/scuro/test_algebra_signatures.py b/src/main/python/tests/scuro/test_algebra_signatures.py new file mode 100644 index 00000000000..9cd10e02107 --- /dev/null +++ b/src/main/python/tests/scuro/test_algebra_signatures.py @@ -0,0 +1,86 @@ +# ------------------------------------------------------------- +# +# 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 pytest + +from systemds.scuro.drsearch.algebra import ANY_TYPE, operator_signature +from systemds.scuro.drsearch.operator_registry import Registry, load_builtin_operators +from systemds.scuro.modality.type import ModalityType +from systemds.scuro.representations.aggregated_representation import ( + AggregatedRepresentation, +) +from systemds.scuro.representations.concatenation import Concatenation +from systemds.scuro.representations.lstm import LSTM +from systemds.scuro.representations.mlp_averaging import MLPAveraging +from systemds.scuro.representations.sum import Sum +from systemds.scuro.representations.window_aggregation import WindowAggregation + +load_builtin_operators(ANY_TYPE, include_fusion=True) +REGISTERED = sorted( + { + (op, kind, modality) + for kind, lists in [ + ("encode", Registry._representations), + ("context", Registry._context_operators), + ("reduce", Registry._dimensionality_reduction_operators), + ] + for modality, ops in lists.items() + for op in ops + }, + key=lambda t: (t[0].__name__, t[1], t[2].name), +) + [(op, "fuse", None) for op in Registry().get_fusion_operators()] + + +@pytest.mark.parametrize( + "op, kind, modality", + REGISTERED, + ids=[ + f"{op.__name__}-{kind}-{getattr(m, 'name', '')}" for op, kind, m in REGISTERED + ], +) +def test_every_registered_operator_has_a_signature(op, kind, modality): + sig = operator_signature(op) + assert sig.kind == kind + assert modality is None or modality in sig.input_types + if kind in ("encode", "fuse"): + assert isinstance(sig.output_type, ModalityType) + + +@pytest.mark.parametrize( + "op, params, expected", + [ + (Sum, {}, ("fuse", "n", "vector", "vector")), + ( + Concatenation, + {"preserve_leading_axis": False}, + ("fuse", "n", "vector", "vector"), + ), + (Concatenation, {"preserve_leading_axis": True}, ("fuse", "n", "seq", "seq")), + (LSTM, {}, ("fuse", "n", "seq", "vector")), + (AggregatedRepresentation, {}, ("aggregate", 1, "seq", "vector")), + (WindowAggregation, {}, ("context", 1, "seq", "windows")), + (MLPAveraging, {}, ("reduce", 1, "vector", "vector")), + (None, {}, ("task", 1, "any", "any")), + ], +) +def test_rank_rules(op, params, expected): + sig = operator_signature(op, params) + assert (sig.kind, sig.arity, sig.input_rank, sig.output_rank) == expected diff --git a/src/main/python/tests/scuro/test_algebra_stats.py b/src/main/python/tests/scuro/test_algebra_stats.py new file mode 100644 index 00000000000..5c18d50756f --- /dev/null +++ b/src/main/python/tests/scuro/test_algebra_stats.py @@ -0,0 +1,52 @@ +# ------------------------------------------------------------- +# +# 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 numpy as np +import pytest + +from systemds.scuro.drsearch.operator_registry import Registry +from systemds.scuro.representations.representation import RepresentationStats +from tests.scuro.test_algebra_laws import _fuse + +# Other test modules register fusion operators of their own. +OPERATORS = [ + op + for op in Registry().get_fusion_operators() + if op.__module__.startswith("systemds.") and not op().needs_training +] +OPERANDS = { + "equal": [(5, 6), (5, 6)], + "unequal": [(5, 6), (5, 3)], + "sequence": [(5, 4, 6), (5, 3)], +} + + +@pytest.mark.parametrize("case", OPERANDS) +@pytest.mark.parametrize("op", OPERATORS, ids=lambda op: op.__name__) +def test_predicted_stats_match_execution(op, case): + rng = np.random.default_rng(0) + operands = [ + rng.standard_normal(shape).astype(np.float32) for shape in OPERANDS[case] + ] + stats = [RepresentationStats(len(x), x.shape[1:]) for x in operands] + fused = _fuse(op, {}, *operands) + assert tuple(op().get_output_stats(stats).output_shape) == fused.shape[1:] + assert op().estimate_peak_memory_bytes(stats)["cpu_peak_bytes"] > fused.nbytes diff --git a/src/main/python/tests/scuro/test_algebra_subtree_sharing.py b/src/main/python/tests/scuro/test_algebra_subtree_sharing.py new file mode 100644 index 00000000000..1051c111521 --- /dev/null +++ b/src/main/python/tests/scuro/test_algebra_subtree_sharing.py @@ -0,0 +1,115 @@ +# ------------------------------------------------------------- +# +# 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 unittest.mock import patch + +import numpy as np +import pytest + +from systemds.scuro.drsearch import multimodal_ga_optimizer as ga +from systemds.scuro.drsearch.algebra import AlgebraConfig +from systemds.scuro.drsearch.representation_dag import RepresentationDAGBuilder +from systemds.scuro.drsearch.trial_result_cache import TrialResultCache +from systemds.scuro.representations.concatenation import Concatenation +from systemds.scuro.representations.lstm import LSTM +from tests.scuro.data_generator import TestSVM, TestTask +from tests.scuro.test_multimodal_ga_optimizer import ( + _build_real_optimizer, + _make_real_representation, +) + +LSTM_PARAMS = {"width": 8, "epochs": 2} +SPECS = [("accuracy", "max")] + + +class _RecordingSVM(TestSVM): + inputs = [] + + def fit(self, X, y, X_test, y_test): + self.inputs.append(np.array(X)) + return super().fit(X, y, X_test, y_test) + + +@pytest.fixture +def setup(monkeypatch): + monkeypatch.setattr(ga, "_SUBTREE_CACHE", TrialResultCache(2**30, 0.0)) + task = TestTask("share_task", "share_model", 40) + task.model = _RecordingSVM("share_model") + modalities = [_make_real_representation(i, 40, 4) for i in range(3)] + _RecordingSVM.inputs = [] + return task, modalities + + +def _dag(leaves, root_op=None): + b = RepresentationDAGBuilder() + ids = [b.create_leaf_node(i, 0) for i in leaves] + node = b.create_operation_node(LSTM, ids[:2], dict(LSTM_PARAMS)) + if root_op is not None: + node = b.create_operation_node(root_op, [node, ids[2]], {}) + return b.build(node) + + +def _evaluate(task, modalities, dag, share): + _RecordingSVM.inputs = [] + _, payload = ga._evaluate_genome_body(dag, task, modalities, SPECS, share) + return payload, _RecordingSVM.inputs + + +def test_fitted_subtree_is_reused_bit_identically(setup): + task, modalities = setup + _evaluate(task, modalities, _dag([0, 1]), True) + shared, shared_inputs = _evaluate( + task, modalities, _dag([0, 1, 2], Concatenation), True + ) + _, fresh_inputs = _evaluate(task, modalities, _dag([0, 1, 2], Concatenation), False) + + assert shared["subtree_cache"]["hits"] == len(task.evaluation_splits) + assert len(shared_inputs) == len(fresh_inputs) == len(task.evaluation_splits) + assert all(np.array_equal(a, b) for a, b in zip(shared_inputs, fresh_inputs)) + + +def test_operand_order_is_part_of_the_key(setup): + task, modalities = setup + _evaluate(task, modalities, _dag([0, 1]), True) + swapped, _ = _evaluate(task, modalities, _dag([1, 0]), True) + assert swapped["subtree_cache"]["hits"] == 0 + + +@pytest.mark.parametrize("share", [False, True]) +def test_switch_controls_call_and_stats(share): + optimizer, task = _build_real_optimizer( + fusion_ops=[LSTM], algebra=AlgebraConfig(share_fused_subtrees=share) + ) + genome = optimizer._random_genome(task.model.name) + payload = { + "train_score": {}, + "val_score": {"accuracy": 0.5}, + "test_score": {}, + "runtime": 1.0, + "task_time": 0.5, + "representation_time": 0.5, + "subtree_cache": {"hits": 3, "misses": 2, "saved_seconds": 4.0, "bytes": 10}, + } + with patch.object( + ga, "_evaluate_genome_body", return_value=((0.5,), payload) + ) as body: + optimizer._evaluate_genome(genome, task) + assert len(body.call_args.args) == 4 + share + assert optimizer.algebra_stats["subtree_cache_hits"] == 3 diff --git a/src/main/python/tests/scuro/test_algebra_typecheck.py b/src/main/python/tests/scuro/test_algebra_typecheck.py new file mode 100644 index 00000000000..b59f0e25e10 --- /dev/null +++ b/src/main/python/tests/scuro/test_algebra_typecheck.py @@ -0,0 +1,121 @@ +# ------------------------------------------------------------- +# +# 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 types import SimpleNamespace + +import numpy as np +import pytest + +from systemds.scuro.drsearch.algebra import AlgebraConfig, typecheck +from systemds.scuro.drsearch.multimodal_ga_optimizer import ( + DagGenome, + MultimodalDeapOptimizer, +) +from systemds.scuro.drsearch.representation_dag import ( + CSEAwareDAGBuilder, + RepresentationDAGBuilder, +) +from systemds.scuro.modality.type import ModalityType +from systemds.scuro.representations.bow import BoW +from systemds.scuro.representations.concatenation import Concatenation +from systemds.scuro.representations.mfcc import MFCC +from systemds.scuro.representations.representation import RepresentationStats +from systemds.scuro.representations.sum import Sum +from tests.scuro.test_multimodal_ga_optimizer import _FakeUnimodalResults, _make_task + +TEXT, EMB = ModalityType.TEXT, ModalityType.EMBEDDING +SEQ, VEC = RepresentationStats(5, (4, 8)), RepresentationStats(5, (3,)) + + +def _check(op, params, leaves, with_stats=True): + builder = RepresentationDAGBuilder() + ids = [builder.create_leaf_node(f"m{i}", 0) for i in range(len(leaves))] + dag = builder.build(builder.create_operation_node(op, ids, params)) + types = {(f"m{i}", 0): t for i, (t, _) in enumerate(leaves)} + stats = {(f"m{i}", 0): s for i, (_, s) in enumerate(leaves)} + violations = typecheck(dag, types, stats if with_stats else None) + return sorted((v.severity, v.kind) for v in violations) + + +@pytest.mark.parametrize( + "op, params, leaves, expected", + [ + (BoW, {}, [(TEXT, None)], []), + (MFCC, {}, [(TEXT, None)], [("error", "input_type")]), + (Sum, {}, [(EMB, VEC)], [("error", "arity")]), + (BoW, {}, [(TEXT, None), (TEXT, None)], [("error", "arity")]), + (Sum, {}, [(EMB, VEC), (EMB, VEC)], []), + ( + Sum, + {}, + [(EMB, SEQ), (EMB, VEC)], + [("implicit", "hidden_aggregation"), ("implicit", "zero_padding")], + ), + ( + Concatenation, + {}, + [(EMB, SEQ), (EMB, VEC)], + [("implicit", "hidden_aggregation")], + ), + (Concatenation, {"preserve_leading_axis": True}, [(EMB, SEQ), (EMB, SEQ)], []), + ], +) +def test_violations(op, params, leaves, expected): + assert _check(op, params, leaves) == expected + + +def test_shapes_unknown_without_leaf_stats(): + assert _check(Sum, {}, [(EMB, SEQ), (EMB, VEC)], with_stats=False) == [] + + +@pytest.mark.parametrize( + "mode, raises, counted", + [("off", False, 0), ("warn", False, 1), ("strict", True, 1)], +) +def test_builder_modes(mode, raises, counted): + modality = SimpleNamespace(modality_id="t", modality_type=TEXT, stats=None) + builder = CSEAwareDAGBuilder(AlgebraConfig(typecheck=mode), leaves=[modality]) + root = builder.create_operation_node(MFCC, [builder.create_leaf_node("t")], {}) + if raises: + with pytest.raises(ValueError, match="input_type"): + builder.build(root) + else: + builder.build(root) + assert builder.algebra_stats["typecheck_input_type"] == counted + + +def test_ga_counts_hidden_operators(): + reps = { + "m0": [SimpleNamespace(modality_type=EMB, data=np.zeros((5, 4, 8)))], + "m1": [SimpleNamespace(modality_type=EMB, data=np.zeros((5, 3)))], + } + optimizer = MultimodalDeapOptimizer( + [SimpleNamespace(modality_id=m) for m in reps], + _FakeUnimodalResults(reps), + [_make_task()], + debug=False, + algebra=AlgebraConfig(typecheck="strict"), + ) + genome = DagGenome([("m0", 0), ("m1", 0)], (0, 1), {"": Sum}) + assert optimizer._oversized_reason("task0", genome) is None + stats = optimizer.algebra_stats + assert stats["typecheck_hidden_aggregation"] == 1 + assert stats["typecheck_zero_padding:Sum"] == 1 diff --git a/src/main/python/tests/scuro/test_data_loaders.py b/src/main/python/tests/scuro/test_data_loaders.py index 0d5ace01a47..4affe2d3b7f 100644 --- a/src/main/python/tests/scuro/test_data_loaders.py +++ b/src/main/python/tests/scuro/test_data_loaders.py @@ -19,6 +19,8 @@ # # ------------------------------------------------------------- +import json +import os import shutil import unittest import numpy as np @@ -26,6 +28,7 @@ from tests.scuro.data_generator import setup_data from systemds.scuro.dataloader.audio_loader import AudioLoader, AudioStats from systemds.scuro.dataloader.image_loader import ImageLoader, ImageStats +from systemds.scuro.dataloader.json_loader import JSONLoader from systemds.scuro.dataloader.text_loader import TextLoader, TextStats from systemds.scuro.dataloader.video_loader import VideoLoader, VideoStats from systemds.scuro.modality.type import ModalityType @@ -183,5 +186,36 @@ def test_image_loader_stats(self): # self.assertEqual(stats.num_signals, 2) +class TestJSONLoaderFields(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.test_file_path = "test_json_fields_data" + os.makedirs(cls.test_file_path, exist_ok=True) + cls.indices = ["0", "1"] + cls.json_path = os.path.join(cls.test_file_path, "dialogue.json") + with open(cls.json_path, "w") as f: + json.dump( + { + index: {"utterance": f"utt {index}", "context": ["a", "b"]} + for index in cls.indices + }, + f, + ) + + @classmethod + def tearDownClass(cls): + shutil.rmtree(cls.test_file_path) + + def test_single_field(self): + data, _ = JSONLoader(self.json_path, self.indices, "utterance").load() + self.assertEqual(data, ["utt 0", "utt 1"]) + + def test_list_of_fields_is_joined_in_order(self): + loader = JSONLoader(self.json_path, self.indices, ["context", "utterance"]) + data, _ = loader.load() + self.assertEqual(data, ["a b utt 0", "a b utt 1"]) + self.assertEqual(loader.stats.max_length, len("a b utt 0")) + + if __name__ == "__main__": unittest.main() diff --git a/src/main/python/tests/scuro/test_multimodal_ga_optimizer.py b/src/main/python/tests/scuro/test_multimodal_ga_optimizer.py index f79157f0247..e5eac1c5cfb 100644 --- a/src/main/python/tests/scuro/test_multimodal_ga_optimizer.py +++ b/src/main/python/tests/scuro/test_multimodal_ga_optimizer.py @@ -36,6 +36,7 @@ _collect_leaf_indices, _failure_fitness, _objective_value, + _task_uses_gpu, ) from systemds.scuro.drsearch.operator_registry import Registry, register_fusion_operator from systemds.scuro.modality.type import ModalityType @@ -110,10 +111,20 @@ def _fake_success_body(_dag, _task, _modalities, _objective_specs, value=0.5): class _SynchronousPool: instances = [] - def __init__(self, n_workers, dispatch, ctx=None, threads_per_worker=1): + def __init__( + self, + n_workers, + dispatch, + ctx=None, + threads_per_worker=1, + gpu_devices=None, + gpu_slots_per_device=1, + gpu_demand_fraction=1.0, + ): self.n_workers = n_workers self.dispatch = dispatch self.threads_per_worker = threads_per_worker + self.gpu_devices = list(gpu_devices or []) self.pending = None self.next_job_id = 0 self.shutdown_called = False @@ -123,7 +134,18 @@ def __init__(self, n_workers, dispatch, ctx=None, threads_per_worker=1): def has_idle_worker(self): return self.pending is None - def submit(self, kind, payload, gpu_id=None): + @property + def has_idle_cpu_worker(self): + return self.pending is None + + @property + def gpu_worker_devices(self): + return [] + + def has_idle_worker_for(self, gpu_id, allow_gpu_worker_for_cpu=False): + return self.pending is None + + def submit(self, kind, payload, gpu_id=None, allow_gpu_worker_for_cpu=False): job_id = self.next_job_id self.next_job_id += 1 try: @@ -1216,5 +1238,53 @@ def test_real_fusion_multi_objective_end_to_end(self): self.assertGreaterEqual(result.runtime, 0.0) +class _LanePool: + """Idle-worker bookkeeping of PersistentWorkerPool, without processes.""" + + def __init__(self, idle_gpu, idle_cpu): + self.idle_gpu = set(idle_gpu) + self.idle_cpu = idle_cpu + self.gpu_worker_devices = [0, 1] + + @property + def has_idle_cpu_worker(self): + return self.idle_cpu + + def has_idle_worker_for(self, gpu_id, allow_gpu_worker_for_cpu=False): + if gpu_id is not None: + return gpu_id in self.idle_gpu + return self.idle_cpu or (allow_gpu_worker_for_cpu and bool(self.idle_gpu)) + + +class TestGpuLanes(unittest.TestCase): + def _lane(self, pool, uses_gpu): + return MultimodalDeapOptimizer._idle_lane( + SimpleNamespace(_worker_pool=pool), uses_gpu + ) + + def test_gpu_task_prefers_idle_gpu_worker(self): + self.assertEqual(self._lane(_LanePool({1}, idle_cpu=True), True), 1) + + def test_gpu_task_overflows_to_cpu_worker(self): + self.assertIsNone(self._lane(_LanePool(set(), idle_cpu=True), True)) + + def test_cpu_task_may_take_gpu_worker(self): + self.assertIsNone(self._lane(_LanePool({0}, idle_cpu=False), False)) + + def test_no_idle_worker(self): + self.assertIs(self._lane(_LanePool(set(), idle_cpu=False), True), False) + self.assertIs(self._lane(_LanePool(set(), idle_cpu=False), False), False) + + def test_task_uses_gpu_from_flag_or_device(self): + def task(**model_attrs): + return SimpleNamespace(model=SimpleNamespace(**model_attrs)) + + self.assertTrue(_task_uses_gpu(task(uses_gpu=True))) + self.assertFalse(_task_uses_gpu(task(uses_gpu=False, device="cuda"))) + self.assertTrue(_task_uses_gpu(task(device=SimpleNamespace(type="cuda")))) + self.assertFalse(_task_uses_gpu(task(device=SimpleNamespace(type="cpu")))) + self.assertFalse(_task_uses_gpu(task())) + + if __name__ == "__main__": unittest.main() diff --git a/src/main/python/tests/scuro/test_unimodal_optimizer.py b/src/main/python/tests/scuro/test_unimodal_optimizer.py index f27c721aa25..6e002c1688e 100644 --- a/src/main/python/tests/scuro/test_unimodal_optimizer.py +++ b/src/main/python/tests/scuro/test_unimodal_optimizer.py @@ -47,6 +47,7 @@ from systemds.scuro.representations.bow import BoW from systemds.scuro.representations.bert import Bert from systemds.scuro.representations.word2vec import W2V +from systemds.scuro.representations.window_aggregation import Window from tests.scuro.test_unimodal_representations import ( PHYSIOLOGICAL_REPRESENTATIONS, TIMESERIES_REPRESENTATIONS, @@ -162,6 +163,43 @@ def test_robust_results_ignore_non_finite_scores(self): self.assertEqual(indices, [0, 5]) self.assertTrue(all(np.isfinite(entry.robustness_score) for entry in robust)) + def test_robust_results_penalise_fold_variation(self): + modality = SimpleNamespace(modality_id="modality") + task = SimpleNamespace(model=SimpleNamespace(name="task")) + results = UnimodalResults([modality], [task], k=2) + entries = [ + ResultEntry( + val_score={"accuracy": 0.90}, + val_fold_scores={"accuracy": [0.70, 0.80, 0.90, 1.00, 1.10]}, + dag=SimpleNamespace( + nodes=[SimpleNamespace(operation=UnimodalOptimizer)] + ), + ), + ResultEntry( + val_score={"accuracy": 0.86}, + val_fold_scores={"accuracy": [0.85, 0.86, 0.87, 0.85, 0.87]}, + dag=SimpleNamespace( + nodes=[SimpleNamespace(operation=UnimodalOptimizer)] + ), + ), + ] + results.results[modality.modality_id][task.model.name] = entries + + robust, indices = results.get_k_most_robust_results( + modality, task, "accuracy", one_se_parsimony=False + ) + + self.assertEqual(robust, [entries[1], entries[0]]) + self.assertEqual(indices, [1, 0]) + self.assertAlmostEqual( + entries[0].robustness_score, + 0.90 - np.std(entries[0].val_fold_scores["accuracy"], ddof=1), + ) + self.assertAlmostEqual( + entries[1].robustness_score, + 0.86 - np.std(entries[1].val_fold_scores["accuracy"], ddof=1), + ) + def test_bow_and_tfidf_require_dimensionality_reduction_before_task(self): text_data, text_md = ModalityRandomDataGenerator().create_text_data( self.num_instances, 10 @@ -194,6 +232,84 @@ def test_bow_and_tfidf_require_dimensionality_reduction_before_task(self): task_input = dag.get_node_by_id(task_node.inputs[0]) self.assertIs(task_input.operation, MLPAveraging) + def _window_inputs(self, modality, representations): + """(input operation, input rank) of every window node in the search.""" + with patch.object( + Registry, + "_representations", + registry_for(modality.modality_type, representations), + ): + optimizer = UnimodalOptimizer( + [modality], self.tasks, False, enable_checkpointing=False + ) + dags, _, _ = optimizer._build_execution_dags_for_modality(modality) + stats = {} + for dag in dags: + stats.update(optimizer._node_output_stats(dag)) + inputs = [] + for dag in dags: + for node in dag.nodes: + if node.operation is None or not issubclass(node.operation, Window): + continue + source = dag.get_node_by_id(node.inputs[0]) + inputs.append( + ( + source.operation, + len(stats[source.node_id].output_shape), + ) + ) + return optimizer, inputs + + def test_windows_are_not_applied_to_rank_one_representation_outputs(self): + text_data, text_md = ModalityRandomDataGenerator().create_text_data( + self.num_instances, 10 + ) + text = UnimodalModality( + TestDataLoader( + self.indices, None, ModalityType.TEXT, text_data, str, text_md + ) + ) + audio_data, audio_md = ModalityRandomDataGenerator().create_audio_data( + self.num_instances, 4000 + ) + audio = UnimodalModality( + TestDataLoader( + self.indices, None, ModalityType.AUDIO, audio_data, np.float32, audio_md + ) + ) + # W2V pools to one vector per instance; Concatenation flattens the + # (frames, features) outputs of Spectral and ZeroCrossing to rank 1. + for modality, representations in ( + (text, [W2V]), + (audio, [Spectral, ZeroCrossing]), + ): + with self.subTest(modality=modality.modality_type.name): + optimizer, inputs = self._window_inputs(modality, representations) + representation_inputs = [ + (operation, rank) + for operation, rank in inputs + if operation is not None + ] + self.assertEqual(representation_inputs, []) + self.assertTrue( + any("no time axis" in p["reason"] for p in optimizer.pruned) + ) + + def test_windows_are_kept_for_representations_with_a_time_axis(self): + video_data, video_md = ModalityRandomDataGenerator().create_visual_modality( + self.num_instances, 10, 10, 10 + ) + video = UnimodalModality( + TestDataLoader( + self.indices, None, ModalityType.VIDEO, video_data, np.float32, video_md + ) + ) + _, inputs = self._window_inputs(video, [ResNet]) + self.assertGreater(len(inputs), 0) + for operation, rank in inputs: + self.assertIs(operation, ResNet) + self.assertGreaterEqual(rank, 2) + def test_unimodal_optimizer_for_image_modality(self): image_data, image_md = ModalityRandomDataGenerator().create_visual_modality( self.num_instances, 1, 10, 10 @@ -453,3 +569,71 @@ def optimize_unimodal_representation_for_modality(self, modalities): modalities[0], self.tasks[0], "accuracy" ) assert len(result) == 1 + + +class TestSelectionStrategies(unittest.TestCase): + """The three selection strategies must genuinely disagree. + + They are only worth carrying separately through HPO and fusion if they can + pick different winners out of the same pool, so the fixture is built so + each strategy has a different answer: A wins on raw score, B on robustness + (low fold variation, small DAG), and B again once runtime counts, with A + demoted for being slow. + """ + + @staticmethod + def _entry(accuracy, std, runtime, n_nodes): + dag = SimpleNamespace( + nodes=[SimpleNamespace(operation=object()) for _ in range(n_nodes)] + ) + return SimpleNamespace( + dag=dag, + val_score={"accuracy": accuracy}, + val_fold_scores={"accuracy": [accuracy - std, accuracy + std]}, + val_score_stds={"accuracy": std}, + representation_time=runtime, + task_time=0.0, + ) + + def setUp(self): + self.modality = SimpleNamespace(modality_id="m1") + self.task = SimpleNamespace(model=SimpleNamespace(name="t1")) + self.entries = { + "noisy_slow": self._entry(0.90, 0.10, 100.0, 5), + "stable_fast": self._entry(0.88, 0.01, 10.0, 2), + "weak_fastest": self._entry(0.85, 0.05, 1.0, 1), + } + self.results = UnimodalResults([self.modality], [self.task], k=3) + self.results.results["m1"]["t1"] = list(self.entries.values()) + + def _winner(self, strategy, **kwargs): + view = self.results.with_selection_strategy(strategy, **kwargs) + ranked, _ = view.get_k_best_results( + self.modality, self.task, "accuracy", cache_needed=False + ) + return ranked[0] + + def test_performance_takes_the_validation_argmax(self): + self.assertIs(self._winner("performance"), self.entries["noisy_slow"]) + + def test_robustness_penalises_fold_variation(self): + self.assertIs(self._winner("robustness"), self.entries["stable_fast"]) + + def test_tradeoff_weighs_runtime_against_accuracy(self): + self.assertIs(self._winner("tradeoff"), self.entries["stable_fast"]) + + def test_tradeoff_with_all_weight_on_accuracy_matches_performance(self): + self.assertIs( + self._winner("tradeoff", tradeoff_weights=(1.0, 0.0)), + self.entries["noisy_slow"], + ) + + def test_views_share_the_candidate_pool_and_leave_the_source_alone(self): + view = self.results.with_selection_strategy("robustness") + self.assertIs(view.results, self.results.results) + self.assertIs(view.cache, self.results.cache) + self.assertEqual(self.results.selection_strategy, "performance") + + def test_unknown_strategy_is_rejected(self): + with self.assertRaises(ValueError): + self.results.with_selection_strategy("cheapest") From 32ae92cc775cda2e41d09ef87568e66008de8e10 Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Mon, 12 Oct 2026 11:29:10 +0200 Subject: [PATCH 16/17] fix python version error --- src/main/python/systemds/scuro/dataloader/json_loader.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/main/python/systemds/scuro/dataloader/json_loader.py b/src/main/python/systemds/scuro/dataloader/json_loader.py index d2802894ecb..c8361643046 100644 --- a/src/main/python/systemds/scuro/dataloader/json_loader.py +++ b/src/main/python/systemds/scuro/dataloader/json_loader.py @@ -81,7 +81,7 @@ def extract(self, file: str, index: Optional[Union[str, List[str]]] = None): self.data.append(text) self.metadata.append( - self.modality_type.create_metadata(len(text), text) | record + {**self.modality_type.create_metadata(len(text), text), **record} ) def get_stats(self, source_path: str): From 42426bfb9e4fd9eec5061623ec97a208a4074b2b Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Mon, 12 Oct 2026 13:12:18 +0200 Subject: [PATCH 17/17] replace library --- .../tests/scuro/test_algebra_canonical.py | 52 ++--- .../tests/scuro/test_algebra_explicit_prep.py | 206 ++++++++++-------- .../python/tests/scuro/test_algebra_laws.py | 64 +++--- .../tests/scuro/test_algebra_signatures.py | 72 +++--- .../python/tests/scuro/test_algebra_stats.py | 37 +++- .../scuro/test_algebra_subtree_sharing.py | 116 +++++----- .../tests/scuro/test_algebra_typecheck.py | 130 ++++++----- 7 files changed, 376 insertions(+), 301 deletions(-) diff --git a/src/main/python/tests/scuro/test_algebra_canonical.py b/src/main/python/tests/scuro/test_algebra_canonical.py index f1db174cb4a..9801d4817a3 100644 --- a/src/main/python/tests/scuro/test_algebra_canonical.py +++ b/src/main/python/tests/scuro/test_algebra_canonical.py @@ -20,8 +20,7 @@ # ------------------------------------------------------------- import json - -import pytest +import unittest from systemds.scuro.drsearch.algebra import ALL_LAWS, canonical_id from systemds.scuro.drsearch.representation_dag import RepresentationDAGBuilder @@ -242,28 +241,31 @@ def make(e): return builder.build(make(expr)) -@pytest.mark.parametrize( - "left,right,laws,invariant,equal", [c[1:] for c in CASES], ids=[c[0] for c in CASES] -) -def test_canonical_pairs(left, right, laws, invariant, equal): - laws = frozenset(laws) - left, right = build(left), build(right) - ids = [canonical_id(dag, laws, invariant) for dag in (left, right)] - assert (ids[0] == ids[1]) == equal - for dag, dag_id in zip((left, right), ids): - spec = json.loads(json.dumps(dag.to_spec())) - assert canonical_id(spec, laws, invariant) == dag_id - - -def test_dag_method_defaults(): - dag = build(C("a", S("c", "b"))) - assert dag.canonical_id() == canonical_id(dag, ALL_LAWS, False) - assert dag.canonical_id(laws=NONE) != dag.canonical_id() +class TestCanonicalId(unittest.TestCase): + def test_canonical_pairs(self): + for name, left, right, laws, invariant, equal in CASES: + with self.subTest(name): + laws = frozenset(laws) + left, right = build(left), build(right) + ids = [canonical_id(dag, laws, invariant) for dag in (left, right)] + self.assertEqual(ids[0] == ids[1], equal) + for dag, dag_id in zip((left, right), ids): + spec = json.loads(json.dumps(dag.to_spec())) + self.assertEqual(canonical_id(spec, laws, invariant), dag_id) + + def test_dag_method_defaults(self): + dag = build(C("a", S("c", "b"))) + self.assertEqual(dag.canonical_id(), canonical_id(dag, ALL_LAWS, False)) + self.assertNotEqual(dag.canonical_id(laws=NONE), dag.canonical_id()) + + def test_sub_plan_signature(self): + dag = build(C("a", S("b", "c"))) + sub = dag.get_node_by_id(dag.get_node_by_id(dag.root_node_id).inputs[1]) + self.assertEqual( + canonical_id(dag, ALL_LAWS, False, sub.node_id), + canonical_id(build(S("c", "b")), ALL_LAWS, False), + ) -def test_sub_plan_signature(): - dag = build(C("a", S("b", "c"))) - sub = dag.get_node_by_id(dag.get_node_by_id(dag.root_node_id).inputs[1]) - assert canonical_id(dag, ALL_LAWS, False, sub.node_id) == canonical_id( - build(S("c", "b")), ALL_LAWS, False - ) +if __name__ == "__main__": + unittest.main() diff --git a/src/main/python/tests/scuro/test_algebra_explicit_prep.py b/src/main/python/tests/scuro/test_algebra_explicit_prep.py index 3f57605a813..0f6ff056567 100644 --- a/src/main/python/tests/scuro/test_algebra_explicit_prep.py +++ b/src/main/python/tests/scuro/test_algebra_explicit_prep.py @@ -19,12 +19,13 @@ # # ------------------------------------------------------------- +import itertools import tempfile +import unittest from collections import Counter from types import SimpleNamespace import numpy as np -import pytest from systemds.scuro.drsearch.algebra import ( ALL_LAWS, @@ -125,94 +126,117 @@ def _ga_dag(algebra, reps, genome): return optimizer._genome_to_dag(genome) -@pytest.mark.parametrize("algebra", [AGG, PAD, BOTH]) -@pytest.mark.parametrize("op", FUSIONS) -@pytest.mark.parametrize("leaves", [(0, 1), (1, 2)]) -def test_ga_genome_to_dag(algebra, op, leaves): - reps = _reps() - genome = DagGenome([(m, 0) for m in leaves], (0, 1), {"": op}) - _assert_equivalent( - _ga_dag(AlgebraConfig(), reps, genome), - _ga_dag(algebra, reps, genome), - reps, - _expected(algebra, op, leaves), - ) - - -@pytest.mark.parametrize("algebra", [AGG, BOTH]) -def test_exhaustive_enumerator(algebra): - reps = _reps() - runs = [] - for explicit in (AlgebraConfig(), algebra): - with tempfile.TemporaryDirectory() as tmp: - optimizer = MultimodalOptimizer( - [_FakeModality(m) for m in reps], - _FakeUnimodalResults(reps), - [_make_task()], - debug=False, - checkpoint_dir=tmp, - algebra=AlgebraConfig( - exhaustive_dedup_permutations=True, - explicit_fusion_prep=explicit.explicit_fusion_prep, - explicit_padding=explicit.explicit_padding, - ), - ) - optimizer.fusion_operators = FUSIONS - runs.append( - list(optimizer._generate_fusion_dags([0, 1], {0: 0, 1: 0}, "task0")) +class TestExplicitPrep(unittest.TestCase): + def test_ga_genome_to_dag(self): + cases = itertools.product([AGG, PAD, BOTH], FUSIONS, [(0, 1), (1, 2)]) + for algebra, op, leaves in cases: + with self.subTest(algebra=algebra, op=op.__name__, leaves=leaves): + reps = _reps() + genome = DagGenome([(m, 0) for m in leaves], (0, 1), {"": op}) + _assert_equivalent( + _ga_dag(AlgebraConfig(), reps, genome), + _ga_dag(algebra, reps, genome), + reps, + _expected(algebra, op, leaves), + ) + + def test_exhaustive_enumerator(self): + for algebra in (AGG, BOTH): + with self.subTest(algebra=algebra): + reps = _reps() + runs = [] + for explicit in (AlgebraConfig(), algebra): + with tempfile.TemporaryDirectory() as tmp: + optimizer = MultimodalOptimizer( + [_FakeModality(m) for m in reps], + _FakeUnimodalResults(reps), + [_make_task()], + debug=False, + checkpoint_dir=tmp, + algebra=AlgebraConfig( + exhaustive_dedup_permutations=True, + explicit_fusion_prep=explicit.explicit_fusion_prep, + explicit_padding=explicit.explicit_padding, + ), + ) + optimizer.fusion_operators = FUSIONS + runs.append( + list( + optimizer._generate_fusion_dags( + [0, 1], {0: 0, 1: 0}, "task0" + ) + ) + ) + self.assertEqual(len(runs[0]), 2 * len(FUSIONS)) + self.assertEqual(len(runs[1]), 2 * len(FUSIONS)) + for legacy, explicit, op in zip(*runs, FUSIONS * 2): + _assert_equivalent( + legacy, explicit, reps, _expected(algebra, op, (0, 1)) + ) + self.assertEqual( + optimizer.algebra_stats["explicit_prep_plans"], 2 * len(FUSIONS) + ) + + def test_counts_and_typechecks_padded_plan(self): + reps = _reps() + dag = _ga_dag(BOTH, reps, DagGenome([(0, 0), (1, 0)], (0, 1), {"": Sum})) + stats = Counter() + count_explicit_prep(stats, dag) + self.assertEqual( + stats, + { + "explicit_prep_plans": 1, + "explicit_padding_plans": 1, + "explicit_aggregations": 1, + "explicit_paddings": 1, + "explicit_padding_ratio_sum": 8 / 3, + "explicit_padding_ratio_max": 8 / 3, + }, ) - assert len(runs[0]) == len(runs[1]) == 2 * len(FUSIONS) - for legacy, explicit, op in zip(*runs, FUSIONS * 2): - _assert_equivalent(legacy, explicit, reps, _expected(algebra, op, (0, 1))) - assert optimizer.algebra_stats["explicit_prep_plans"] == 2 * len(FUSIONS) - - -def test_counts_and_typechecks_padded_plan(): - reps = _reps() - dag = _ga_dag(BOTH, reps, DagGenome([(0, 0), (1, 0)], (0, 1), {"": Sum})) - stats = Counter() - count_explicit_prep(stats, dag) - assert stats == { - "explicit_prep_plans": 1, - "explicit_padding_plans": 1, - "explicit_aggregations": 1, - "explicit_paddings": 1, - "explicit_padding_ratio_sum": 8 / 3, - "explicit_padding_ratio_max": 8 / 3, - } - leaves = {(m, 0): r[0] for m, r in reps.items()} - types = {k: r.modality_type for k, r in leaves.items()} - assert typecheck(dag, types, {k: sample_stats(r) for k, r in leaves.items()}) == [] - - -def test_zero_pad_is_not_registered(): - registry = Registry() - lists = [ - registry._representations, - registry._context_operators, - registry._dimensionality_reduction_operators, - ] - assert all(ZeroPad not in ops for lst in lists for ops in lst.values()) - - -def test_unbound_genome_to_dag_and_bare_checker(): - reps = _reps() - genome = DagGenome([(0, 0), (1, 0)], (0, 1), {"": Sum}) - assert _prep_nodes(MultimodalDeapOptimizer._genome_to_dag(None, genome)) == [] - checker = object.__new__(MultimodalDeapOptimizer) - checker.max_feature_dim = None - checker.k_best_representations = {"task0": reps} - assert checker._oversized_reason("task0", genome) is None - - -@pytest.mark.parametrize( - "params, wrapped", [({}, True), ({"preserve_leading_axis": True}, False)] -) -def test_cse_builder(params, wrapped): - leaf = SimpleNamespace(modality_id="t", stats=RepresentationStats(5, (4, 8))) - builder = CSEAwareDAGBuilder(BOTH, leaves=[leaf]) - leaf_id = builder.create_leaf_node("t") - root = builder.create_operation_node(Concatenation, [leaf_id, leaf_id], params) - builder.build(root) - assert (builder.get_node(root).inputs != [leaf_id, leaf_id]) == wrapped - assert builder.algebra_stats["explicit_aggregations"] == 2 * wrapped + leaves = {(m, 0): r[0] for m, r in reps.items()} + types = {k: r.modality_type for k, r in leaves.items()} + leaf_stats = {k: sample_stats(r) for k, r in leaves.items()} + self.assertEqual(typecheck(dag, types, leaf_stats), []) + + def test_zero_pad_is_not_registered(self): + registry = Registry() + lists = [ + registry._representations, + registry._context_operators, + registry._dimensionality_reduction_operators, + ] + for ops in (ops for lst in lists for ops in lst.values()): + self.assertNotIn(ZeroPad, ops) + + def test_unbound_genome_to_dag_and_bare_checker(self): + reps = _reps() + genome = DagGenome([(0, 0), (1, 0)], (0, 1), {"": Sum}) + dag = MultimodalDeapOptimizer._genome_to_dag(None, genome) + self.assertEqual(_prep_nodes(dag), []) + checker = object.__new__(MultimodalDeapOptimizer) + checker.max_feature_dim = None + checker.k_best_representations = {"task0": reps} + self.assertIsNone(checker._oversized_reason("task0", genome)) + + def test_cse_builder(self): + for params, wrapped in (({}, True), ({"preserve_leading_axis": True}, False)): + with self.subTest(params=params): + leaf = SimpleNamespace( + modality_id="t", stats=RepresentationStats(5, (4, 8)) + ) + builder = CSEAwareDAGBuilder(BOTH, leaves=[leaf]) + leaf_id = builder.create_leaf_node("t") + root = builder.create_operation_node( + Concatenation, [leaf_id, leaf_id], params + ) + builder.build(root) + self.assertEqual( + builder.get_node(root).inputs != [leaf_id, leaf_id], wrapped + ) + self.assertEqual( + builder.algebra_stats["explicit_aggregations"], 2 * wrapped + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/main/python/tests/scuro/test_algebra_laws.py b/src/main/python/tests/scuro/test_algebra_laws.py index 9ef8f7d2c61..c0d55ddc6b6 100644 --- a/src/main/python/tests/scuro/test_algebra_laws.py +++ b/src/main/python/tests/scuro/test_algebra_laws.py @@ -19,8 +19,10 @@ # # ------------------------------------------------------------- +import itertools +import unittest + import numpy as np -import pytest from systemds.scuro.drsearch.operator_registry import Registry from systemds.scuro.modality.modality import Modality @@ -107,30 +109,36 @@ def _concat_perm(f, a, b, c, exact): } -def test_every_declared_law_is_checked(): - assert {law for _, _, law in DECLARED} <= LAW_CHECKS.keys() - - -@pytest.mark.parametrize("rank", [1, 2]) -@pytest.mark.parametrize("dtypes", DTYPES.values(), ids=DTYPES.keys()) -@pytest.mark.parametrize("widths", WIDTHS.values(), ids=WIDTHS.keys()) -@pytest.mark.parametrize( - "op,params,law", - DECLARED, - ids=[f"{op.__name__}{params or ''}-{law}" for op, params, law in DECLARED], -) -def test_declared_law_holds(op, params, law, widths, dtypes, rank): - LAW_CHECKS[law]( - lambda *xs: _fuse(op, params, *xs), - *_operands(widths, dtypes, rank), - exact=op is Concatenation, - ) - - -def test_window_axis_concatenation_is_order_sensitive(): - params = {"preserve_leading_axis": True} - a, b, _ = _operands(WIDTHS["equal"], DTYPES["float32"], rank=2) - assert not np.array_equal( - _fuse(Concatenation, params, a, b), _fuse(Concatenation, params, b, a) - ) - assert not {"comm", "concat_perm"} & Concatenation.algebraic_laws(params) +class TestAlgebraLaws(unittest.TestCase): + def test_every_declared_law_is_checked(self): + self.assertLessEqual({law for _, _, law in DECLARED}, set(LAW_CHECKS)) + + def test_declared_law_holds(self): + cases = itertools.product(DECLARED, WIDTHS, DTYPES, (1, 2)) + for (op, params, law), widths, dtypes, rank in cases: + with self.subTest( + op=f"{op.__name__}{params or ''}", + law=law, + widths=widths, + dtypes=dtypes, + rank=rank, + ): + LAW_CHECKS[law]( + lambda *xs: _fuse(op, params, *xs), + *_operands(WIDTHS[widths], DTYPES[dtypes], rank), + exact=op is Concatenation, + ) + + def test_window_axis_concatenation_is_order_sensitive(self): + params = {"preserve_leading_axis": True} + a, b, _ = _operands(WIDTHS["equal"], DTYPES["float32"], rank=2) + self.assertFalse( + np.array_equal( + _fuse(Concatenation, params, a, b), _fuse(Concatenation, params, b, a) + ) + ) + self.assertFalse({"comm", "concat_perm"} & Concatenation.algebraic_laws(params)) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/main/python/tests/scuro/test_algebra_signatures.py b/src/main/python/tests/scuro/test_algebra_signatures.py index 9cd10e02107..d18097861a1 100644 --- a/src/main/python/tests/scuro/test_algebra_signatures.py +++ b/src/main/python/tests/scuro/test_algebra_signatures.py @@ -19,7 +19,7 @@ # # ------------------------------------------------------------- -import pytest +import unittest from systemds.scuro.drsearch.algebra import ANY_TYPE, operator_signature from systemds.scuro.drsearch.operator_registry import Registry, load_builtin_operators @@ -48,39 +48,43 @@ key=lambda t: (t[0].__name__, t[1], t[2].name), ) + [(op, "fuse", None) for op in Registry().get_fusion_operators()] +RANK_CASES = [ + (Sum, {}, ("fuse", "n", "vector", "vector")), + ( + Concatenation, + {"preserve_leading_axis": False}, + ("fuse", "n", "vector", "vector"), + ), + (Concatenation, {"preserve_leading_axis": True}, ("fuse", "n", "seq", "seq")), + (LSTM, {}, ("fuse", "n", "seq", "vector")), + (AggregatedRepresentation, {}, ("aggregate", 1, "seq", "vector")), + (WindowAggregation, {}, ("context", 1, "seq", "windows")), + (MLPAveraging, {}, ("reduce", 1, "vector", "vector")), + (None, {}, ("task", 1, "any", "any")), +] -@pytest.mark.parametrize( - "op, kind, modality", - REGISTERED, - ids=[ - f"{op.__name__}-{kind}-{getattr(m, 'name', '')}" for op, kind, m in REGISTERED - ], -) -def test_every_registered_operator_has_a_signature(op, kind, modality): - sig = operator_signature(op) - assert sig.kind == kind - assert modality is None or modality in sig.input_types - if kind in ("encode", "fuse"): - assert isinstance(sig.output_type, ModalityType) +class TestOperatorSignatures(unittest.TestCase): + def test_every_registered_operator_has_a_signature(self): + for op, kind, modality in REGISTERED: + with self.subTest( + op=op.__name__, kind=kind, modality=getattr(modality, "name", None) + ): + sig = operator_signature(op) + self.assertEqual(sig.kind, kind) + if modality is not None: + self.assertIn(modality, sig.input_types) + if kind in ("encode", "fuse"): + self.assertIsInstance(sig.output_type, ModalityType) -@pytest.mark.parametrize( - "op, params, expected", - [ - (Sum, {}, ("fuse", "n", "vector", "vector")), - ( - Concatenation, - {"preserve_leading_axis": False}, - ("fuse", "n", "vector", "vector"), - ), - (Concatenation, {"preserve_leading_axis": True}, ("fuse", "n", "seq", "seq")), - (LSTM, {}, ("fuse", "n", "seq", "vector")), - (AggregatedRepresentation, {}, ("aggregate", 1, "seq", "vector")), - (WindowAggregation, {}, ("context", 1, "seq", "windows")), - (MLPAveraging, {}, ("reduce", 1, "vector", "vector")), - (None, {}, ("task", 1, "any", "any")), - ], -) -def test_rank_rules(op, params, expected): - sig = operator_signature(op, params) - assert (sig.kind, sig.arity, sig.input_rank, sig.output_rank) == expected + def test_rank_rules(self): + for op, params, expected in RANK_CASES: + with self.subTest(op=getattr(op, "__name__", None), params=params): + sig = operator_signature(op, params) + self.assertEqual( + (sig.kind, sig.arity, sig.input_rank, sig.output_rank), expected + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/main/python/tests/scuro/test_algebra_stats.py b/src/main/python/tests/scuro/test_algebra_stats.py index 5c18d50756f..619c686e37b 100644 --- a/src/main/python/tests/scuro/test_algebra_stats.py +++ b/src/main/python/tests/scuro/test_algebra_stats.py @@ -19,8 +19,10 @@ # # ------------------------------------------------------------- +import itertools +import unittest + import numpy as np -import pytest from systemds.scuro.drsearch.operator_registry import Registry from systemds.scuro.representations.representation import RepresentationStats @@ -39,14 +41,25 @@ } -@pytest.mark.parametrize("case", OPERANDS) -@pytest.mark.parametrize("op", OPERATORS, ids=lambda op: op.__name__) -def test_predicted_stats_match_execution(op, case): - rng = np.random.default_rng(0) - operands = [ - rng.standard_normal(shape).astype(np.float32) for shape in OPERANDS[case] - ] - stats = [RepresentationStats(len(x), x.shape[1:]) for x in operands] - fused = _fuse(op, {}, *operands) - assert tuple(op().get_output_stats(stats).output_shape) == fused.shape[1:] - assert op().estimate_peak_memory_bytes(stats)["cpu_peak_bytes"] > fused.nbytes +class TestPredictedStats(unittest.TestCase): + def test_predicted_stats_match_execution(self): + for op, case in itertools.product(OPERATORS, OPERANDS): + with self.subTest(op=op.__name__, case=case): + rng = np.random.default_rng(0) + operands = [ + rng.standard_normal(shape).astype(np.float32) + for shape in OPERANDS[case] + ] + stats = [RepresentationStats(len(x), x.shape[1:]) for x in operands] + fused = _fuse(op, {}, *operands) + self.assertEqual( + tuple(op().get_output_stats(stats).output_shape), fused.shape[1:] + ) + self.assertGreater( + op().estimate_peak_memory_bytes(stats)["cpu_peak_bytes"], + fused.nbytes, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/main/python/tests/scuro/test_algebra_subtree_sharing.py b/src/main/python/tests/scuro/test_algebra_subtree_sharing.py index 1051c111521..fbaa7f7721c 100644 --- a/src/main/python/tests/scuro/test_algebra_subtree_sharing.py +++ b/src/main/python/tests/scuro/test_algebra_subtree_sharing.py @@ -18,10 +18,10 @@ # under the License. # # ------------------------------------------------------------- +import unittest from unittest.mock import patch import numpy as np -import pytest from systemds.scuro.drsearch import multimodal_ga_optimizer as ga from systemds.scuro.drsearch.algebra import AlgebraConfig @@ -47,16 +47,6 @@ def fit(self, X, y, X_test, y_test): return super().fit(X, y, X_test, y_test) -@pytest.fixture -def setup(monkeypatch): - monkeypatch.setattr(ga, "_SUBTREE_CACHE", TrialResultCache(2**30, 0.0)) - task = TestTask("share_task", "share_model", 40) - task.model = _RecordingSVM("share_model") - modalities = [_make_real_representation(i, 40, 4) for i in range(3)] - _RecordingSVM.inputs = [] - return task, modalities - - def _dag(leaves, root_op=None): b = RepresentationDAGBuilder() ids = [b.create_leaf_node(i, 0) for i in leaves] @@ -72,44 +62,66 @@ def _evaluate(task, modalities, dag, share): return payload, _RecordingSVM.inputs -def test_fitted_subtree_is_reused_bit_identically(setup): - task, modalities = setup - _evaluate(task, modalities, _dag([0, 1]), True) - shared, shared_inputs = _evaluate( - task, modalities, _dag([0, 1, 2], Concatenation), True - ) - _, fresh_inputs = _evaluate(task, modalities, _dag([0, 1, 2], Concatenation), False) - - assert shared["subtree_cache"]["hits"] == len(task.evaluation_splits) - assert len(shared_inputs) == len(fresh_inputs) == len(task.evaluation_splits) - assert all(np.array_equal(a, b) for a, b in zip(shared_inputs, fresh_inputs)) - - -def test_operand_order_is_part_of_the_key(setup): - task, modalities = setup - _evaluate(task, modalities, _dag([0, 1]), True) - swapped, _ = _evaluate(task, modalities, _dag([1, 0]), True) - assert swapped["subtree_cache"]["hits"] == 0 - - -@pytest.mark.parametrize("share", [False, True]) -def test_switch_controls_call_and_stats(share): - optimizer, task = _build_real_optimizer( - fusion_ops=[LSTM], algebra=AlgebraConfig(share_fused_subtrees=share) - ) - genome = optimizer._random_genome(task.model.name) - payload = { - "train_score": {}, - "val_score": {"accuracy": 0.5}, - "test_score": {}, - "runtime": 1.0, - "task_time": 0.5, - "representation_time": 0.5, - "subtree_cache": {"hits": 3, "misses": 2, "saved_seconds": 4.0, "bytes": 10}, - } - with patch.object( - ga, "_evaluate_genome_body", return_value=((0.5,), payload) - ) as body: - optimizer._evaluate_genome(genome, task) - assert len(body.call_args.args) == 4 + share - assert optimizer.algebra_stats["subtree_cache_hits"] == 3 +class TestSubtreeSharing(unittest.TestCase): + def setUp(self): + cache = patch.object(ga, "_SUBTREE_CACHE", TrialResultCache(2**30, 0.0)) + cache.start() + self.addCleanup(cache.stop) + self.task = TestTask("share_task", "share_model", 40) + self.task.model = _RecordingSVM("share_model") + self.modalities = [_make_real_representation(i, 40, 4) for i in range(3)] + _RecordingSVM.inputs = [] + + def test_fitted_subtree_is_reused_bit_identically(self): + task, modalities = self.task, self.modalities + splits = len(task.evaluation_splits) + _evaluate(task, modalities, _dag([0, 1]), True) + shared, shared_inputs = _evaluate( + task, modalities, _dag([0, 1, 2], Concatenation), True + ) + _, fresh_inputs = _evaluate( + task, modalities, _dag([0, 1, 2], Concatenation), False + ) + + self.assertEqual(shared["subtree_cache"]["hits"], splits) + self.assertEqual(len(shared_inputs), splits) + self.assertEqual(len(fresh_inputs), splits) + for a, b in zip(shared_inputs, fresh_inputs): + np.testing.assert_array_equal(a, b) + + def test_operand_order_is_part_of_the_key(self): + _evaluate(self.task, self.modalities, _dag([0, 1]), True) + swapped, _ = _evaluate(self.task, self.modalities, _dag([1, 0]), True) + self.assertEqual(swapped["subtree_cache"]["hits"], 0) + + def test_switch_controls_call_and_stats(self): + payload = { + "train_score": {}, + "val_score": {"accuracy": 0.5}, + "test_score": {}, + "runtime": 1.0, + "task_time": 0.5, + "representation_time": 0.5, + "subtree_cache": { + "hits": 3, + "misses": 2, + "saved_seconds": 4.0, + "bytes": 10, + }, + } + for share in (False, True): + with self.subTest(share=share): + optimizer, task = _build_real_optimizer( + fusion_ops=[LSTM], algebra=AlgebraConfig(share_fused_subtrees=share) + ) + genome = optimizer._random_genome(task.model.name) + with patch.object( + ga, "_evaluate_genome_body", return_value=((0.5,), payload) + ) as body: + optimizer._evaluate_genome(genome, task) + self.assertEqual(len(body.call_args.args), 4 + share) + self.assertEqual(optimizer.algebra_stats["subtree_cache_hits"], 3) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/main/python/tests/scuro/test_algebra_typecheck.py b/src/main/python/tests/scuro/test_algebra_typecheck.py index b59f0e25e10..4e34a3be095 100644 --- a/src/main/python/tests/scuro/test_algebra_typecheck.py +++ b/src/main/python/tests/scuro/test_algebra_typecheck.py @@ -19,10 +19,10 @@ # # ------------------------------------------------------------- +import unittest from types import SimpleNamespace import numpy as np -import pytest from systemds.scuro.drsearch.algebra import AlgebraConfig, typecheck from systemds.scuro.drsearch.multimodal_ga_optimizer import ( @@ -55,67 +55,79 @@ def _check(op, params, leaves, with_stats=True): return sorted((v.severity, v.kind) for v in violations) -@pytest.mark.parametrize( - "op, params, leaves, expected", - [ - (BoW, {}, [(TEXT, None)], []), - (MFCC, {}, [(TEXT, None)], [("error", "input_type")]), - (Sum, {}, [(EMB, VEC)], [("error", "arity")]), - (BoW, {}, [(TEXT, None), (TEXT, None)], [("error", "arity")]), - (Sum, {}, [(EMB, VEC), (EMB, VEC)], []), - ( - Sum, - {}, - [(EMB, SEQ), (EMB, VEC)], - [("implicit", "hidden_aggregation"), ("implicit", "zero_padding")], - ), - ( - Concatenation, - {}, - [(EMB, SEQ), (EMB, VEC)], - [("implicit", "hidden_aggregation")], - ), - (Concatenation, {"preserve_leading_axis": True}, [(EMB, SEQ), (EMB, SEQ)], []), - ], -) -def test_violations(op, params, leaves, expected): - assert _check(op, params, leaves) == expected +VIOLATION_CASES = [ + (BoW, {}, [(TEXT, None)], []), + (MFCC, {}, [(TEXT, None)], [("error", "input_type")]), + (Sum, {}, [(EMB, VEC)], [("error", "arity")]), + (BoW, {}, [(TEXT, None), (TEXT, None)], [("error", "arity")]), + (Sum, {}, [(EMB, VEC), (EMB, VEC)], []), + ( + Sum, + {}, + [(EMB, SEQ), (EMB, VEC)], + [("implicit", "hidden_aggregation"), ("implicit", "zero_padding")], + ), + ( + Concatenation, + {}, + [(EMB, SEQ), (EMB, VEC)], + [("implicit", "hidden_aggregation")], + ), + (Concatenation, {"preserve_leading_axis": True}, [(EMB, SEQ), (EMB, SEQ)], []), +] -def test_shapes_unknown_without_leaf_stats(): - assert _check(Sum, {}, [(EMB, SEQ), (EMB, VEC)], with_stats=False) == [] +class TestTypecheck(unittest.TestCase): + def test_violations(self): + for i, (op, params, leaves, expected) in enumerate(VIOLATION_CASES): + with self.subTest(case=i, op=op.__name__): + self.assertEqual(_check(op, params, leaves), expected) + def test_shapes_unknown_without_leaf_stats(self): + violations = _check(Sum, {}, [(EMB, SEQ), (EMB, VEC)], with_stats=False) + self.assertEqual(violations, []) -@pytest.mark.parametrize( - "mode, raises, counted", - [("off", False, 0), ("warn", False, 1), ("strict", True, 1)], -) -def test_builder_modes(mode, raises, counted): - modality = SimpleNamespace(modality_id="t", modality_type=TEXT, stats=None) - builder = CSEAwareDAGBuilder(AlgebraConfig(typecheck=mode), leaves=[modality]) - root = builder.create_operation_node(MFCC, [builder.create_leaf_node("t")], {}) - if raises: - with pytest.raises(ValueError, match="input_type"): - builder.build(root) - else: - builder.build(root) - assert builder.algebra_stats["typecheck_input_type"] == counted + def test_builder_modes(self): + for mode, raises, counted in [ + ("off", False, 0), + ("warn", False, 1), + ("strict", True, 1), + ]: + with self.subTest(mode=mode): + modality = SimpleNamespace( + modality_id="t", modality_type=TEXT, stats=None + ) + builder = CSEAwareDAGBuilder( + AlgebraConfig(typecheck=mode), leaves=[modality] + ) + root = builder.create_operation_node( + MFCC, [builder.create_leaf_node("t")], {} + ) + if raises: + with self.assertRaisesRegex(ValueError, "input_type"): + builder.build(root) + else: + builder.build(root) + self.assertEqual(builder.algebra_stats["typecheck_input_type"], counted) + + def test_ga_counts_hidden_operators(self): + reps = { + "m0": [SimpleNamespace(modality_type=EMB, data=np.zeros((5, 4, 8)))], + "m1": [SimpleNamespace(modality_type=EMB, data=np.zeros((5, 3)))], + } + optimizer = MultimodalDeapOptimizer( + [SimpleNamespace(modality_id=m) for m in reps], + _FakeUnimodalResults(reps), + [_make_task()], + debug=False, + algebra=AlgebraConfig(typecheck="strict"), + ) + genome = DagGenome([("m0", 0), ("m1", 0)], (0, 1), {"": Sum}) + self.assertIsNone(optimizer._oversized_reason("task0", genome)) + stats = optimizer.algebra_stats + self.assertEqual(stats["typecheck_hidden_aggregation"], 1) + self.assertEqual(stats["typecheck_zero_padding:Sum"], 1) -def test_ga_counts_hidden_operators(): - reps = { - "m0": [SimpleNamespace(modality_type=EMB, data=np.zeros((5, 4, 8)))], - "m1": [SimpleNamespace(modality_type=EMB, data=np.zeros((5, 3)))], - } - optimizer = MultimodalDeapOptimizer( - [SimpleNamespace(modality_id=m) for m in reps], - _FakeUnimodalResults(reps), - [_make_task()], - debug=False, - algebra=AlgebraConfig(typecheck="strict"), - ) - genome = DagGenome([("m0", 0), ("m1", 0)], (0, 1), {"": Sum}) - assert optimizer._oversized_reason("task0", genome) is None - stats = optimizer.algebra_stats - assert stats["typecheck_hidden_aggregation"] == 1 - assert stats["typecheck_zero_padding:Sum"] == 1 +if __name__ == "__main__": + unittest.main()