[SYSTEMDS-3913] Improve Representation Handling in Scuro

src/main/python/systemds/scuro/dataloader/audio_loader.py

@@ -20,12 +20,11 @@
 # -------------------------------------------------------------
 from dataclasses import dataclass
 import math
-import os
 from typing import List, Optional, Union
 import librosa
 import numpy as np
 
-from systemds.scuro.dataloader.base_loader import BaseLoader
+from systemds.scuro.dataloader.base_loader import BaseLoader, LazyFileSequence
 from systemds.scuro.modality.type import ModalityType
 
 
@@ -51,34 +50,53 @@ def __init__(
         chunk_size: Optional[int] = None,
         normalize: bool = True,
         load=True,
+        sampling_rate: Optional[int] = None,
     ):
         super().__init__(
             source_path, indices, data_type, chunk_size, ModalityType.AUDIO
         )
         self.normalize = normalize
         self.load_data_from_file = load
+        self.sampling_rate = sampling_rate
         files = self.get_file_names(self.indices)
         self.stats = self.cached_stats(
             files,
-            {"dtype": str(data_type), "normalize": normalize},
+            {
+                "dtype": str(data_type),
+                "normalize": normalize,
+                "sampling_rate": sampling_rate,
+            },
             lambda: self.get_stats(source_path),
         )
 
-    def extract(self, file: str, index: Optional[Union[str, List[str]]] = None):
+    def _load(self, indices: List[str]):
+        if self.data is not None and len(self.data) == len(indices):
+            return self.data, self.metadata
+
+        file_names = self.get_file_names(indices)
+        if isinstance(file_names, str):
+            file_names = [file_names]
+        decoded = LazyFileSequence(file_names, self._decode_file).materialize()
+        self.data = [audio for audio, _ in decoded]
+        self.metadata = [metadata for _, metadata in decoded]
+        return self.data, self.metadata
+
+    def _decode_file(self, file: str):
         self.file_sanity_check(file)
         if not self.load_data_from_file:
-            import numpy as np
-
             audio = np.array([0])
             sr = 1000
         else:
-            audio, sr = librosa.load(file, dtype=self._data_type)
+            audio, sr = librosa.load(file, sr=self.sampling_rate, dtype=self._data_type)
 
             if self.normalize:
                 audio = librosa.util.normalize(audio)
 
-        self.metadata.append(self.modality_type.create_metadata(sr, audio))
+        return audio, self.modality_type.create_metadata(sr, audio)
 
+    def extract(self, file: str, index: Optional[Union[str, List[str]]] = None):
+        audio, metadata = self._decode_file(file)
+        self.metadata.append(metadata)
         self.data.append(audio)
 
     def get_stats(self, source_path: str):
@@ -92,10 +110,12 @@ def get_stats(self, source_path: str):
             self.file_sanity_check(file)
             try:
                 info = soundfile.info(file)
-                sr = 22050
+                sr = self.sampling_rate or info.samplerate
                 length = math.ceil(info.frames * sr / info.samplerate)
             except (RuntimeError, TypeError):
-                audio, sr = librosa.load(file, dtype=self._data_type)
+                audio, sr = librosa.load(
+                    file, sr=self.sampling_rate, dtype=self._data_type
+                )
                 length = audio.shape[0]
             num_instances += 1
             sampling_rate = max(sampling_rate, sr)

src/main/python/systemds/scuro/dataloader/base_loader.py

@@ -24,13 +24,24 @@
 import pickle
 from abc import ABC, abstractmethod
 from collections.abc import Sequence
+from concurrent.futures import ThreadPoolExecutor
 from typing import Callable, Iterator, List, Optional, Tuple, Union
 import math
 from numbers import Integral
 
 import numpy as np
 
 
+def decode_threads() -> int:
+    explicit = os.environ.get("SCURO_DECODE_THREADS")
+    if explicit:
+        try:
+            return max(1, int(explicit))
+        except ValueError:
+            pass
+    return max(1, min(16, os.cpu_count() or 1))
+
+
 class LazyFileSequence(Sequence):
     """List-like file references decoded only when a sample is requested."""
 
@@ -49,6 +60,13 @@ def __len__(self):
     def subset(self, indices):
         return type(self)([self.file_names[i] for i in indices], self.decoder)
 
+    def materialize(self, num_threads: Optional[int] = None) -> list:
+        num_threads = decode_threads() if num_threads is None else int(num_threads)
+        if num_threads <= 1 or len(self) <= 1:
+            return [self.decoder(file_name) for file_name in self.file_names]
+        with ThreadPoolExecutor(max_workers=min(num_threads, len(self))) as pool:
+            return list(pool.map(self.decoder, self.file_names))
+
 
 class BaseLoader(ABC):
     def __init__(
@@ -254,6 +272,13 @@ def get_file_names(self, indices=None):
     def extract(self, file: str, index: Optional[Union[str, List[str]]] = None):
         pass
 
+    def estimate_decoded_bytes(self, metadata=None) -> Optional[int]:
+        """
+        Bytes of all instances once decoded (of the instances `metadata`
+        describes, if given), or None if not known before decoding.
+        """
+        return None
+
     @staticmethod
     def file_sanity_check(file):
         """
@@ -271,6 +296,7 @@ def file_sanity_check(file):
     def resolve_data_type(data_type):
         if isinstance(data_type, str):
             if data_type.lower() in [
+                "uint8",
                 "float16",
                 "float32",
                 "float64",
@@ -282,6 +308,7 @@ def resolve_data_type(data_type):
             else:
                 raise ValueError(f"Unsupported data_type string: {data_type}")
         elif data_type in [
+            np.uint8,
             np.float16,
             np.float32,
             np.float64,

src/main/python/systemds/scuro/dataloader/image_loader.py

@@ -46,7 +46,7 @@ def __init__(
         self,
         source_path: str,
         indices: List[str],
-        data_type: Union[np.dtype, str] = np.float16,
+        data_type: Union[np.dtype, str] = np.uint8,
         chunk_size: Optional[int] = None,
         load=True,
         ext=".jpg",
@@ -139,6 +139,10 @@ def get_stats(self, source_path: str):
             average_channels,
         )
 
+    def estimate_decoded_bytes(self, metadata=None) -> int:
+        metadata = self._all_metadata if metadata is None else metadata
+        return int(sum(int(md["height"]) * int(md["width"]) * 3 for md in metadata))
+
     def estimate_peak_memory_bytes(self) -> dict:
         n = self.chunk_size if self.chunk_size is not None else 1
         per_instance = (

src/main/python/systemds/scuro/dataloader/video_loader.py

@@ -69,7 +69,7 @@ def __init__(
         self,
         source_path: str,
         indices: List[str],
-        data_type: Union[np.dtype, str] = np.float16,
+        data_type: Union[np.dtype, str] = np.uint8,
         chunk_size: Optional[int] = None,
         load=True,
         fps=None,
@@ -158,7 +158,8 @@ def _decode_file(self, file: str):
             frame_interval = self._frame_interval(source_fps)
             stored_fps = source_fps / frame_interval if source_fps else source_fps
 
-            scale_denominator = np.dtype(self._data_type).type(255.0)
+            data_type = np.dtype(self._data_type)
+            scale_denominator = data_type.type(255.0) if data_type.kind == "f" else None
             frames = []
             idx = 0
             while True:
@@ -168,7 +169,10 @@ def _decode_file(self, file: str):
                 if idx % frame_interval == 0:
                     frame = self._fit_frame(frame)
                     frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-                    frame = frame.astype(self._data_type) / scale_denominator
+                    if scale_denominator is not None:
+                        frame = frame.astype(data_type) / scale_denominator
+                    elif data_type != frame.dtype:
+                        frame = frame.astype(data_type)
                     frames.append(frame)
                 idx += 1
         finally:
@@ -259,6 +263,20 @@ def get_stats(self, source_path: str):
             shape_variance,
         )
 
+    def estimate_decoded_bytes(self, metadata=None) -> int:
+        itemsize = np.dtype(self._data_type).itemsize
+        metadata = self._all_metadata if metadata is None else metadata
+        return int(
+            sum(
+                int(md["length"])
+                * int(md["height"])
+                * int(md["width"])
+                * int(md["num_channels"])
+                for md in metadata
+            )
+            * itemsize
+        )
+
     def estimate_peak_memory_bytes(self) -> dict:
         stats = self.stats
         n = self.chunk_size if self.chunk_size is not None else 1

src/main/python/systemds/scuro/drsearch/algebra.py

@@ -144,7 +144,7 @@ def _output_type(op_cls):
 
 class Violation(NamedTuple):
     severity: str  # error | implicit
-    kind: str  # input_type | arity | hidden_aggregation | zero_padding
+    kind: str  # input_type | arity | rank | hidden_aggregation | zero_padding
     node_id: str
     operator: str
     detail: str
@@ -176,6 +176,24 @@ def check_fusion(node):
             if len(set(widths)) > 1:
                 flag("implicit", "zero_padding", node, f"operand widths {widths}")
 
+    def rank(nid):
+        """seq or vector, None if unknown. Raw modalities declare a max_length."""
+        node = nodes[nid]
+        if node.inputs and node.operation is not None:
+            declared = operator_signature(node.operation, node.parameters).output_rank
+            if declared != "any":
+                return "vector" if declared == "vector" else "seq"
+        s = stats(nid)
+        if s is None or not s.output_shape:
+            return None
+        seq = len(s.output_shape) > 1 or getattr(s, "max_length", None) is not None
+        return "seq" if seq else "vector"
+
+    def check_rank(node, sig):
+        got = rank(node.inputs[0])
+        if got is not None and (got == "vector") != (sig.input_rank == "vector"):
+            flag("error", "rank", node, f"{got} input, expects {sig.input_rank}")
+
     def out_type(nid):
         if nid in types:
             return types[nid]
@@ -193,6 +211,8 @@ def out_type(nid):
         for t in in_types:
             if t is not None and t not in sig.input_types:
                 flag("error", "input_type", node, f"{t.name} input")
+        if sig.arity == 1 and len(node.inputs) == 1 and sig.input_rank != "any":
+            check_rank(node, sig)
         if sig.kind == "fuse" and leaf_stats is not None:
             check_fusion(node)
         types[nid] = sig.output_type or in_types[0]
@@ -419,7 +439,8 @@ def compute(nid, laws):
         if not inputs:
             return ("leaf", modality_id, representation_index)
 
-        children = [sig(i, laws & EXACT_LAWS if _fitted(op) else laws) for i in inputs]
+        below = laws & EXACT_LAWS if _fitted(op) else laws - {"concat_perm"}
+        children = [sig(i, below) for i in inputs]
         params = dict(params)
         pushed = (
             params.pop("_pushdown_aggregation", None)
@@ -462,7 +483,7 @@ def compute(nid, laws):
 def count_fusion_plans(n, fusion_ops, laws, invariant) -> dict:
     heads = _fusion_heads(fusion_ops, laws)
     f = len(heads)
-    reorders = sum(_reorders(op_laws, invariant) for *_, op_laws in heads)
+    reorders = sum(_reorders(op_laws, False) for *_, op_laws in heads)
     ordered = factorial(2 * n - 2) // factorial(n - 1) * f ** (n - 1)
     if n == 1:
         canonical = 1
@@ -503,19 +524,21 @@ def _fitted(op) -> bool:
 
 def _commutative_plans(n, heads, invariant) -> int:
     plans = {(1, True): 1, (1, False): 1}
+
+    def count(m, free, invariant):
+        total = 0
+        for *_, op, _, op_laws in heads:
+            below = free and not _fitted(op)
+            splits = sum(
+                comb(m, k) * plans[k, below] * plans[m - k, below] for k in range(1, m)
+            )
+            total += splits // (2 if free and _reorders(op_laws, invariant) else 1)
+        return total
+
     for m in range(2, n + 1):
         for free in (True, False):
-            plans[m, free] = 0
-            for *_, op, _, op_laws in heads:
-                below = free and not _fitted(op)
-                splits = sum(
-                    comb(m, k) * plans[k, below] * plans[m - k, below]
-                    for k in range(1, m)
-                )
-                plans[m, free] += splits // (
-                    2 if free and _reorders(op_laws, invariant) else 1
-                )
-    return plans[n, True]
+            plans[m, free] = count(m, free, False)
+    return count(n, True, invariant) if n > 1 else 1
 
 
 def _reorders(op_laws, invariant) -> bool: