[SYSTEMDS-3978] Add Algebra Rules to Scuro
This patch adds rules for the DARE-Algebra to Scuro that enable better node sharing, and ordering.
Assisted-by: AI
diff --git a/src/main/python/systemds/scuro/dataloader/json_loader.py b/src/main/python/systemds/scuro/dataloader/json_loader.py
index d280289..c836164 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 @@
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):
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 0000000..880e7af
--- /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 bfb62bc..7805ac6 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 @@
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 @@
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 @@
budget=n_calls,
initial_config=None,
rep_name=rep_name,
+ tied_params=tied_params,
)
all_results.append(baseline)
@@ -924,6 +989,7 @@
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 @@
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 @@
)
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 b7472a3..1f78ff6 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 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 @@
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 @@
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 @@
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 79a04ca..0c18233 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 @@
_TIMING_OBJECTIVES = {"runtime", "task_time", "representation_time"}
+_MAX_ERROR_SAMPLES = 5
ObjectiveSpec = Tuple[str, str] # (name, "max" | "min")
@@ -187,6 +202,14 @@
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 @@
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 @@
"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 _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 @@
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 @@
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 @@
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 @@
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 @@
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 @@
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 @@
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 @@
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 @@
)
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 @@
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 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 @@
]
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 @@
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 @@
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 @@
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 54cddfd..22476bc 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 @@
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 @@
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 @@
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_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 @@
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 @@
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 @@
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 a498c27..4be1b66 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 @@
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 @@
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 @@
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 @@
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 @@
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 @@
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 @@
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 @@
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 @@
"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 @@
"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 78ef5b9..b69c50f 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 @@
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 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 @@
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_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 d84d092..efc4692 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 @@
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 @@
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 @@
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 @@
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_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 @@
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 97de3da..05b70d0 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 @@
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 @@
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 @@
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 ced501b..8b717e5 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 @@
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 @@
)
# 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 @@
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 62cc03e..87cc3e9 100644
--- a/src/main/python/systemds/scuro/modality/transformed.py
+++ b/src/main/python/systemds/scuro/modality/transformed.py
@@ -185,10 +185,13 @@
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 fcd4150..965058e 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 366bb9f..2c54913 100644
--- a/src/main/python/systemds/scuro/representations/aggregate.py
+++ b/src/main/python/systemds/scuro/representations/aggregate.py
@@ -51,18 +51,22 @@
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 06c435d..ade8b03 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 a64354b..55f841e 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 cee3038..f599aba 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 5cbe128..95bdec2 100644
--- a/src/main/python/systemds/scuro/representations/fusion.py
+++ b/src/main/python/systemds/scuro/representations/fusion.py
@@ -55,22 +55,28 @@
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 9513760..8e9cf0f 100644
--- a/src/main/python/systemds/scuro/representations/glove.py
+++ b/src/main/python/systemds/scuro/representations/glove.py
@@ -49,8 +49,10 @@
@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 0697a9f..a88d7d6 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 97b318e..46e0842 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 @@
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 2e0df4c..f221a0c 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 @@
"""
super().__init__("RowMax")
self.needs_alignment = True
- self.associative = True
- self.commutative = True
def execute(
self,
@@ -64,3 +64,14 @@
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 60599e8..b2d82cc 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 @@
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 @@
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 @@
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 @@
)
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 @@
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 @@
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 @@
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 @@
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 @@
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)
+ keep = py_cpu_nms(dets, detector.nms_threshold)
+ dets = np.concatenate((dets[keep], points[keep]), axis=1)
- 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
+ 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 58cac2e..0b9ce2e 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 @@
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 5bf4a6d..2172e6a 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 @@
# "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 6f82937..214b56a 100644
--- a/src/main/python/systemds/scuro/representations/vgg.py
+++ b/src/main/python/systemds/scuro/representations/vgg.py
@@ -56,6 +56,7 @@
@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 @@
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 e8189e9..60b5b46 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 @@
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 0000000..505545e
--- /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 0000000..9801d48
--- /dev/null
+++ b/src/main/python/tests/scuro/test_algebra_canonical.py
@@ -0,0 +1,271 @@
+# -------------------------------------------------------------
+#
+# 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 unittest
+
+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))
+
+
+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),
+ )
+
+
+if __name__ == "__main__":
+ unittest.main()
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 0000000..df86a55
--- /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 0000000..249720e
--- /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 0000000..0f6ff05
--- /dev/null
+++ b/src/main/python/tests/scuro/test_algebra_explicit_prep.py
@@ -0,0 +1,242 @@
+# -------------------------------------------------------------
+#
+# 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 itertools
+import tempfile
+import unittest
+from collections import Counter
+from types import SimpleNamespace
+
+import numpy as np
+
+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)
+
+
+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,
+ },
+ )
+ 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_ga_dedup.py b/src/main/python/tests/scuro/test_algebra_ga_dedup.py
new file mode 100644
index 0000000..8bc457d
--- /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 0000000..c0d55dd
--- /dev/null
+++ b/src/main/python/tests/scuro/test_algebra_laws.py
@@ -0,0 +1,144 @@
+# -------------------------------------------------------------
+#
+# 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 itertools
+import unittest
+
+import numpy as np
+
+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,
+}
+
+
+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_order_invariant.py b/src/main/python/tests/scuro/test_algebra_order_invariant.py
new file mode 100644
index 0000000..1f6c126
--- /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 0000000..8861d89
--- /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 0000000..d180978
--- /dev/null
+++ b/src/main/python/tests/scuro/test_algebra_signatures.py
@@ -0,0 +1,90 @@
+# -------------------------------------------------------------
+#
+# 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 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()]
+
+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")),
+]
+
+
+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)
+
+ 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
new file mode 100644
index 0000000..619c686
--- /dev/null
+++ b/src/main/python/tests/scuro/test_algebra_stats.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 itertools
+import unittest
+
+import numpy as np
+
+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)],
+}
+
+
+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
new file mode 100644
index 0000000..fbaa7f7
--- /dev/null
+++ b/src/main/python/tests/scuro/test_algebra_subtree_sharing.py
@@ -0,0 +1,127 @@
+# -------------------------------------------------------------
+#
+# 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
+
+import numpy as np
+
+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)
+
+
+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
+
+
+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
new file mode 100644
index 0000000..4e34a3b
--- /dev/null
+++ b/src/main/python/tests/scuro/test_algebra_typecheck.py
@@ -0,0 +1,133 @@
+# -------------------------------------------------------------
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+# -------------------------------------------------------------
+
+import unittest
+from types import SimpleNamespace
+
+import numpy as np
+
+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)
+
+
+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)], []),
+]
+
+
+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, [])
+
+ 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)
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/src/main/python/tests/scuro/test_data_loaders.py b/src/main/python/tests/scuro/test_data_loaders.py
index 0d5ace0..4affe2d 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 @@
# 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 f79157f..e5eac1c 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 @@
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 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 @@
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 f27c721..6e002c1 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 @@
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 @@
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 @@
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")
