[SYSTEMDS-3977] Split Aware Execution

src/main/python/systemds/scuro/__init__.py

@@ -18,152 +18,113 @@
 # under the License.
 #
 # -------------------------------------------------------------
-from systemds.scuro.dataloader.base_loader import BaseLoader
-from systemds.scuro.dataloader.audio_loader import AudioLoader
-from systemds.scuro.dataloader.image_loader import ImageLoader
-from systemds.scuro.dataloader.video_loader import VideoLoader
-from systemds.scuro.dataloader.text_loader import TextLoader
-from systemds.scuro.dataloader.json_loader import JSONLoader
-from systemds.scuro.dataloader.tabular_loader import TabularLoader
-from systemds.scuro.representations.representation import Representation
-from systemds.scuro.representations.aggregate import Aggregation
-from systemds.scuro.representations.aggregated_representation import (
-    AggregatedRepresentation,
-)
-from systemds.scuro.representations.average import Average
-from systemds.scuro.representations.bert import (
-    Bert,
-    RoBERTa,
-    DistillBERT,
-    ALBERT,
-    ELECTRA,
-)
-from systemds.scuro.representations.bow import BoW
-from systemds.scuro.representations.concatenation import Concatenation
-from systemds.scuro.representations.context import Context
-from systemds.scuro.representations.fusion import Fusion
-from systemds.scuro.representations.glove import GloVe
-from systemds.scuro.representations.lstm import LSTM
-from systemds.scuro.representations.max import RowMax
-from systemds.scuro.representations.mel_spectrogram import MelSpectrogram
-from systemds.scuro.representations.multimodal_attention_fusion import (
-    AttentionFusion,
-)
-from systemds.scuro.representations.timeseries_representations import (
-    Mean,
-    Max,
-    Min,
-    Kurtosis,
-    Skew,
-    Std,
-    RMS,
-    ACF,
-    FrequencyMagnitude,
-    SpectralCentroid,
-    Quantile,
-    ZeroCrossingRate,
-    BandpowerFFT,
-)
-from systemds.scuro.representations.mfcc import MFCC
-from systemds.scuro.representations.hadamard import Hadamard
-from systemds.scuro.representations.optical_flow import OpticalFlow
-from systemds.scuro.representations.representation import Representation
-from systemds.scuro.representations.representation_dataloader import NPY
-from systemds.scuro.representations.representation_dataloader import JSON
-from systemds.scuro.representations.representation_dataloader import Pickle
-from systemds.scuro.representations.resnet import ResNet
-from systemds.scuro.representations.openface import OpenFace
+from importlib import import_module
 
-try:
-    from systemds.scuro.representations.image_bind import ImageBind
-except ImportError as _imagebind_import_error:  # pragma: no cover
-    import warnings as _warnings
-
-    _warnings.warn(
-        f"ImageBind representation unavailable, it is excluded from the search "
-        f"({_imagebind_import_error})."
-    )
-    ImageBind = None
-
-from systemds.scuro.representations.spectrogram import Spectrogram
-from systemds.scuro.representations.sum import Sum
-from systemds.scuro.representations.swin_video_transformer import SwinVideoTransformer
-from systemds.scuro.representations.tabular_features import TabularFeatures
-from systemds.scuro.representations.tfidf import TfIdf
-from systemds.scuro.representations.unimodal import UnimodalRepresentation
-from systemds.scuro.representations.wav2vec import Wav2Vec
-from systemds.scuro.representations.window_aggregation import (
-    WindowAggregation,
-    DynamicWindow,
-    StaticWindow,
-)
-from systemds.scuro.representations.physiological_window import (
-    AdaptiveWindow,
-    PhysiologicalEventWindow,
-)
-from systemds.scuro.representations.word2vec import W2V
-from systemds.scuro.representations.x3d import X3D
-from systemds.scuro.representations.color_histogram import ColorHistogram
-from systemds.scuro.models.model import Model
-from systemds.scuro.models.discrete_model import DiscreteModel
-from systemds.scuro.modality.joined import JoinedModality
-from systemds.scuro.modality.joined_transformed import JoinedTransformedModality
-from systemds.scuro.modality.modality import Modality
-from systemds.scuro.utils.identifier import Identifier
-from systemds.scuro.modality.transformed import TransformedModality
-from systemds.scuro.modality.type import ModalityType
-from systemds.scuro.modality.unimodal_modality import UnimodalModality
-from systemds.scuro.drsearch.task import Task
-from systemds.scuro.drsearch.operator_registry import Registry
-from systemds.scuro.drsearch.optimization_data import OptimizationData
-from systemds.scuro.representations.covarep_audio_features import (
-    RMSE,
-    Spectral,
-    ZeroCrossing,
-    Pitch,
-)
-from systemds.scuro.drsearch.multimodal_optimizer import MultimodalOptimizer
-from systemds.scuro.drsearch.unimodal_optimizer import UnimodalOptimizer
-from systemds.scuro.representations.vgg import VGG19
-from systemds.scuro.representations.clip import CLIPText, CLIPVisual
-from systemds.scuro.representations.text_context import (
-    SentenceBoundarySplit,
-    OverlappingSplit,
-)
-from systemds.scuro.representations.text_context_with_indices import (
-    SentenceBoundarySplitIndices,
-    OverlappingSplitIndices,
-)
-from systemds.scuro.representations.elmo import ELMoRepresentation
-from systemds.scuro.representations.dimensionality_reduction import (
-    DimensionalityReduction,
-)
-from systemds.scuro.representations.mlp_averaging import MLPAveraging
-from systemds.scuro.representations.mlp_learned_dim_reduction import (
-    MLPLearnedDimReduction,
+_EXPORT_GROUPS = (
+    ("dataloader.base_loader", "BaseLoader"),
+    ("dataloader.audio_loader", "AudioLoader"),
+    ("dataloader.image_loader", "ImageLoader"),
+    ("dataloader.video_loader", "VideoLoader"),
+    ("dataloader.text_loader", "TextLoader"),
+    ("dataloader.json_loader", "JSONLoader"),
+    ("dataloader.tabular_loader", "TabularLoader"),
+    (
+        "representations.representation",
+        "FitRequirement Representation RepresentationFitState",
+    ),
+    ("representations.aggregate", "Aggregation"),
+    ("representations.aggregated_representation", "AggregatedRepresentation"),
+    ("representations.average", "Average"),
+    ("representations.bert", "Bert RoBERTa DistillBERT ALBERT ELECTRA"),
+    ("representations.bow", "BoW"),
+    ("representations.concatenation", "Concatenation"),
+    ("representations.context", "Context"),
+    ("representations.fusion", "Fusion"),
+    ("representations.glove", "GloVe"),
+    ("representations.lstm", "LSTM"),
+    ("representations.max", "RowMax"),
+    ("representations.mel_spectrogram", "MelSpectrogram"),
+    ("representations.multimodal_attention_fusion", "AttentionFusion"),
+    (
+        "representations.timeseries_representations",
+        "Mean Max Min Kurtosis Skew Std RMS ACF FrequencyMagnitude SpectralCentroid Quantile ZeroCrossingRate BandpowerFFT",
+    ),
+    ("representations.mfcc", "MFCC"),
+    ("representations.hadamard", "Hadamard"),
+    ("representations.optical_flow", "OpticalFlow"),
+    ("representations.representation_dataloader", "NPY JSON Pickle"),
+    ("representations.resnet", "ResNet"),
+    ("representations.openface", "OpenFace"),
+    ("representations.image_bind", "ImageBind"),
+    ("representations.spectrogram", "Spectrogram"),
+    ("representations.sum", "Sum"),
+    ("representations.swin_video_transformer", "SwinVideoTransformer"),
+    ("representations.tabular_features", "TabularFeatures"),
+    ("representations.tfidf", "TfIdf"),
+    ("representations.unimodal", "UnimodalRepresentation"),
+    ("representations.wav2vec", "Wav2Vec"),
+    (
+        "representations.window_aggregation",
+        "WindowAggregation DynamicWindow StaticWindow",
+    ),
+    ("representations.physiological_window", "AdaptiveWindow PhysiologicalEventWindow"),
+    ("representations.word2vec", "W2V"),
+    ("representations.x3d", "X3D"),
+    ("representations.color_histogram", "ColorHistogram"),
+    ("models.model", "Model"),
+    ("models.discrete_model", "DiscreteModel"),
+    ("modality.joined", "JoinedModality"),
+    ("modality.joined_transformed", "JoinedTransformedModality"),
+    ("modality.modality", "Modality"),
+    ("modality.indexed", "IndexedModality"),
+    ("utils.identifier", "Identifier"),
+    ("modality.transformed", "TransformedModality"),
+    ("modality.type", "ModalityType"),
+    ("modality.unimodal_modality", "UnimodalModality"),
+    ("drsearch.task", "Task"),
+    ("drsearch.evaluation_split", "EvaluationSplit"),
+    ("drsearch.operator_registry", "Registry"),
+    ("drsearch.optimization_data", "OptimizationData"),
+    ("drsearch.multimodal_optimizer", "MultimodalOptimizer"),
+    ("drsearch.unimodal_optimizer", "UnimodalOptimizer"),
+    ("representations.covarep_audio_features", "RMSE Spectral ZeroCrossing Pitch"),
+    ("representations.vgg", "VGG19"),
+    ("representations.clip", "CLIPText CLIPVisual"),
+    ("representations.text_context", "SentenceBoundarySplit OverlappingSplit"),
+    (
+        "representations.text_context_with_indices",
+        "SentenceBoundarySplitIndices OverlappingSplitIndices",
+    ),
+    ("representations.elmo", "ELMoRepresentation"),
+    ("representations.dimensionality_reduction", "DimensionalityReduction"),
+    ("representations.mlp_averaging", "MLPAveraging"),
+    ("representations.mlp_learned_dim_reduction", "MLPLearnedDimReduction"),
+    (
+        "representations.physiological_representations",
+        "SDNN RMSSD pNN RRPerMinute HRVBandPower HRVVLF HRVLF HRVHF HRVLFHF PoincareSD1 PoincareSD2 SCLSlope SCLDynamicRange SCRPeaksPerMinute SCRAverageAmplitude SCRAverageDuration",
+    ),
 )
+_EXPORTS = {
+    name: (module_name, name)
+    for module_name, names in _EXPORT_GROUPS
+    for name in names.split()
+}
 
 
-from systemds.scuro.representations.physiological_representations import (
-    SDNN,
-    RMSSD,
-    pNN,
-    RRPerMinute,
-    HRVBandPower,
-    HRVVLF,
-    HRVLF,
-    HRVHF,
-    HRVLFHF,
-    PoincareSD1,
-    PoincareSD2,
-    # SampleEntropy,
-    # DFAAlpha,
-    SCLSlope,
-    SCLDynamicRange,
-    SCRPeaksPerMinute,
-    SCRAverageAmplitude,
-    SCRAverageDuration,
-)
+def __getattr__(name):
+    try:
+        module_name, attribute = _EXPORTS[name]
+    except KeyError as exc:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from exc
+    try:
+        value = getattr(import_module(f"{__name__}.{module_name}"), attribute)
+    except ImportError:
+        if name != "ImageBind":
+            raise
+        value = None
+    globals()[name] = value
+    return value
+
 
 __all__ = [
     "BaseLoader",
@@ -256,7 +217,6 @@
     "MLPAveraging",
     "MLPLearnedDimReduction",
     "DimensionalityReduction",
-    "MultimodalGAPymooOptimizer",
     "SDNN",
     "RMSSD",
     "pNN",

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

@@ -19,6 +19,7 @@
 #
 # -------------------------------------------------------------
 from dataclasses import dataclass
+import math
 import os
 from typing import List, Optional, Union
 import librosa
@@ -56,7 +57,12 @@ def __init__(
         )
         self.normalize = normalize
         self.load_data_from_file = load
-        self.stats = self.get_stats(source_path)
+        files = self.get_file_names(self.indices)
+        self.stats = self.cached_stats(
+            files,
+            {"dtype": str(data_type), "normalize": normalize},
+            lambda: self.get_stats(source_path),
+        )
 
     def extract(self, file: str, index: Optional[Union[str, List[str]]] = None):
         self.file_sanity_check(file)
@@ -80,15 +86,20 @@ def get_stats(self, source_path: str):
         max_length = 0
         avg_length = 0
         num_instances = 0
-        for file in os.listdir(source_path):
-            file_name = file.split(".")[0]
-            if file_name not in self.indices:
-                continue
-            self.file_sanity_check(source_path + file)
-            audio, sr = librosa.load(source_path + file, dtype=self._data_type)
+        import soundfile
+
+        for file in self.get_file_names(self.indices):
+            self.file_sanity_check(file)
+            try:
+                info = soundfile.info(file)
+                sr = 22050
+                length = math.ceil(info.frames * sr / info.samplerate)
+            except (RuntimeError, TypeError):
+                audio, sr = librosa.load(file, dtype=self._data_type)
+                length = audio.shape[0]
             num_instances += 1
             sampling_rate = max(sampling_rate, sr)
-            max_length = max(max_length, audio.shape[0])
-            avg_length += audio.shape[0]
+            max_length = max(max_length, length)
+            avg_length += length
         avg_length /= num_instances
         return AudioStats(sampling_rate, max_length, avg_length, num_instances, True)

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

@@ -18,7 +18,10 @@
 # under the License.
 #
 # -------------------------------------------------------------
+import hashlib
+import json
 import os
+import pickle
 from abc import ABC, abstractmethod
 from collections.abc import Sequence
 from typing import Callable, Iterator, List, Optional, Tuple, Union
@@ -77,6 +80,50 @@ def __init__(
         self.stats = None
         self.chunk_size = chunk_size
 
+    def cached_stats(self, files, options, compute):
+        if os.environ.get("SCURO_DISABLE_STATS_CACHE", "0") == "1":
+            return compute()
+        if isinstance(files, (str, os.PathLike)):
+            files = [files]
+        fingerprint = []
+        for file in files:
+            path = os.path.realpath(os.fspath(file))
+            stat = os.stat(path)
+            fingerprint.append((path, stat.st_size, stat.st_mtime_ns))
+        key = json.dumps(
+            {
+                "schema": 2,
+                "loader": f"{type(self).__module__}.{type(self).__qualname__}",
+                "files": fingerprint,
+                "options": options,
+            },
+            sort_keys=True,
+            default=str,
+        ).encode()
+        digest = hashlib.sha256(key).hexdigest()
+        cache_root = os.environ.get("SCURO_STATS_CACHE_DIR")
+        if cache_root is None:
+            cache_root = os.path.join(
+                os.environ.get("XDG_CACHE_HOME", os.path.expanduser("~/.cache")),
+                "scuro",
+                "loader_stats",
+            )
+        cache_path = os.path.join(cache_root, f"{digest}.pkl")
+        try:
+            with open(cache_path, "rb") as cache_file:
+                return pickle.load(cache_file)
+        except (FileNotFoundError, EOFError, OSError, pickle.PickleError):
+            stats = compute()
+        try:
+            os.makedirs(cache_root, exist_ok=True)
+            temporary_path = f"{cache_path}.{os.getpid()}.tmp"
+            with open(temporary_path, "wb") as cache_file:
+                pickle.dump(stats, cache_file, protocol=pickle.HIGHEST_PROTOCOL)
+            os.replace(temporary_path, cache_path)
+        except OSError:
+            pass
+        return stats
+
     @property
     def chunk_size(self):
         return self._chunk_size