You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@systemds.apache.org by cd...@apache.org on 2026/04/10 12:53:30 UTC
(systemds) branch main updated: [SYSTEMDS-3940] Rework unimodal optimizer
This is an automated email from the ASF dual-hosted git repository.
cdionysio pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/systemds.git
The following commit(s) were added to refs/heads/main by this push:
new a78b38fae5 [SYSTEMDS-3940] Rework unimodal optimizer
a78b38fae5 is described below
commit a78b38fae5e030f4ed35ad256598d6663c4b1a56
Author: Christina Dionysio <di...@tu-berlin.de>
AuthorDate: Fri Apr 10 14:53:25 2026 +0200
[SYSTEMDS-3940] Rework unimodal optimizer
In this patch a new unimodal optimizing strategy is introduced. The dags are executed in parallel via a memory aware node scheduler and executor.
The scheduler tracks the dependencies between the nodes and records the memory used by the system in order to run new nodes that fit these memory limits. Every representation needs to provide a memory estimate so the scheduler knows which representation nodes can be scheduled. The memory estimates consist of a cpu and gpu estimate and the scheduler is aware of both those attributes.
(The hyperparameter tuner and multimodal optimizer still need to be adapted to this new approach, therefore the tests are currently disabled)
---
.github/workflows/python.yml | 2 -
.../systemds/scuro/dataloader/audio_loader.py | 34 ++
.../systemds/scuro/dataloader/base_loader.py | 25 +-
.../systemds/scuro/dataloader/image_loader.py | 58 ++
.../systemds/scuro/dataloader/json_loader.py | 85 ++-
.../systemds/scuro/dataloader/text_loader.py | 38 ++
.../systemds/scuro/dataloader/timeseries_loader.py | 30 +-
.../systemds/scuro/dataloader/video_loader.py | 71 ++-
.../systemds/scuro/drsearch/dag_group_executor.py | 170 ++++++
.../systemds/scuro/drsearch/dag_group_scheduler.py | 159 ++++++
.../scuro/drsearch/hyperparameter_tuner.py | 154 ++++--
.../scuro/drsearch/modality_shared_memory.py | 403 ++++++++++++++
.../scuro/drsearch/multimodal_optimizer.py | 87 ++-
.../systemds/scuro/drsearch/node_executor.py | 615 +++++++++++++++++++++
.../systemds/scuro/drsearch/node_scheduler.py | 444 +++++++++++++++
.../systemds/scuro/drsearch/operator_registry.py | 32 +-
src/main/python/systemds/scuro/drsearch/ranking.py | 2 +-
.../systemds/scuro/drsearch/representation_dag.py | 493 +++++++++++++++--
.../drsearch/representation_dag_visualizer.py | 42 +-
src/main/python/systemds/scuro/drsearch/task.py | 98 +++-
.../systemds/scuro/drsearch/unimodal_optimizer.py | 487 +++++++++-------
src/main/python/systemds/scuro/modality/joined.py | 10 +-
.../python/systemds/scuro/modality/modality.py | 4 +-
.../python/systemds/scuro/modality/transformed.py | 64 ++-
src/main/python/systemds/scuro/modality/type.py | 9 +
.../systemds/scuro/modality/unimodal_modality.py | 62 ++-
.../systemds/scuro/representations/aggregate.py | 75 ++-
.../representations/aggregated_representation.py | 99 +++-
.../python/systemds/scuro/representations/bert.py | 281 +++++++++-
.../python/systemds/scuro/representations/bow.py | 41 +-
.../python/systemds/scuro/representations/clip.py | 329 +++++++++--
.../scuro/representations/color_histogram.py | 60 +-
.../scuro/representations/concatenation.py | 33 +-
.../systemds/scuro/representations/context.py | 15 +-
.../representations/covarep_audio_features.py | 123 ++++-
.../representations/dimensionality_reduction.py | 4 +-
.../python/systemds/scuro/representations/elmo.py | 14 +-
.../representations/frequency_representation.py | 76 +++
.../python/systemds/scuro/representations/glove.py | 31 +-
.../systemds/scuro/representations/hadamard.py | 27 +-
.../systemds/scuro/representations/image_bind.py | 2 +-
.../python/systemds/scuro/representations/lstm.py | 17 +-
.../scuro/representations/mel_spectrogram.py | 64 ++-
.../python/systemds/scuro/representations/mfcc.py | 64 ++-
.../scuro/representations/mlp_averaging.py | 106 +++-
.../representations/mlp_learned_dim_reduction.py | 34 +-
.../systemds/scuro/representations/optical_flow.py | 4 +-
.../scuro/representations/representation.py | 24 +
.../systemds/scuro/representations/resnet.py | 188 +++++--
.../systemds/scuro/representations/spectrogram.py | 92 ++-
.../python/systemds/scuro/representations/sum.py | 27 +-
.../representations/swin_video_transformer.py | 32 +-
.../representations/text_context_with_indices.py | 97 +++-
.../python/systemds/scuro/representations/tfidf.py | 35 +-
.../representations/timeseries_representations.py | 69 ++-
.../systemds/scuro/representations/unimodal.py | 2 +-
.../python/systemds/scuro/representations/vgg.py | 131 ++++-
.../scuro/representations/video_context.py | 72 +++
.../systemds/scuro/representations/wav2vec.py | 10 +-
.../scuro/representations/window_aggregation.py | 291 ++++++++--
.../systemds/scuro/representations/word2vec.py | 29 +-
.../python/systemds/scuro/representations/x3d.py | 54 +-
.../python/systemds/scuro/utils/checkpointing.py | 164 ++++++
.../python/systemds/scuro/utils/memory_utility.py | 194 +++++++
.../systemds/scuro/utils/static_variables.py | 117 +++-
.../python/systemds/scuro/utils/torch_dataset.py | 12 +-
src/main/python/tests/scuro/data_generator.py | 331 +++++++----
src/main/python/tests/scuro/test_data_loaders.py | 176 ++++--
src/main/python/tests/scuro/test_hp_tuner.py | 297 +++++-----
.../python/tests/scuro/test_multimodal_fusion.py | 131 ++---
src/main/python/tests/scuro/test_scheduler.py | 230 ++++++++
.../python/tests/scuro/test_unimodal_optimizer.py | 109 ++--
.../tests/scuro/test_unimodal_representations.py | 30 +-
.../python/tests/scuro/test_window_operations.py | 43 ++
74 files changed, 6993 insertions(+), 1171 deletions(-)
diff --git a/.github/workflows/python.yml b/.github/workflows/python.yml
index ea8c9f485e..73d53ea679 100644
--- a/.github/workflows/python.yml
+++ b/.github/workflows/python.yml
@@ -158,8 +158,6 @@ jobs:
env:
TORCH_HOME: ${{ github.workspace }}/.torch
run: |
- df -h
- exit
( while true; do echo "."; sleep 25; done ) &
KA=$!
pip install --upgrade pip wheel setuptools
diff --git a/src/main/python/systemds/scuro/dataloader/audio_loader.py b/src/main/python/systemds/scuro/dataloader/audio_loader.py
index 8b224e6f25..e0b9e61a07 100644
--- a/src/main/python/systemds/scuro/dataloader/audio_loader.py
+++ b/src/main/python/systemds/scuro/dataloader/audio_loader.py
@@ -18,6 +18,8 @@
# under the License.
#
# -------------------------------------------------------------
+from dataclasses import dataclass
+import os
from typing import List, Optional, Union
import librosa
import numpy as np
@@ -26,6 +28,19 @@ from systemds.scuro.dataloader.base_loader import BaseLoader
from systemds.scuro.modality.type import ModalityType
+@dataclass
+class AudioStats:
+ sampling_rate: int
+ max_length: int
+ avg_length: float
+ num_instances: int
+ output_shape_is_known: bool
+
+ @property
+ def output_shape(self):
+ return (self.max_length,)
+
+
class AudioLoader(BaseLoader):
def __init__(
self,
@@ -41,6 +56,7 @@ class AudioLoader(BaseLoader):
)
self.normalize = normalize
self.load_data_from_file = load
+ self.stats = self.get_stats(source_path)
def extract(self, file: str, index: Optional[Union[str, List[str]]] = None):
self.file_sanity_check(file)
@@ -59,3 +75,21 @@ class AudioLoader(BaseLoader):
self.metadata[file] = self.modality_type.create_metadata(sr, audio)
self.data.append(audio)
+
+ def get_stats(self, source_path: str):
+ sampling_rate = 0
+ 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)
+ num_instances += 1
+ sampling_rate = max(sampling_rate, sr)
+ max_length = max(max_length, audio.shape[0])
+ avg_length += audio.shape[0]
+ avg_length /= num_instances
+ return AudioStats(sampling_rate, max_length, avg_length, num_instances, True)
diff --git a/src/main/python/systemds/scuro/dataloader/base_loader.py b/src/main/python/systemds/scuro/dataloader/base_loader.py
index 777c72d5d5..88decd641f 100644
--- a/src/main/python/systemds/scuro/dataloader/base_loader.py
+++ b/src/main/python/systemds/scuro/dataloader/base_loader.py
@@ -20,7 +20,7 @@
# -------------------------------------------------------------
import os
from abc import ABC, abstractmethod
-from typing import List, Optional, Union
+from typing import Iterator, List, Optional, Tuple, Union
import math
import numpy as np
@@ -55,7 +55,7 @@ class BaseLoader(ABC):
self._chunk_size = None
self._data_type = data_type
self._ext = ext
-
+ self.stats = None
if chunk_size:
self.chunk_size = chunk_size
@@ -98,6 +98,24 @@ class BaseLoader(ABC):
return self._load(self.indices)
+ def iter_loaded_chunks(
+ self, reset: bool = True
+ ) -> Iterator[Tuple[list, dict, List[str]]]:
+ if reset:
+ self.reset()
+
+ if not self._chunk_size:
+ data, metadata = self._load(self.indices)
+ yield data, metadata, self.indices
+ return
+
+ while self._next_chunk < self._num_chunks:
+ chunk_start = self._next_chunk * self._chunk_size
+ chunk_end = (self._next_chunk + 1) * self._chunk_size
+ chunk_indices = self.indices[chunk_start:chunk_end]
+ data, metadata = self._load_next_chunk()
+ yield data, metadata, chunk_indices
+
def update_chunk_sizes(self, other):
if not self._chunk_size and not other.chunk_size:
return
@@ -125,6 +143,9 @@ class BaseLoader(ABC):
return self._load(next_chunk_indices)
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):
self.extract(file_names, indices)
diff --git a/src/main/python/systemds/scuro/dataloader/image_loader.py b/src/main/python/systemds/scuro/dataloader/image_loader.py
index 21ad27bf04..498ae77a89 100644
--- a/src/main/python/systemds/scuro/dataloader/image_loader.py
+++ b/src/main/python/systemds/scuro/dataloader/image_loader.py
@@ -18,6 +18,8 @@
# under the License.
#
# -------------------------------------------------------------
+from dataclasses import dataclass
+import os
from typing import List, Optional, Union
import numpy as np
@@ -27,6 +29,18 @@ import cv2
from systemds.scuro.modality.type import ModalityType
+@dataclass
+class ImageStats:
+ max_width: int
+ max_height: int
+ max_channels: int
+ num_instances: int
+ output_shape: tuple
+ average_width: int
+ average_height: int
+ average_channels: int
+
+
class ImageLoader(BaseLoader):
def __init__(
self,
@@ -41,6 +55,7 @@ class ImageLoader(BaseLoader):
source_path, indices, data_type, chunk_size, ModalityType.IMAGE, ext
)
self.load_data_from_file = load
+ self.stats = self.get_stats(source_path)
def extract(self, file: str, index: Optional[Union[str, List[str]]] = None):
self.file_sanity_check(file)
@@ -61,3 +76,46 @@ class ImageLoader(BaseLoader):
)
self.data.append(image)
+
+ def get_stats(self, source_path: str):
+ max_width = 0
+ max_height = 0
+ max_channels = 0
+ num_instances = 0
+ average_width = 0
+ average_height = 0
+ average_channels = 0
+ for file in self.indices:
+ path = os.path.join(source_path, f"{file}{self._ext}")
+ # if self.chunk_size is None:
+ # self.extract(path)
+ # md = self.metadata[path]
+ # max_width = max(max_width, md["width"])
+ # max_height = max(max_height, md["height"])
+ # max_channels = max(max_channels, md["num_channels"])
+ # num_instances += 1
+ # else:
+ self.file_sanity_check(path)
+ image = cv2.imread(path, cv2.IMREAD_COLOR)
+ image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+ height, width, channels = image.shape
+ max_width = max(max_width, width)
+ max_height = max(max_height, height)
+ max_channels = max(max_channels, channels)
+ num_instances += 1
+ average_width += width
+ average_height += height
+ average_channels += channels
+ average_width = average_width / num_instances
+ average_height = average_height / num_instances
+ average_channels = average_channels / num_instances
+ return ImageStats(
+ max_width,
+ max_height,
+ max_channels,
+ num_instances,
+ (average_width, average_height, average_channels),
+ average_width,
+ average_height,
+ average_channels,
+ )
diff --git a/src/main/python/systemds/scuro/dataloader/json_loader.py b/src/main/python/systemds/scuro/dataloader/json_loader.py
index 53e98e7e19..f5ffd89ea3 100644
--- a/src/main/python/systemds/scuro/dataloader/json_loader.py
+++ b/src/main/python/systemds/scuro/dataloader/json_loader.py
@@ -18,15 +18,26 @@
# under the License.
#
# -------------------------------------------------------------
+from dataclasses import dataclass
import json
-
+import os
import numpy as np
-
+from typing import Tuple
from systemds.scuro.modality.type import ModalityType
from systemds.scuro.dataloader.base_loader import BaseLoader
from typing import Optional, List, Union
+@dataclass
+class JSONStats:
+ num_instances: int
+ max_length: int
+ avg_length: float
+ max_words: int
+ avg_words: float
+ output_shape: Tuple[int]
+
+
class JSONLoader(BaseLoader):
def __init__(
self,
@@ -41,6 +52,7 @@ class JSONLoader(BaseLoader):
source_path, indices, data_type, chunk_size, ModalityType.TEXT, ext
)
self.field = field
+ self.stats = self.get_stats(source_path)
def extract(self, file: str, index: Optional[Union[str, List[str]]] = None):
self.file_sanity_check(file)
@@ -58,3 +70,72 @@ class JSONLoader(BaseLoader):
text = " ".join(text) if isinstance(text, list) else text
self.data.append(text)
self.metadata[idx] = self.modality_type.create_metadata(len(text), text)
+
+ def get_stats(self, source_path: str):
+ self.file_sanity_check(source_path)
+ num_instances = 0
+ max_length = 0
+ avg_length = 0
+ max_words = 0
+ avg_words = 0
+ if os.path.isfile(source_path):
+ with open(source_path) as f:
+ json_file = json.load(f)
+ for id in self.indices:
+ try:
+ text = json_file[id][self.field]
+ except:
+ text = json_file[self.field]
+
+ text = " ".join(text) if isinstance(text, list) else text
+ num_instances += 1
+ max_length = max(max_length, len(text))
+ avg_length += len(text)
+ else:
+ for id in self.indices:
+ with open(os.path.join(source_path, f"{id}{self._ext}")) as f:
+ json_file = json.load(f)
+ text = json_file[self.field]
+
+ text = " ".join(text) if isinstance(text, list) else text
+ num_instances += 1
+ max_length = max(max_length, len(text)) # number of characters
+ max_words = max(max_words, len(text.split()))
+ avg_words += len(text.split())
+ avg_length += len(text)
+
+ avg_length /= num_instances
+ avg_words /= num_instances
+ return JSONStats(
+ num_instances,
+ max_length,
+ avg_length,
+ max_words,
+ avg_words,
+ (max_length,),
+ )
+
+ def estimate_peak_memory_bytes(self) -> dict:
+ s = self.stats
+ n = max(1, s.num_instances)
+
+ avg_len = s.avg_length / n if s.avg_length > s.max_length else s.avg_length
+ avg_len = max(1.0, avg_len)
+
+ bytes_per_char = 2
+ str_overhead = 49
+ ptr_size = 8
+ list_header = 56
+ list_overalloc = 1.125
+ metadata_per_instance = 192
+
+ resident_strings = n * (str_overhead + bytes_per_char * avg_len)
+ resident_list = list_header + int(ptr_size * n * list_overalloc)
+ resident_metadata = n * metadata_per_instance
+ resident_total = resident_strings + resident_list + resident_metadata
+
+ parse_factor = 2.0
+ transient_parse = parse_factor * (n * bytes_per_char * avg_len)
+
+ cpu_peak = int(resident_total + transient_parse)
+ return {"cpu_peak_bytes": cpu_peak, "gpu_peak_bytes": 0}
diff --git a/src/main/python/systemds/scuro/dataloader/text_loader.py b/src/main/python/systemds/scuro/dataloader/text_loader.py
index ff75a4f318..8b987f6845 100644
--- a/src/main/python/systemds/scuro/dataloader/text_loader.py
+++ b/src/main/python/systemds/scuro/dataloader/text_loader.py
@@ -18,12 +18,24 @@
# under the License.
#
# -------------------------------------------------------------
+from dataclasses import dataclass
+import os
from systemds.scuro.dataloader.base_loader import BaseLoader
from typing import Optional, Pattern, List, Union
from systemds.scuro.modality.type import ModalityType
import re
+@dataclass
+class TextStats:
+ num_instances: int
+ max_length: int
+ avg_length: float
+ max_words: int
+ avg_words: float
+ output_shape: tuple
+
+
class TextLoader(BaseLoader):
def __init__(
self,
@@ -35,6 +47,7 @@ class TextLoader(BaseLoader):
):
super().__init__(source_path, indices, data_type, chunk_size, ModalityType.TEXT)
self.prefix = prefix
+ self.stats = self.get_stats(source_path)
def extract(self, file: str, index: Optional[Union[str, List[str]]] = None):
self.file_sanity_check(file)
@@ -47,3 +60,28 @@ class TextLoader(BaseLoader):
len(line.split()), line
)
self.data.append(line)
+
+ def get_stats(self, source_path: str):
+ num_instances = 0
+ max_length = 0
+ avg_length = 0
+ max_words = 0
+ avg_words = 0
+ for file in os.listdir(source_path):
+ self.file_sanity_check(source_path + file)
+ with open(source_path + file) as text_file:
+ for line in text_file:
+ if self.prefix:
+ line = re.sub(self.prefix, "", line)
+ line = line.replace("\n", "")
+ length = len(line.split())
+ num_instances += 1
+ max_length = max(max_length, length)
+ avg_length += length
+ max_words = max(max_words, len(line.split()))
+ avg_words += len(line.split())
+ avg_length /= num_instances
+ avg_words /= num_instances
+ return TextStats(
+ num_instances, max_length, avg_length, max_words, avg_words, (max_length,)
+ )
diff --git a/src/main/python/systemds/scuro/dataloader/timeseries_loader.py b/src/main/python/systemds/scuro/dataloader/timeseries_loader.py
index 1e6bfa8fc2..6e40e8eb08 100644
--- a/src/main/python/systemds/scuro/dataloader/timeseries_loader.py
+++ b/src/main/python/systemds/scuro/dataloader/timeseries_loader.py
@@ -18,6 +18,8 @@
# under the License.
#
# -------------------------------------------------------------
+from dataclasses import dataclass
+import os
import numpy as np
from typing import List, Optional, Union
import h5py
@@ -27,6 +29,15 @@ from systemds.scuro.dataloader.base_loader import BaseLoader
from systemds.scuro.modality.type import ModalityType
+@dataclass
+class TimeseriesStats:
+ max_length: int
+ num_instances: int
+ num_signals: int
+ output_shape: tuple
+ output_shape_is_known: bool
+
+
class TimeseriesLoader(BaseLoader):
def __init__(
self,
@@ -46,7 +57,7 @@ class TimeseriesLoader(BaseLoader):
self.sampling_rate = sampling_rate
self.normalize = normalize
self.file_format = file_format.lower()
-
+ self.stats = self.get_stats(source_path)
if self.file_format not in ["npy", "mat", "hdf5", "txt"]:
raise ValueError(f"Unsupported file format: {self.file_format}")
@@ -93,11 +104,11 @@ class TimeseriesLoader(BaseLoader):
return data
def _load_npy(self, file: str) -> np.ndarray:
- data = np.load(file).astype(self._data_type)
+ data = np.load(file).astype(self._data_type, copy=False)
return data
def _load_txt(self, file: str) -> np.ndarray:
- data = np.loadtxt(file).astype(self._data_type)
+ data = np.loadtxt(file, dtype=self._data_type)
return data
def _load_txt_with_header(self, file: str) -> np.ndarray:
@@ -127,3 +138,16 @@ class TimeseriesLoader(BaseLoader):
selected = [name for name in self.signal_names if name in df.columns]
data = df[selected].to_numpy(dtype=self._data_type)
return data
+
+ def get_stats(self, source_path: str):
+ pass # TODO: Implement this
+ # self.file_sanity_check(source_path)
+ # max_length = 0
+ # num_instances = 0
+ # num_signals = 0
+ # for file in os.listdir(source_path):
+ # data = self._load_npy(source_path + file)
+ # max_length = max(max_length, data.shape[0])
+ # num_instances += 1
+ # num_signals = max(num_signals, data.shape[1])
+ # return TimeseriesStats(max_length, num_instances, num_signals)
diff --git a/src/main/python/systemds/scuro/dataloader/video_loader.py b/src/main/python/systemds/scuro/dataloader/video_loader.py
index 2fee7cbf5a..e57f685e03 100644
--- a/src/main/python/systemds/scuro/dataloader/video_loader.py
+++ b/src/main/python/systemds/scuro/dataloader/video_loader.py
@@ -18,6 +18,8 @@
# under the License.
#
# -------------------------------------------------------------
+from dataclasses import dataclass
+import os
from typing import List, Optional, Union
import numpy as np
@@ -27,6 +29,28 @@ import cv2
from systemds.scuro.modality.type import ModalityType
+@dataclass
+class VideoStats:
+ fps: int
+ max_length: int
+ max_width: int
+ max_height: int
+ max_channels: int
+ num_instances: int
+
+ @property
+ def output_shape(self):
+ """
+ Approximate output shape for raw video tensors.
+
+ This is used by generic resource estimation logic which expects
+ a stats object to expose an ``output_shape`` iterable describing
+ the per-instance tensor shape. For videos we approximate this as
+ (max_length, max_height, max_width, max_num_channels).
+ """
+ return (self.max_length, self.max_height, self.max_width, self.max_channels)
+
+
class VideoLoader(BaseLoader):
def __init__(
self,
@@ -42,6 +66,7 @@ class VideoLoader(BaseLoader):
)
self.load_data_from_file = load
self.fps = fps
+ self.stats = self.get_stats(source_path)
def extract(self, file: str, index: Optional[Union[str, List[str]]] = None):
self.file_sanity_check(file)
@@ -75,8 +100,52 @@ class VideoLoader(BaseLoader):
break
if idx % frame_interval == 0:
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
- frame = frame.astype(self._data_type) / 255.0
+ frame = frame.astype(self._data_type, copy=False) / 255.0
frames.append(frame)
idx += 1
self.data.append(np.stack(frames))
+
+ def get_stats(self, source_path: str):
+ self.file_sanity_check(source_path)
+ fps = 0
+ max_length = 0
+ max_width = 0
+ max_height = 0
+ max_num_channels = 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)
+ cap = cv2.VideoCapture(source_path + file)
+
+ length = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+ width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+ height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+ num_channels = 3
+ max_length = max(max_length, length)
+ max_width = max(max_width, width)
+ max_height = max(max_height, height)
+ max_num_channels = max(max_num_channels, num_channels)
+ num_instances += 1
+ return VideoStats(
+ fps, max_length, max_width, max_height, max_num_channels, num_instances
+ )
+
+ def estimate_peak_memory_bytes(self) -> dict:
+ s = self.stats
+ if self.chunk_size is not None:
+ n = self.chunk_size
+ else:
+ n = s.num_instances
+ return {
+ "cpu_peak_bytes": n
+ * s.output_shape[0]
+ * s.output_shape[1]
+ * s.output_shape[2]
+ * s.output_shape[3]
+ * 4,
+ "gpu_peak_bytes": 0,
+ }
diff --git a/src/main/python/systemds/scuro/drsearch/dag_group_executor.py b/src/main/python/systemds/scuro/drsearch/dag_group_executor.py
new file mode 100644
index 0000000000..57007cb06b
--- /dev/null
+++ b/src/main/python/systemds/scuro/drsearch/dag_group_executor.py
@@ -0,0 +1,170 @@
+# -------------------------------------------------------------
+#
+# 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 concurrent.futures import ProcessPoolExecutor, as_completed
+import multiprocessing as mp
+import os
+import pickle
+
+import time
+from typing import Any, Dict, List, Optional
+from systemds.scuro import Modality
+from systemds.scuro.drsearch.representation_dag import (
+ LRUCache,
+ RepresentationDag,
+ group_dags_by_dependencies,
+)
+from systemds.scuro.utils.checkpointing import CheckpointManager
+from systemds.scuro.drsearch.dag_group_scheduler import DAGGroupScheduler
+
+
+def _process_dag_group(
+ dag_group_pickle: bytes,
+ modality_pickle: bytes,
+ tasks_pickle: bytes,
+ modality_id: int,
+ dag_group_idx: int,
+) -> Dict[str, Any]:
+ checkpoint_manager = CheckpointManager(
+ checkpoint_dir=os.getcwd(),
+ prefix=f"unimodal_checkpoint_group_{modality_id}_{dag_group_idx}_",
+ checkpoint_every=1,
+ resume=False,
+ )
+ results = []
+
+ dag_group = pickle.loads(dag_group_pickle)
+ modality = pickle.loads(modality_pickle)
+ tasks = pickle.loads(tasks_pickle)
+
+ group_cache = LRUCache(max_size=6)
+
+ for i, dag in enumerate(dag_group):
+ representation = dag.execute([modality], external_cache=group_cache)
+
+ for task in tasks:
+ start = time.perf_counter()
+ scores = task.run(representation.data)
+ end = time.perf_counter()
+
+ results.append(
+ {
+ "scores": scores,
+ "transform_time": representation.transform_time,
+ "task_name": task.model.name,
+ "task_time": end - start,
+ "dag": dag,
+ "modality_id": modality_id,
+ }
+ )
+
+ checkpoint_manager.increment(modality_id, 1, dag_group_idx=dag_group_idx)
+ checkpoint_manager.checkpoint_if_due(results)
+
+ return {"results": results}
+
+
+class DAGGroupExecutor:
+ def __init__(
+ self,
+ dags: List[RepresentationDag],
+ modalities: List[Modality],
+ tasks: List[Any],
+ checkpoint_manager: Optional[CheckpointManager] = None,
+ max_workers: Optional[int] = None,
+ ):
+ self.dags = dags
+ self.dag_groups = group_dags_by_dependencies(dags)
+ self.modalities = modalities
+ self.tasks = tasks
+ self.max_workers = max_workers or mp.cpu_count()
+ self.checkpoint_manager = checkpoint_manager
+ self.scheduler = DAGGroupScheduler(
+ dag_groups=self.dag_groups, modality=modalities[0]
+ )
+
+ def run(self):
+ results = []
+ ctx = mp.get_context("spawn")
+ max_workers = min(len(self.dag_groups), self.max_workers)
+
+ modality_pickle = pickle.dumps(
+ self.modalities[0]
+ ) # TODO: handle multiple modalities
+ tasks_pickle = pickle.dumps(self.tasks)
+
+ pending_dag_groups = set(range(len(self.dag_groups)))
+ running_dag_groups = {}
+ all_groups_succeeded = True
+ with ProcessPoolExecutor(max_workers=max_workers, mp_context=ctx) as executor:
+ while pending_dag_groups or running_dag_groups:
+ pending_resources = [
+ (
+ i,
+ self.scheduler.group_resources[i][0],
+ self.scheduler.group_resources[i][1],
+ )
+ for i in pending_dag_groups
+ ]
+ ready_to_execute = self.scheduler.get_runnable(
+ pending_resources, max_concurrent=max_workers
+ )
+ for group_id, gpu_id in ready_to_execute:
+ pending_dag_groups.remove(group_id)
+ dag_group = self.dag_groups[group_id]
+ cpu_mem, gpu_mem = self.scheduler.group_resources[group_id]
+
+ future = executor.submit(
+ _process_dag_group,
+ pickle.dumps(dag_group),
+ modality_pickle,
+ tasks_pickle,
+ self.modalities[0].modality_id,
+ group_id,
+ )
+ running_dag_groups[future] = (group_id, cpu_mem, gpu_mem, gpu_id)
+ if not running_dag_groups:
+ break
+ done = next(as_completed(running_dag_groups), None)
+ if done is None:
+ break
+ group_id, cpu_mem, gpu_mem, gpu_id = running_dag_groups.pop(done)
+ self.scheduler.release(cpu_mem, gpu_mem, gpu_id)
+
+ try:
+ result_dict = future.result()
+
+ for result_entry in result_dict["results"]:
+ results.append(
+ {
+ "scores": result_entry["scores"],
+ "transform_time": result_entry["transform_time"],
+ "task_name": result_entry["task_name"],
+ "task_time": result_entry["task_time"],
+ "dag": result_entry["dag"],
+ "modality_id": self.modalities[0].modality_id,
+ }
+ )
+ except Exception as e:
+ all_groups_succeeded = False
+ print(
+ f"Error processing DAG group {group_id} for modality {self.modalities[0].modality_id}: {e}"
+ )
+ return results, all_groups_succeeded
diff --git a/src/main/python/systemds/scuro/drsearch/dag_group_scheduler.py b/src/main/python/systemds/scuro/drsearch/dag_group_scheduler.py
new file mode 100644
index 0000000000..def12219f5
--- /dev/null
+++ b/src/main/python/systemds/scuro/drsearch/dag_group_scheduler.py
@@ -0,0 +1,159 @@
+# -------------------------------------------------------------
+#
+# 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 typing import List, Tuple, Optional, Dict, Any
+import torch
+import psutil
+
+from systemds.scuro.drsearch.representation_dag import RepresentationDag
+from systemds.scuro.modality.modality import Modality
+
+
+def get_peak_memory_from_dag_group(
+ dag_group: List[RepresentationDag], modality: Modality
+) -> tuple[float, float]:
+ peak_memory_cpu = 0.0
+ peak_memory_gpu = 0.0
+ leaf_memory_bytes = modality.estimate_memory_bytes()
+ for dag in dag_group:
+ prev_stats = modality.get_stats()
+ for node in dag.nodes[1:]:
+ peak_memory = node.operation().estimate_peak_memory_bytes(prev_stats)
+ peak_memory_cpu = max(peak_memory_cpu, peak_memory["cpu_peak_bytes"])
+ peak_memory_gpu = max(peak_memory_gpu, peak_memory["gpu_peak_bytes"])
+ prev_stats = node.operation().get_output_shape(prev_stats)
+
+ return (peak_memory_cpu + leaf_memory_bytes, peak_memory_gpu)
+
+
+class DAGGroupScheduler:
+ def __init__(
+ self,
+ cpu_margin: float = 0.8,
+ gpu_margin: float = 0.8,
+ shared_state: Optional[Dict[str, Any]] = None,
+ lock=None,
+ dag_groups: List[List[RepresentationDag]] = None,
+ modality: Modality = None,
+ ):
+ self._margin = (cpu_margin, gpu_margin)
+ self._n_gpu = (
+ torch.cuda.device_count() if torch and torch.cuda.is_available() else 0
+ )
+ self._shared = shared_state
+ self._lock = lock
+ if self._shared is None:
+ self._shared = {
+ "cpu_in_use": 0.0,
+ "gpu_in_use": {i: 0.0 for i in range(max(1, self._n_gpu))},
+ }
+ else:
+ self._shared.setdefault("cpu_in_use", 0.0)
+ self._shared.setdefault("gpu_in_use", {})
+ self.group_resources = []
+ for dag_group in dag_groups:
+ cpu_mem, gpu_mem = get_peak_memory_from_dag_group(dag_group, modality)
+ self.group_resources.append((cpu_mem, gpu_mem))
+
+ def _avail_cpu(self) -> float:
+ available_memory = (psutil.virtual_memory().available) if psutil else 4096.0
+ return max(
+ 0.0,
+ available_memory * self._margin[0] - float(self._shared["cpu_in_use"]),
+ )
+
+ def _avail_gpu(self, i: int) -> float:
+ available_gpu_memory = 0.0
+ try:
+ free, _ = torch.cuda.mem_get_info(i)
+ available_gpu_memory = free
+ except Exception:
+ pass
+ g = self._shared["gpu_in_use"]
+ in_use = g.get(str(i), g.get(i, 0.0))
+ return max(0.0, available_gpu_memory * self._margin[1] - in_use)
+
+ def _update(
+ self, cpu_delta: float, gpu_delta: float, gpu_id: Optional[int]
+ ) -> None:
+ self._shared["cpu_in_use"] = max(0.0, self._shared["cpu_in_use"] + cpu_delta)
+ if gpu_id is not None and gpu_delta != 0:
+ g = self._shared["gpu_in_use"]
+ k = str(gpu_id)
+ g[k] = max(0.0, g.get(k, 0) + gpu_delta)
+
+ def reserve(self, cpu_mb: float, gpu_mb: float) -> Tuple[bool, Optional[int]]:
+ if self._lock:
+ self._lock.acquire()
+ try:
+ if self._avail_cpu() < cpu_mb:
+ return (False, None)
+ if gpu_mb <= 0:
+ self._update(cpu_mb, 0, None)
+ return (True, None)
+ if self._n_gpu == 0 and self._avail_cpu() >= gpu_mb + cpu_mb:
+ self._update(cpu_mb + gpu_mb, 0, None)
+ return (True, None)
+
+ for i in range(self._n_gpu):
+ if i is not None and self._avail_gpu(i) >= gpu_mb:
+ self._update(cpu_mb, gpu_mb, i)
+ return (True, i)
+
+ return (False, None)
+ finally:
+ if self._lock:
+ self._lock.release()
+
+ def release(self, cpu_mb: float, gpu_mb: float, gpu_id: Optional[int]) -> None:
+ if self._lock:
+ self._lock.acquire()
+ try:
+ if self._n_gpu == 0:
+ self._update(-cpu_mb - gpu_mb, 0, None)
+ return
+ if gpu_id is not None:
+ self._update(-cpu_mb, -gpu_mb, gpu_id)
+ else:
+ self._update(-cpu_mb - gpu_mb, 0, None)
+ finally:
+ if self._lock:
+ self._lock.release()
+
+ def get_runnable(
+ self,
+ pending_resources: List[Tuple[int, float, float, Optional[int]]],
+ max_concurrent: Optional[int] = None,
+ largest_first: bool = True,
+ ) -> List[Tuple[int, Optional[int]]]:
+ order = sorted(
+ range(len(pending_resources)),
+ key=lambda i: pending_resources[i][1] + pending_resources[i][2],
+ reverse=largest_first,
+ )
+ runnable = []
+ for i in order:
+ if max_concurrent is not None and len(runnable) >= max_concurrent:
+ break
+ idx, cpu_bytes, gpu_bytes = pending_resources[i]
+ ok, assigned = self.reserve(cpu_bytes, gpu_bytes)
+ if ok:
+ runnable.append((idx, assigned))
+ return runnable
diff --git a/src/main/python/systemds/scuro/drsearch/hyperparameter_tuner.py b/src/main/python/systemds/scuro/drsearch/hyperparameter_tuner.py
index 0a129c8eb4..64141dc929 100644
--- a/src/main/python/systemds/scuro/drsearch/hyperparameter_tuner.py
+++ b/src/main/python/systemds/scuro/drsearch/hyperparameter_tuner.py
@@ -19,6 +19,7 @@
#
# -------------------------------------------------------------
from typing import Dict, List, Tuple, Any, Optional
+import os
from skopt.space import Real, Integer, Categorical
import numpy as np
import logging
@@ -33,6 +34,7 @@ from systemds.scuro.drsearch.representation_dag import (
from systemds.scuro.modality.modality import Modality
from systemds.scuro.drsearch.task import PerformanceMeasure
import pickle
+from systemds.scuro.utils.checkpointing import CheckpointManager
def get_params_for_node(node_id, params):
@@ -113,6 +115,8 @@ class HyperparameterTuner:
maximize_metric: bool = True,
save_results: bool = False,
debug: bool = False,
+ checkpoint_every: Optional[int] = None,
+ resume: bool = True,
):
self.tasks = tasks
self.unimodal_optimization_results = optimization_results
@@ -130,6 +134,14 @@ class HyperparameterTuner:
self.extract_k_best_modalities_per_task()
self.debug = debug
self.logger = logging.getLogger(__name__)
+ self.checkpoint_every = checkpoint_every
+ self.resume = resume
+ self._checkpoint_manager = CheckpointManager(
+ os.getcwd(),
+ "hyperparam_checkpoint_",
+ checkpoint_every=self.checkpoint_every,
+ resume=self.resume,
+ )
if debug:
logging.basicConfig(
level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
@@ -176,17 +188,54 @@ class HyperparameterTuner:
self.k_best_cache[task.model.name].extend(cached_data)
self.representations = representations
+ def _count_results_by_task(self, results: Dict[str, Any]) -> Dict[str, int]:
+ counts = {}
+ for task_name, task_results in results.items():
+ task_count = 0
+ for _, result_list in task_results.items():
+ task_count += len(result_list)
+ counts[task_name] = task_count
+ return counts
+
+ def resume_from_checkpoint(self):
+ loaded = self._checkpoint_manager.resume_from_checkpoint(
+ "eval_count_by_task", self._count_results_by_task
+ )
+ if loaded:
+ results, _, _ = loaded
+ self.optimization_results.results = results
+
def tune_unimodal_representations(self, max_eval_per_rep: Optional[int] = None):
+ self.resume_from_checkpoint()
for task in self.tasks:
reps = self.k_best_representations[task.model.name]
- self.optimization_results.add_result(
- Parallel(n_jobs=self.n_jobs)(
- delayed(self.tune_dag_representation)(
- rep.dag, rep.dag.root_node_id, task, max_eval_per_rep
- )
- for rep in reps
- )
+ skip_remaining = self._checkpoint_manager.skip_remaining_by_key.get(
+ task.model.name, 0
)
+ if skip_remaining >= len(reps):
+ continue
+
+ chunk_size = self.checkpoint_every or len(reps)
+ for start_idx in range(skip_remaining, len(reps), chunk_size):
+ rep_chunk = reps[start_idx : start_idx + chunk_size]
+ try:
+ results = []
+ for rep in rep_chunk:
+ results.append(
+ self.tune_dag_representation(
+ rep.dag, rep.dag.root_node_id, task, max_eval_per_rep
+ )
+ )
+ self.optimization_results.add_result(results)
+ self._checkpoint_manager.increment(task.model.name, len(results))
+ self._checkpoint_manager.checkpoint_if_due(
+ self.optimization_results.results, "eval_count_by_task"
+ )
+ except Exception:
+ self._checkpoint_manager.save_checkpoint(
+ self.optimization_results.results, "eval_count_by_task", {}
+ )
+ raise
if self.save_results:
self.save_tuning_results()
@@ -287,13 +336,13 @@ class HyperparameterTuner:
opt = Optimizer(
search_space, random_state=42, n_initial_points=min(10, n_calls // 2)
)
-
+ self.n_jobs = 2
n_batch = min(abs(self.n_jobs), n_calls) if self.n_jobs != 0 else 1
for _ in range(0, n_calls, n_batch):
points = opt.ask(n_points=n_batch)
- results = Parallel(n_jobs=self.n_jobs)(
- delayed(evaluate_point)(p) for p in points
- )
+ results = Parallel(
+ n_jobs=self.n_jobs, max_nbytes=None, mmap_mode=None, backend="threading"
+ )(delayed(evaluate_point)(p) for p in points)
objective_values = [result[0] for result in results]
all_results.extend(result[1] for result in results)
opt.tell(points, objective_values)
@@ -369,6 +418,7 @@ class HyperparameterTuner:
optimize_unimodal: bool = True,
max_eval_per_rep: Optional[int] = None,
):
+ self.resume_from_checkpoint()
self.optimization_results.setup_mm(optimize_unimodal)
for task in self.tasks:
@@ -389,43 +439,59 @@ class HyperparameterTuner:
reverse=self.maximize_metric,
)[:k]
best_optimization_results = best_results
+ skip_remaining = self._checkpoint_manager.skip_remaining_by_key.get(
+ task.model.name, 0
+ )
for representation in best_optimization_results:
- if optimize_unimodal:
- dag = copy.deepcopy(representation.dag)
- index = 0
- for i, node in enumerate(representation.dag.nodes):
- if not node.inputs:
- leaf_node_id = node.node_id
- leaf_nodes = self.representations[task.model.name][
- node.modality_id
- ][node.representation_index].dag.nodes
- for leaf_idx, node in enumerate(dag.nodes):
- if node.node_id == leaf_node_id:
- dag.nodes[leaf_idx : leaf_idx + 1] = leaf_nodes
- index = leaf_idx + len(leaf_nodes) - 1
- break
-
- for node in dag.nodes:
- try:
- idx = node.inputs.index(leaf_node_id)
- node.inputs[idx] = dag.nodes[index].node_id
- break
- except ValueError:
- continue
-
- result = self.tune_dag_representation(
- dag, dag.root_node_id, task, max_eval_per_rep
+ if skip_remaining > 0:
+ skip_remaining -= 1
+ continue
+ try:
+ if optimize_unimodal:
+ dag = copy.deepcopy(representation.dag)
+ index = 0
+ for i, node in enumerate(representation.dag.nodes):
+ if not node.inputs:
+ leaf_node_id = node.node_id
+ leaf_nodes = self.representations[task.model.name][
+ node.modality_id
+ ][node.representation_index].dag.nodes
+ for leaf_idx, node in enumerate(dag.nodes):
+ if node.node_id == leaf_node_id:
+ dag.nodes[leaf_idx : leaf_idx + 1] = leaf_nodes
+ index = leaf_idx + len(leaf_nodes) - 1
+ break
+
+ for node in dag.nodes:
+ try:
+ idx = node.inputs.index(leaf_node_id)
+ node.inputs[idx] = dag.nodes[index].node_id
+ break
+ except ValueError:
+ continue
+
+ result = self.tune_dag_representation(
+ dag, dag.root_node_id, task, max_eval_per_rep
+ )
+ else:
+ result = self.tune_dag_representation(
+ representation.dag,
+ representation.dag.root_node_id,
+ task,
+ max_eval_per_rep,
+ mm_opt=True,
+ )
+ self.optimization_results.add_result([result])
+ self._checkpoint_manager.increment(task.model.name, 1)
+ self._checkpoint_manager.checkpoint_if_due(
+ self.optimization_results.results, "eval_count_by_task"
)
- else:
- result = self.tune_dag_representation(
- representation.dag,
- representation.dag.root_node_id,
- task,
- max_eval_per_rep,
- mm_opt=True,
+ except Exception:
+ self._checkpoint_manager.save_checkpoint(
+ self.optimization_results.results, "eval_count_by_task", {}
)
- self.optimization_results.add_result([result])
+ raise
if self.save_results:
self.save_tuning_results()
diff --git a/src/main/python/systemds/scuro/drsearch/modality_shared_memory.py b/src/main/python/systemds/scuro/drsearch/modality_shared_memory.py
new file mode 100644
index 0000000000..a98592b620
--- /dev/null
+++ b/src/main/python/systemds/scuro/drsearch/modality_shared_memory.py
@@ -0,0 +1,403 @@
+# -------------------------------------------------------------
+#
+# 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 typing import Any, List, Tuple
+import numpy as np
+from multiprocessing import shared_memory
+
+SHARED_MEMORY_MIN_BYTES = 1 * 1024 * 1024
+
+
+class SharedStringList:
+ def __init__(
+ self, shm_name: str, offsets: List[Tuple[int, int]], payload_nbytes: int
+ ):
+ self.shm_name = shm_name
+ self.offsets = offsets
+ self.payload_nbytes = int(payload_nbytes)
+ self._shm = None
+
+ def _ensure_attached(self):
+ if self._shm is None:
+ self._shm = shared_memory.SharedMemory(name=self.shm_name)
+
+ def __len__(self):
+ return len(self.offsets)
+
+ def __getitem__(self, idx):
+ self._ensure_attached()
+ if isinstance(idx, slice):
+ return [self[i] for i in range(*idx.indices(len(self)))]
+ start, length = self.offsets[idx]
+ buf = self._shm.buf[start : start + length]
+ return buf.tobytes().decode("utf-8")
+
+ def __iter__(self):
+ for i in range(len(self)):
+ yield self[i]
+
+ def close(self):
+ if self._shm is not None:
+ self._shm.close()
+ self._shm = None
+
+ def __del__(self):
+ try:
+ self.close()
+ except Exception:
+ pass
+
+ def __getstate__(self):
+ state = self.__dict__.copy()
+ state["_shm"] = None
+ return state
+
+
+class SharedGroupedArrayList:
+ def __init__(
+ self,
+ shm_name: str,
+ dtype_str: str,
+ offsets: List[tuple],
+ total_elems: int,
+ group_bounds: List[Tuple[int, int]],
+ ):
+ self.shm_name = shm_name
+ self.dtype_str = dtype_str
+ self.offsets = offsets
+ self.total_elems = int(total_elems)
+ self.group_bounds = group_bounds
+ self._dtype = np.dtype(dtype_str)
+ self._shm = None
+ self._buffer = None
+
+ def _ensure_attached(self):
+ if self._shm is None:
+ self._shm = shared_memory.SharedMemory(name=self.shm_name)
+ self._buffer = np.ndarray(
+ (self.total_elems,), dtype=self._dtype, buffer=self._shm.buf
+ )
+
+ def _leaf_view(self, leaf_idx: int) -> np.ndarray:
+ start, size, shape = self.offsets[leaf_idx]
+ view = self._buffer[start : start + size].reshape(shape)
+ view.setflags(write=False)
+ return view
+
+ def __len__(self):
+ return len(self.group_bounds)
+
+ def __getitem__(self, idx):
+ self._ensure_attached()
+ if isinstance(idx, slice):
+ return [self[i] for i in range(*idx.indices(len(self)))]
+ lo, hi = self.group_bounds[idx]
+ return [self._leaf_view(j) for j in range(lo, hi)]
+
+ def __iter__(self):
+ for i in range(len(self)):
+ yield self[i]
+
+ def close(self):
+ if self._shm is not None:
+ self._shm.close()
+ self._shm = None
+ self._buffer = None
+
+ def __del__(self):
+ try:
+ self.close()
+ except Exception:
+ pass
+
+ def __getstate__(self):
+ state = self.__dict__.copy()
+ state["_shm"] = None
+ state["_buffer"] = None
+ return state
+
+
+class SharedNDArray:
+ def __init__(self, shm_name: str, dtype_str: str, shape: tuple):
+ self.shm_name = shm_name
+ self.dtype_str = dtype_str
+ self.shape = tuple(shape)
+ self._dtype = np.dtype(dtype_str)
+ self._shm = None
+ self._arr = None
+
+ def _ensure_attached(self):
+ if self._shm is None:
+ self._shm = shared_memory.SharedMemory(name=self.shm_name)
+ self._arr = np.ndarray(self.shape, dtype=self._dtype, buffer=self._shm.buf)
+ self._arr.setflags(write=False)
+
+ def to_numpy(self, copy: bool = False):
+ self._ensure_attached()
+ return self._arr.copy() if copy else self._arr
+
+ def __array__(self, dtype=None):
+ arr = self.to_numpy(copy=False)
+ return arr.astype(dtype, copy=False) if dtype is not None else arr
+
+ @property
+ def dtype(self):
+ return self._dtype
+
+ @property
+ def ndim(self):
+ return len(self.shape)
+
+ def __len__(self):
+ return self.shape[0] if len(self.shape) > 0 else 0
+
+ def __getitem__(self, idx):
+ self._ensure_attached()
+ return self._arr[idx]
+
+ def __iter__(self):
+ self._ensure_attached()
+ return iter(self._arr)
+
+ def close(self):
+ if self._shm is not None:
+ self._shm.close()
+ self._shm = None
+ self._arr = None
+
+ def __del__(self):
+ try:
+ self.close()
+ except Exception:
+ pass
+
+ def __getstate__(self):
+ state = self.__dict__.copy()
+ state["_shm"] = None
+ state["_arr"] = None
+ return state
+
+
+class SharedArrayList(list):
+ def __init__(
+ self,
+ shm_name: str,
+ dtype_str: str,
+ offsets: List[tuple],
+ total_elems: int,
+ ):
+ super().__init__()
+ self.shm_name = shm_name
+ self.dtype_str = dtype_str
+ self.offsets = offsets
+ self.total_elems = int(total_elems)
+ self._dtype = np.dtype(dtype_str)
+ self._shm = None
+ self._buffer = None
+
+ def _ensure_attached(self):
+ if self._shm is None:
+ self._shm = shared_memory.SharedMemory(name=self.shm_name)
+ self._buffer = np.ndarray(
+ (self.total_elems,), dtype=self._dtype, buffer=self._shm.buf
+ )
+
+ def __len__(self):
+ return len(self.offsets)
+
+ def __getitem__(self, idx):
+ self._ensure_attached()
+ if isinstance(idx, slice):
+ return [self[i] for i in range(*idx.indices(len(self)))]
+ start, size, shape = self.offsets[idx]
+ view = self._buffer[start : start + size].reshape(shape)
+ view.setflags(write=False)
+ return view
+
+ def __iter__(self):
+ for i in range(len(self)):
+ yield self[i]
+
+ def close(self):
+ if self._shm is not None:
+ self._shm.close()
+ self._shm = None
+ self._buffer = None
+
+ def __del__(self):
+ try:
+ self.close()
+ except Exception:
+ pass
+
+ def __getstate__(self):
+ state = self.__dict__.copy()
+ state["_shm"] = None
+ state["_buffer"] = None
+ return state
+
+
+def _is_string_list_shared_memory_candidate(data: Any) -> bool:
+ if not isinstance(data, list) or not data:
+ return False
+ return all(type(x) is str for x in data)
+
+
+def _estimate_ndarray_list_nbytes(data: List[np.ndarray]) -> int:
+ return int(sum(arr.nbytes for arr in data if isinstance(arr, np.ndarray)))
+
+
+def _is_shared_ndarray_candidate(data: Any) -> bool:
+ return (
+ isinstance(data, np.ndarray)
+ and data.dtype.kind in {"f", "i", "u", "b"}
+ and data.ndim >= 1
+ and data.size > 0
+ )
+
+
+def _is_shared_memory_candidate(data: Any) -> bool:
+ if not isinstance(data, list) or not data:
+ return False
+ if not all(isinstance(x, np.ndarray) for x in data):
+ return False
+ if not all(x.dtype.kind in {"f", "i", "u", "b"} for x in data):
+ return False
+ dtypes = {x.dtype.str for x in data}
+ return len(dtypes) == 1
+
+
+def _is_nested_shared_memory_candidate(data: Any) -> bool:
+ if not isinstance(data, list) or not data:
+ return False
+ if not all(isinstance(x, list) for x in data):
+ return False
+ dtypes = set()
+ for group in data:
+ for arr in group:
+ if not isinstance(arr, np.ndarray):
+ return False
+ if arr.dtype.kind not in {"f", "i", "u", "b"}:
+ return False
+ if arr.ndim < 1 or arr.size == 0:
+ return False
+ dtypes.add(arr.dtype.str)
+ if not dtypes:
+ return False
+ return len(dtypes) == 1
+
+
+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:
+ dtype = data[0].dtype
+ offsets = []
+ total_elems = 0
+ for arr in data:
+ arr_size = int(arr.size)
+ offsets.append((total_elems, arr_size, tuple(arr.shape)))
+ total_elems += arr_size
+
+ shm = shared_memory.SharedMemory(create=True, size=data_nbytes)
+ buffer = np.ndarray((total_elems,), dtype=dtype, buffer=shm.buf)
+ cursor = 0
+ for arr in data:
+ flat = np.asarray(arr, dtype=dtype).reshape(-1)
+ n = flat.size
+ buffer[cursor : cursor + n] = flat
+ cursor += n
+
+ data = SharedArrayList(shm.name, dtype.str, offsets, total_elems)
+ resident_bytes = min(
+ resident_bytes, max(2 * 1024 * 1024, len(offsets) * 64)
+ )
+ shm.close()
+ return data, shm.name, data_nbytes, resident_bytes
+ elif _is_shared_ndarray_candidate(data):
+ arr = data
+ data_nbytes = int(arr.nbytes)
+ if data_nbytes >= SHARED_MEMORY_MIN_BYTES:
+ 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")
+ shm_arr.setflags(write=False)
+
+ data = SharedNDArray(shm.name, arr.dtype.str, arr.shape)
+ resident_bytes = min(resident_bytes, 2 * 1024 * 1024)
+
+ shm.close()
+ return data, shm.name, data_nbytes, resident_bytes
+ elif _is_nested_shared_memory_candidate(data):
+ leaves: List[np.ndarray] = []
+ group_bounds: List[Tuple[int, int]] = []
+ for group in data:
+ lo = len(leaves)
+ for arr in group:
+ leaves.append(arr)
+ group_bounds.append((lo, len(leaves)))
+ data_nbytes = _estimate_ndarray_list_nbytes(leaves)
+ if data_nbytes >= SHARED_MEMORY_MIN_BYTES:
+ dtype = leaves[0].dtype
+ offsets = []
+ total_elems = 0
+ for arr in leaves:
+ arr_size = int(arr.size)
+ offsets.append((total_elems, arr_size, tuple(arr.shape)))
+ total_elems += arr_size
+
+ shm = shared_memory.SharedMemory(create=True, size=data_nbytes)
+ buffer = np.ndarray((total_elems,), dtype=dtype, buffer=shm.buf)
+ cursor = 0
+ for arr in leaves:
+ flat = np.asarray(arr, dtype=dtype).reshape(-1)
+ n = flat.size
+ buffer[cursor : cursor + n] = flat
+ cursor += n
+
+ data = SharedGroupedArrayList(
+ shm.name, dtype.str, offsets, total_elems, group_bounds
+ )
+ resident_bytes = min(
+ resident_bytes, max(2 * 1024 * 1024, len(offsets) * 64)
+ )
+ shm.close()
+ return data, shm.name, data_nbytes, resident_bytes
+ elif _is_string_list_shared_memory_candidate(data):
+ encoded = [s.encode("utf-8") for s in data]
+ data_nbytes = int(sum(len(b) for b in encoded))
+ if data_nbytes >= SHARED_MEMORY_MIN_BYTES:
+ shm = shared_memory.SharedMemory(create=True, size=data_nbytes)
+ mv = shm.buf
+ cursor = 0
+ str_offsets: List[Tuple[int, int]] = []
+ for b in encoded:
+ n = len(b)
+ mv[cursor : cursor + n] = b
+ str_offsets.append((cursor, n))
+ cursor += n
+ data = SharedStringList(shm.name, str_offsets, data_nbytes)
+ resident_bytes = min(
+ resident_bytes, max(2 * 1024 * 1024, len(str_offsets) * 32)
+ )
+ shm.close()
+ return data, shm.name, data_nbytes, resident_bytes
+
+ return None, None, 0, resident_bytes
diff --git a/src/main/python/systemds/scuro/drsearch/multimodal_optimizer.py b/src/main/python/systemds/scuro/drsearch/multimodal_optimizer.py
index 7c17353663..596cab0237 100644
--- a/src/main/python/systemds/scuro/drsearch/multimodal_optimizer.py
+++ b/src/main/python/systemds/scuro/drsearch/multimodal_optimizer.py
@@ -43,6 +43,7 @@ import copy
import time
import traceback
from itertools import chain
+from systemds.scuro.utils.checkpointing import CheckpointManager
def _evaluate_dag_worker(dag_pickle, task_pickle, modalities_pickle, debug=False):
@@ -106,6 +107,8 @@ class MultimodalOptimizer:
min_modalities: int = 2,
max_modalities: int = None,
metric: str = "accuracy",
+ checkpoint_every: int = None,
+ resume: bool = True,
):
self.modalities = modalities
self.tasks = tasks
@@ -122,11 +125,40 @@ class MultimodalOptimizer:
unimodal_optimization_results
)
self.optimization_results = []
+ self.checkpoint_every = checkpoint_every
+ self.resume = resume
+ self._checkpoint_manager = CheckpointManager(
+ os.getcwd(),
+ "multimodal_checkpoint_",
+ checkpoint_every=self.checkpoint_every,
+ resume=self.resume,
+ )
+
+ def _count_results_by_task(self, results: Dict[str, List["OptimizationResult"]]):
+ return {
+ task_name: len(task_results) for task_name, task_results in results.items()
+ }
+
+ def resume_from_checkpoint(self):
+ loaded = self._checkpoint_manager.resume_from_checkpoint(
+ "eval_count_by_task", self._count_results_by_task
+ )
+ if loaded:
+ results, _, _ = loaded
+ self.optimization_results = results
+
+ def _ensure_results_initialized(self) -> Dict[str, List["OptimizationResult"]]:
+ if not isinstance(self.optimization_results, dict):
+ self.optimization_results = {}
+ for task in self.tasks:
+ self.optimization_results.setdefault(task.model.name, [])
+ return self.optimization_results
def optimize_parallel(
self, max_combinations: int = None, max_workers: int = 2, batch_size: int = 4
) -> Dict[str, List["OptimizationResult"]]:
- all_results = {}
+ self.resume_from_checkpoint()
+ all_results = self._ensure_results_initialized()
for task in self.tasks:
task_copy = copy.deepcopy(task)
@@ -134,8 +166,12 @@ class MultimodalOptimizer:
print(
f"[DEBUG] Optimizing multimodal fusion for task: {task.model.name}"
)
- all_results[task.model.name] = []
- evaluated_count = 0
+ evaluated_count = self._checkpoint_manager.counts_by_key.get(
+ task.model.name, 0
+ )
+ skip_remaining = self._checkpoint_manager.skip_remaining_by_key.get(
+ task.model.name, 0
+ )
outstanding = set()
stop_generation = False
@@ -167,6 +203,10 @@ class MultimodalOptimizer:
if max_combinations and evaluated_count >= max_combinations:
stop_generation = True
break
+ if skip_remaining > 0:
+ skip_remaining -= 1
+ evaluated_count += 1
+ continue
dag_pickle = pickle.dumps(dag)
fut = ex.submit(
@@ -190,7 +230,18 @@ class MultimodalOptimizer:
except Exception:
if self.debug:
traceback.print_exc()
+ self._checkpoint_manager.save_checkpoint(
+ self.optimization_results,
+ "eval_count_by_task",
+ {},
+ )
evaluated_count += 1
+ self._checkpoint_manager.increment(
+ task.model.name, 1
+ )
+ self._checkpoint_manager.checkpoint_if_due(
+ self.optimization_results, "eval_count_by_task"
+ )
if self.debug and evaluated_count % 100 == 0:
print(
f"[DEBUG] Evaluated {evaluated_count} combinations..."
@@ -211,7 +262,14 @@ class MultimodalOptimizer:
except Exception:
if self.debug:
traceback.print_exc()
+ self._checkpoint_manager.save_checkpoint(
+ self.optimization_results, "eval_count_by_task", {}
+ )
evaluated_count += 1
+ self._checkpoint_manager.increment(task.model.name, 1)
+ self._checkpoint_manager.checkpoint_if_due(
+ self.optimization_results, "eval_count_by_task"
+ )
if self.debug and evaluated_count % 100 == 0:
print(
f"[DEBUG] Evaluated {evaluated_count} combinations..."
@@ -225,8 +283,6 @@ class MultimodalOptimizer:
f"[DEBUG] Task completed: {len(all_results[task.model.name])} valid combinations evaluated"
)
- self.optimization_results = all_results
-
if self.debug:
print(f"[DEBUG] Optimization completed")
@@ -406,15 +462,20 @@ class MultimodalOptimizer:
def optimize(
self, max_combinations: int = None
) -> Dict[str, List["OptimizationResult"]]:
- all_results = {}
+ self.resume_from_checkpoint()
+ all_results = self._ensure_results_initialized()
for task in self.tasks:
if self.debug:
print(
f"[DEBUG] Optimizing multimodal fusion for task: {task.model.name}"
)
- all_results[task.model.name] = []
- evaluated_count = 0
+ evaluated_count = self._checkpoint_manager.counts_by_key.get(
+ task.model.name, 0
+ )
+ skip_remaining = self._checkpoint_manager.skip_remaining_by_key.get(
+ task.model.name, 0
+ )
for modality_subset in self._generate_modality_combinations():
if self.debug:
@@ -427,12 +488,20 @@ class MultimodalOptimizer:
for dag in self._generate_fusion_dags(modality_subset, repr_combo):
if max_combinations and evaluated_count >= max_combinations:
break
+ if skip_remaining > 0:
+ skip_remaining -= 1
+ evaluated_count += 1
+ continue
result = self._evaluate_dag(dag, task)
if result is not None:
all_results[task.model.name].append(result)
evaluated_count += 1
+ self._checkpoint_manager.increment(task.model.name, 1)
+ self._checkpoint_manager.checkpoint_if_due(
+ self.optimization_results, "eval_count_by_task"
+ )
if self.debug and evaluated_count % 100 == 0:
print(f" Evaluated {evaluated_count} combinations...")
@@ -448,8 +517,6 @@ class MultimodalOptimizer:
f"[DEBUG] Task completed: {len(all_results[task.model.name])} valid combinations evaluated"
)
- self.optimization_results = all_results
-
if self.debug:
print(f"[DEBUG] Optimization completed")
diff --git a/src/main/python/systemds/scuro/drsearch/node_executor.py b/src/main/python/systemds/scuro/drsearch/node_executor.py
new file mode 100644
index 0000000000..bf51dfde85
--- /dev/null
+++ b/src/main/python/systemds/scuro/drsearch/node_executor.py
@@ -0,0 +1,615 @@
+# -------------------------------------------------------------
+#
+# 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 concurrent.futures import FIRST_COMPLETED, ProcessPoolExecutor, wait
+from dataclasses import dataclass
+import os
+from multiprocessing import shared_memory
+from systemds.scuro import Modality
+from systemds.scuro.drsearch.modality_shared_memory import add_shared_memory_candidate
+from systemds.scuro.drsearch.node_scheduler import MemoryAwareNodeScheduler
+from systemds.scuro.drsearch.representation_dag import (
+ RepresentationDag,
+ RepresentationNode,
+)
+
+import numpy as np
+from typing import Any, Dict, List, Optional
+import multiprocessing as mp
+import psutil
+import time
+import torch
+from systemds.scuro.drsearch.task import PerformanceMeasure
+from systemds.scuro.representations.context import Context
+from systemds.scuro.representations.dimensionality_reduction import (
+ DimensionalityReduction,
+)
+from systemds.scuro.representations.aggregated_representation import (
+ AggregatedRepresentation,
+)
+from systemds.scuro.representations.representation import RepresentationStats
+from systemds.scuro.representations.unimodal import UnimodalRepresentation
+from systemds.scuro.utils.checkpointing import CheckpointManager
+import threading
+import time
+import psutil
+import os
+from systemds.scuro.utils.static_variables import DEBUG
+
+
+def measure_peak_rss_during(fn, *args, sample_s=0.01, **kwargs):
+ proc = psutil.Process(os.getpid())
+ baseline = proc.memory_info().rss
+ peak = baseline
+ stop = threading.Event()
+
+ def sampler():
+ nonlocal peak
+ while not stop.is_set():
+ rss = proc.memory_info().rss
+ if rss > peak:
+ peak = rss
+ time.sleep(sample_s)
+
+ t = threading.Thread(target=sampler, daemon=True)
+ t.start()
+ try:
+ out = fn(*args, **kwargs)
+ finally:
+ stop.set()
+ t.join()
+
+ return out, (peak - baseline), peak
+
+
+class RefCountResultCache:
+ def __init__(self):
+ self.cache = {}
+ self.ref_count = {}
+ self.memory_usage_per_node = {}
+ self.shared_memory_names = {}
+
+ def get(self, node_id: str) -> Any:
+ return self.cache[node_id]
+
+ def add_result(self, node_id: str, result: Any):
+ resident_bytes = result.calculate_memory_usage()
+ shared_backing_bytes = 0
+
+ if hasattr(result, "data"):
+ try:
+ data, shm_name, data_nbytes, resident_bytes = (
+ add_shared_memory_candidate(result.data, resident_bytes)
+ )
+ if data is not None:
+ result.data = data
+ self.shared_memory_names[node_id] = [shm_name]
+ shared_backing_bytes = data_nbytes
+ except Exception as e:
+ print(
+ f"Failed to move cache entry {node_id} to shared memory, falling back to RAM: {e}"
+ )
+
+ self.cache[node_id] = result
+ self.memory_usage_per_node[node_id] = int(resident_bytes)
+ if DEBUG:
+ print(
+ f"Node {node_id} has a CPU memory usage of {self.memory_usage_per_node[node_id]/1024**3:.5f} GB"
+ + (
+ f" (shared-memory backing: {shared_backing_bytes/1024**3:.5f} GB)"
+ if shared_backing_bytes > 0
+ else ""
+ )
+ )
+
+ def inc_ref(self, node_id: str):
+ if node_id not in self.ref_count:
+ self.ref_count[node_id] = 0
+ self.ref_count[node_id] += 1
+
+ def dec_ref(self, node_id: str):
+ self.ref_count[node_id] -= 1
+ if self.ref_count[node_id] == 0:
+ del self.cache[node_id]
+ del self.ref_count[node_id]
+ del self.memory_usage_per_node[node_id]
+ self._cleanup_shared_memory(node_id)
+
+ def clear(self, node_id: str):
+ if node_id in self.cache:
+ del self.cache[node_id]
+ if node_id in self.ref_count:
+ del self.ref_count[node_id]
+ if node_id in self.memory_usage_per_node:
+ del self.memory_usage_per_node[node_id]
+ self._cleanup_shared_memory(node_id)
+
+ def __len__(self):
+ return len(self.cache)
+
+ def get_memory_total_memory_usage(self):
+ return sum(self.memory_usage_per_node.values())
+
+ def _cleanup_shared_memory(self, node_id: str):
+ names = self.shared_memory_names.pop(node_id, [])
+ for shm_name in names:
+ try:
+ shm = shared_memory.SharedMemory(name=shm_name)
+ shm.close()
+ shm.unlink()
+ except FileNotFoundError:
+ pass
+ except Exception:
+ pass
+
+ def cleanup_all(self):
+ for node_id in list(self.shared_memory_names.keys()):
+ self._cleanup_shared_memory(node_id)
+
+
+def _execute_node_worker(node, input_mods, task, rep_cache, gpu_id):
+ if gpu_id is not None:
+ device = torch.device(f"cuda:{gpu_id}")
+ torch.cuda.set_device(device)
+ torch.cuda.reset_peak_memory_stats(device)
+
+ node_operation = node.operation(params=node.parameters)
+ operation_name = node_operation.name
+ if DEBUG:
+ print(f"Executing node {node.node_id} {operation_name} on GPU {gpu_id}")
+
+ if gpu_id is not None and hasattr(node_operation, "gpu_id"):
+ node_operation.gpu_id = gpu_id
+
+ def _run_node_op():
+ if len(input_mods) == 1:
+ if isinstance(node_operation, Context):
+ return input_mods[0].context(node_operation)
+ elif isinstance(node_operation, DimensionalityReduction):
+ return input_mods[0].dimensionality_reduction(node_operation)
+ elif isinstance(node_operation, AggregatedRepresentation):
+ return node_operation.transform(input_mods[0])
+ elif isinstance(node_operation, UnimodalRepresentation):
+ pushdown_config = node.parameters.get("_pushdown_aggregation", None)
+ agg = None
+ if pushdown_config is not None:
+ agg = AggregatedRepresentation(params=pushdown_config)
+ if rep_cache is not None and node_operation.name in rep_cache:
+ return rep_cache[node_operation.name]
+ return input_mods[0].apply_representation(
+ node_operation, aggregation=agg
+ )
+ return input_mods[0].apply_representation(node_operation)
+ else:
+ fusion_op = node_operation
+ if hasattr(fusion_op, "needs_training") and fusion_op.needs_training:
+ return input_mods[0].combine_with_training(
+ input_mods[1:], fusion_op, task
+ )
+ return input_mods[0].combine(input_mods[1:], fusion_op)
+
+ result, peak_delta_bytes, peak_abs_rss = measure_peak_rss_during(
+ _run_node_op,
+ sample_s=0.01,
+ )
+
+ gpu_peak_bytes = (
+ torch.cuda.max_memory_allocated(device) if gpu_id is not None else 0
+ )
+
+ return {
+ "result": result,
+ "peak_bytes": peak_delta_bytes,
+ "peak_abs_rss_bytes": peak_abs_rss,
+ "gpu_peak_bytes": gpu_peak_bytes,
+ "operation_name": operation_name,
+ }
+
+
+def _execute_task_worker(
+ task_node_id: str, task: Any, data: Any, gpu_id: Optional[int]
+) -> Dict[str, Any]:
+
+ if DEBUG:
+ print(f"Executing task {task_node_id} on GPU {gpu_id}")
+ if gpu_id is not None:
+ device = torch.device(f"cuda:{gpu_id}")
+ torch.cuda.set_device(device)
+ torch.cuda.reset_peak_memory_stats(device)
+
+ if gpu_id is not None and hasattr(task, "model") and hasattr(task.model, "device"):
+ task.model.device = torch.device(f"cuda:{gpu_id}")
+
+ def _run_task():
+ start = time.perf_counter()
+ scores = task.run(data)
+ end = time.perf_counter()
+ return scores, end - start
+
+ gpu_peak_bytes = (
+ torch.cuda.max_memory_allocated(device) if gpu_id is not None else 0
+ )
+ result, peak_delta_bytes, peak_abs_rss = measure_peak_rss_during(
+ _run_task,
+ sample_s=0.01,
+ )
+ if DEBUG:
+ print(
+ f"Task {task_node_id} has a CPU peak memory usage of {peak_delta_bytes/1024**3:.2f} GB, and a GPU peak memory usage of {gpu_peak_bytes/1024**3:.2f} GB"
+ )
+ return {
+ "scores": result[0],
+ "task_time": result[1],
+ "peak_bytes": peak_delta_bytes,
+ "gpu_peak_bytes": gpu_peak_bytes,
+ }
+
+
+class NodeExecutor:
+ def __init__(
+ self,
+ dags: List[RepresentationDag],
+ modalities: List[Modality],
+ tasks: List[Any],
+ checkpoint_manager: Optional[CheckpointManager] = None,
+ max_num_workers: int = -1,
+ result_path: Optional[str] = None,
+ ):
+ available_total_cpu = (
+ float(psutil.virtual_memory().available)
+ - float(psutil.virtual_memory().available) * 0.30
+ )
+ self.dags = dags
+ self.scheduler = MemoryAwareNodeScheduler(
+ dags, modalities, tasks, available_total_cpu
+ )
+ self.checkpoint_manager = CheckpointManager(
+ checkpoint_dir=result_path if result_path is not None else os.getcwd(),
+ prefix=f"node_executor_checkpoint_{modalities[0].modality_id}_",
+ checkpoint_every=1,
+ resume=False,
+ )
+ self.max_num_workers = (
+ min(mp.cpu_count(), max_num_workers)
+ if max_num_workers != -1
+ else mp.cpu_count()
+ )
+ self.modalities = modalities
+ self.tasks = tasks
+ self.result_cache = RefCountResultCache()
+ self.memory_usage_checkpoint = CheckpointManager(
+ checkpoint_dir=result_path if result_path is not None else os.getcwd(),
+ prefix=f"memory_usage_checkpoint_{modalities[0].modality_id}_",
+ checkpoint_every=1,
+ resume=False,
+ )
+
+ def run(self) -> None:
+ task_results = {}
+ memory_usage_data = {}
+
+ self._materialize_leaf_modalities_in_shared_memory()
+
+ ctx = mp.get_context("spawn")
+ with ProcessPoolExecutor(
+ max_workers=self.max_num_workers, mp_context=ctx
+ ) as executor:
+ future_to_node_id = {}
+
+ def submit_node(node_id: str):
+ node = self.scheduler.mapping[node_id]
+ gpu_id = node.gpu_id
+ parent_result = None
+ parent_node_id = self.scheduler.get_valid_parent(node_id)
+ if parent_node_id is not None:
+ parent_result = self.result_cache.get(parent_node_id)
+ if self._is_task_node(node):
+ task_result = ResultEntry(
+ dag=self._get_dag_from_node_ids(node_id),
+ representation_time=parent_result.transform_time,
+ )
+ task_results[node_id] = task_result
+ task_idx = int(node.parameters.get("_task_idx", 0))
+ future = executor.submit(
+ _execute_task_worker,
+ node_id,
+ self.tasks[task_idx],
+ (
+ self.modalities[0].data
+ if parent_result is None
+ else parent_result.data
+ ),
+ gpu_id,
+ )
+ else:
+ future = executor.submit(
+ _execute_node_worker,
+ node,
+ self.modalities if parent_result is None else [parent_result],
+ None,
+ None,
+ gpu_id,
+ )
+ self.scheduler.move_to_running(node_id)
+ future_to_node_id[future] = node_id
+
+ def submit_new_ready_nodes():
+ for node_id in self.scheduler.get_runnable().copy():
+ submit_node(node_id)
+
+ submit_new_ready_nodes()
+
+ while future_to_node_id or not self.scheduler.is_finished():
+ if not future_to_node_id:
+ submit_new_ready_nodes()
+ continue
+
+ done, _ = wait(
+ set(future_to_node_id.keys()), return_when=FIRST_COMPLETED
+ )
+
+ for future in done:
+ node_id = future_to_node_id.pop(future)
+ try:
+ result = future.result()
+ except Exception as e:
+ err_cls = type(e)
+ err_mod = err_cls.__module__
+ if err_mod.startswith("torch"):
+ torch.cuda.empty_cache()
+ print(f"Error executing node {node_id}: {e}")
+ self.scheduler.add_failed_node(node_id)
+ continue
+
+ peak_bytes = result["peak_bytes"]
+ gpu_peak_bytes = result["gpu_peak_bytes"]
+
+ node = self.scheduler.mapping[node_id]
+ if self._is_task_node(node):
+ task_results[node_id].task_time = result["task_time"]
+ task_results[node_id].train_score = result["scores"][
+ 0
+ ].average_scores
+ task_results[node_id].val_score = result["scores"][
+ 1
+ ].average_scores
+ task_results[node_id].test_score = result["scores"][
+ 2
+ ].average_scores
+ self.checkpoint_manager.increment(node_id)
+ self.checkpoint_manager.checkpoint_if_due(task_results)
+ self._checkpoint_memory_usage(
+ node_id,
+ peak_bytes,
+ gpu_peak_bytes,
+ "task",
+ memory_usage_data,
+ None,
+ )
+
+ parent_node_id = self.scheduler.get_valid_parent(node_id)
+ if parent_node_id is not None:
+ self.result_cache.dec_ref(parent_node_id)
+ if (
+ parent_node_id in self.result_cache.ref_count
+ and self.result_cache.ref_count[parent_node_id] == 0
+ ):
+ self.result_cache.clear(parent_node_id)
+ self.scheduler.complete_node(node_id)
+
+ else:
+ transformed_modality = result["result"]
+ actual_stats = self._infer_actual_output_stats(
+ transformed_modality
+ )
+ estimated_stats = self.scheduler.node_stats.get(node_id)
+
+ if actual_stats is not None and (
+ estimated_stats is None
+ or not getattr(
+ estimated_stats, "output_shape_is_known", True
+ )
+ ):
+ self.scheduler.update_node_stats_and_reestimate_descendants(
+ node_id, actual_stats
+ )
+ self._checkpoint_memory_usage(
+ node_id,
+ peak_bytes,
+ gpu_peak_bytes,
+ result["operation_name"],
+ memory_usage_data,
+ transformed_modality.data,
+ )
+ before_bytes = self.result_cache.get_memory_total_memory_usage()
+ self._manage_result_cache(node_id, transformed_modality)
+ after_bytes = self.result_cache.get_memory_total_memory_usage()
+ self.scheduler.update_cpu_memory_in_use(
+ after_bytes - before_bytes
+ )
+ self.scheduler.complete_node(node_id)
+
+ submit_new_ready_nodes()
+ assert len(self.result_cache.ref_count.keys()) == 0
+
+ self.result_cache.cleanup_all()
+ self._cleanup_leaf_shared_memory()
+ return list(task_results.values())
+
+ def _materialize_leaf_modalities_in_shared_memory(self):
+ self._leaf_shm_names = []
+ for modality in self.modalities:
+ if hasattr(modality, "extract_raw_data") and not modality.has_data():
+ modality.extract_raw_data()
+ data, shm_name, _, _ = add_shared_memory_candidate(modality.data)
+ if shm_name is not None:
+ modality.data = data
+ self._leaf_shm_names.append(shm_name)
+
+ def _cleanup_leaf_shared_memory(self):
+ for shm_name in getattr(self, "_leaf_shm_names", []):
+ try:
+ shm = shared_memory.SharedMemory(name=shm_name)
+ shm.close()
+ shm.unlink()
+ except FileNotFoundError:
+ pass
+ except Exception:
+ pass
+ self._leaf_shm_names = []
+
+ def _checkpoint_memory_usage(
+ self,
+ node_id: str,
+ peak_bytes: int,
+ gpu_peak_bytes: int,
+ operation_name: str,
+ data,
+ result,
+ ):
+ self.memory_usage_checkpoint.increment(node_id)
+
+ shape = None
+ if DEBUG:
+ shape = self._print_node_stats(node_id, result, operation_name)
+ if peak_bytes > self.scheduler.node_resources[node_id][0]:
+ print(
+ f"UNDERESTIMATED PEAK MEMORY: Peak bytes: {peak_bytes/1024**3:.2f} GB, Estimated CPU bytes: {self.scheduler.node_resources[node_id][0]/1024**3:.2f} GB for node {node_id}: {operation_name}"
+ )
+ if gpu_peak_bytes > self.scheduler.node_resources[node_id][1]:
+ print(
+ f"UNDERESTIMATED GPU PEAK MEMORY: GPU peak bytes: {gpu_peak_bytes/1024**3:.2f} GB, Estimated GPU bytes: {self.scheduler.node_resources[node_id][1]/1024**3:.2f} GB for node {node_id}: {operation_name}"
+ )
+ if self.scheduler.node_resources[node_id][0] >= peak_bytes * 2:
+ print(
+ f"Peak bytes: {peak_bytes/1024**3:.2f} GB, Estimated CPU bytes: {self.scheduler.node_resources[node_id][0]/1024**3:.2f} GB, 200% of estimated for node {node_id}: {operation_name}"
+ )
+ if self.scheduler.node_resources[node_id][1] > gpu_peak_bytes * 2:
+ print(
+ f"GPU peak bytes: {gpu_peak_bytes/1024**3:.2f} GB, Estimated GPU bytes: {self.scheduler.node_resources[node_id][1]/1024**3:.2f} GB, 200% of estimated for node {node_id}: {operation_name}"
+ )
+ data[node_id] = {
+ "cpu_peak_bytes": peak_bytes,
+ "gpu_peak_bytes": gpu_peak_bytes,
+ "operation_name": operation_name,
+ "estimated_cpu_bytes": self.scheduler.node_resources[node_id][0],
+ "estimated_gpu_bytes": self.scheduler.node_resources[node_id][1],
+ "shape": shape,
+ }
+ self.memory_usage_checkpoint.checkpoint_if_due(data)
+
+ def _print_node_stats(self, node_id: str, result: Any, operation_name: str):
+ if (
+ result is not None
+ and operation_name != "BoW"
+ and not operation_name.endswith("Split")
+ ):
+ node_stats = self.scheduler.node_stats[node_id]
+ shape = None
+ if isinstance(result[0], list):
+ if isinstance(result[0][0], np.ndarray):
+ shape = (len(result[0]), *result[0][0].shape)
+ elif isinstance(result[0][0], list):
+ shape = (len(result[0]), *result[0][0][0].shape)
+ else:
+ shape = (len(result[0]), *result[0][0].shape)
+ else:
+ shape = result[0].shape
+ print(
+ f"Node {node_id} {operation_name} should have shape of {node_stats.num_instances, node_stats.output_shape}, actual shape: {len(result), shape} output shape is known: {node_stats.output_shape_is_known}"
+ )
+ if node_stats.output_shape_is_known:
+ assert (
+ len(result) == node_stats.num_instances
+ ), f"Node {node_id} {operation_name} should have {node_stats.num_instances} instances, actual: {len(result)}"
+ return shape
+
+ def _infer_actual_output_stats(
+ self, transformed_modality: Any
+ ) -> Optional[RepresentationStats]:
+ if transformed_modality is None or not hasattr(transformed_modality, "data"):
+ return None
+
+ data = transformed_modality.data
+
+ if isinstance(data, np.ndarray):
+ if data.ndim == 0:
+ return RepresentationStats(1, (1,), output_shape_is_known=True)
+ num_instances = int(data.shape[0])
+ output_shape = (
+ tuple(int(d) for d in data.shape[1:]) if data.ndim > 1 else (1,)
+ )
+ return RepresentationStats(
+ num_instances, output_shape, output_shape_is_known=True
+ )
+
+ if isinstance(data, list) and len(data) > 0 and isinstance(data[0], np.ndarray):
+ num_instances = len(data)
+ first_shape = tuple(int(d) for d in data[0].shape)
+ same_shape = all(
+ isinstance(x, np.ndarray) and x.shape == data[0].shape for x in data
+ )
+ return RepresentationStats(
+ num_instances,
+ first_shape,
+ output_shape_is_known=bool(same_shape),
+ )
+
+ return None
+
+ def _manage_result_cache(self, node_id: str, result: Any):
+ parent_node_id = self.scheduler.get_valid_parent(node_id)
+ if parent_node_id is not None:
+ self.result_cache.dec_ref(parent_node_id)
+
+ if self.scheduler.get_children(node_id):
+ for _ in self.scheduler.get_children(node_id):
+ self.result_cache.inc_ref(node_id)
+ self.result_cache.add_result(node_id, result)
+
+ if (
+ parent_node_id in self.result_cache.ref_count
+ and self.result_cache.ref_count[parent_node_id] == 0
+ ):
+ self.result_cache.clear(parent_node_id)
+
+ def _get_nodes_by_ids(self, nodes_ids: List[str]) -> List[RepresentationNode]:
+ return [self.scheduler.mapping[node_id] for node_id in nodes_ids]
+
+ def _get_dag_from_node_ids(self, node_id: str) -> RepresentationDag:
+ for dag in self.dags:
+ if dag.root_node_id == node_id:
+ return dag
+ return None
+
+ @staticmethod
+ def _is_task_node(node: RepresentationNode) -> bool:
+ return bool(getattr(node, "parameters", {}).get("_node_kind") == "task")
+
+
+@dataclass
+class ResultEntry:
+ val_score: PerformanceMeasure = None
+ train_score: PerformanceMeasure = None
+ test_score: PerformanceMeasure = None
+ representation_time: float = 0.0
+ task_time: float = 0.0
+ dag: RepresentationDag = None
+ tradeoff_score: float = 0.0
diff --git a/src/main/python/systemds/scuro/drsearch/node_scheduler.py b/src/main/python/systemds/scuro/drsearch/node_scheduler.py
new file mode 100644
index 0000000000..4fe3ac1d03
--- /dev/null
+++ b/src/main/python/systemds/scuro/drsearch/node_scheduler.py
@@ -0,0 +1,444 @@
+# -------------------------------------------------------------
+#
+# 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 __future__ import annotations
+
+from typing import List, Dict, Optional, Any
+from collections import defaultdict, deque
+from collections import deque
+import torch
+from systemds.scuro.drsearch.representation_dag import (
+ RepresentationDag,
+ RepresentationNode,
+)
+from systemds.scuro.modality.modality import Modality
+from systemds.scuro.utils.memory_utility import gpu_memory_info
+from systemds.scuro.utils.static_variables import DEBUG
+
+
+class MemoryAwareNodeScheduler:
+
+ def __init__(
+ self,
+ nodes: List[RepresentationDag],
+ modalities: List[Modality],
+ tasks: List[Any],
+ cpu_memory_budget: float,
+ gpu_memory_budget: Optional[float] = None,
+ ):
+ self.node_stats = {}
+ self.modalities = modalities
+ self.tasks = tasks
+ self.parents = defaultdict(set)
+ self.children = defaultdict(set)
+ self.mapping = {}
+ self.parent_refcounts = defaultdict(int)
+ self.gpu_memory_info = gpu_memory_info()
+ self.memory_budget = {
+ "cpu": cpu_memory_budget,
+ "gpu": {
+ info["index"]: int(info["total_b"]) for info in self.gpu_memory_info
+ },
+ }
+ self.roots = set()
+ self.leaves = set()
+ self.topo_order = []
+ self.nodes = self._get_nodes_from_dags(nodes)
+ self.unresolved_parents = self._get_unresolved_parents()
+ self.node_resources = self._estimate_node_resources()
+ self.success = False
+ self.deadlock = False
+ self.ready_nodes = []
+ self.running_nodes = []
+ self.completed_nodes = []
+ self.failed_nodes = []
+ self.blocked_memory_nodes_perm = []
+ self.cancelled_nodes = []
+ self.n_gpu = (
+ torch.cuda.device_count() if torch and torch.cuda.is_available() else 0
+ )
+ self.memory_stats = {
+ "cpu_in_use": sum([self.node_resources[node][0] for node in self.leaves]),
+ "gpu_in_use": {
+ info["index"]: int(info["total_b"] - info["free_b"])
+ for info in self.gpu_memory_info
+ },
+ }
+ self._initialized = False
+
+ def update_cpu_memory_in_use(self, delta_bytes: int):
+ self.memory_stats["cpu_in_use"] += delta_bytes
+
+ def get_runnable(self) -> List[RepresentationNode]:
+ runnable_nodes = self._get_runnable_nodes()
+
+ for node in runnable_nodes:
+ ok, gpu_id = self._check_memory_constraints(node)
+ if ok:
+ self.mapping[node].gpu_id = gpu_id
+ self.ready_nodes.append(node)
+ self._reserve_memory(node, gpu_id)
+ return self.ready_nodes
+
+ def _get_runnable_nodes(self) -> List[str]:
+ runnable_nodes = []
+ for node in self.topo_order:
+ if (
+ node not in self.leaves
+ and self.unresolved_parents[node] == 0
+ and node not in self.running_nodes
+ and node not in self.completed_nodes
+ and node not in self.ready_nodes
+ ):
+ runnable_nodes.append(node)
+
+ 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]
+ return (-release_bytes, node_id not in self.roots)
+
+ runnable_nodes.sort(key=_score)
+ return runnable_nodes
+
+ def add_failed_node(self, node_id: str):
+ self.failed_nodes.append(node_id)
+ self.running_nodes.remove(node_id)
+
+ self._release_memory(node_id, self.mapping[node_id].gpu_id)
+
+ def move_to_running(self, node_id: str):
+ self.ready_nodes.remove(node_id)
+ self.running_nodes.append(node_id)
+
+ def complete_node(self, node_id: str):
+ self.running_nodes.remove(node_id)
+ self.completed_nodes.append(node_id)
+ self._release_memory(node_id, self.mapping[node_id].gpu_id)
+ self.topo_order.remove(node_id)
+ for child_id in self.children[node_id]:
+ self.parent_refcounts[child_id] -= 1
+ self.unresolved_parents[child_id] -= 1
+
+ for parent_id in self.parents.get(node_id, set()):
+ if self.remaining_children.get(parent_id, 0) > 0:
+ self.remaining_children[parent_id] -= 1
+
+ def is_finished(self) -> bool:
+ if not self._initialized:
+ self._initialized = True
+ return False
+
+ if self.not_enough_memory():
+ self.deadlock = True
+ self.success = False
+ return True
+
+ if self._is_deadlock():
+ self.deadlock = True
+ self.success = False
+ return True
+
+ if self._is_success():
+ self.success = True
+ return True
+
+ return False
+
+ def get_valid_parent(self, node_id: str) -> bool:
+ for parent_id in self.parents[node_id]:
+ if parent_id not in self.leaves:
+ return parent_id
+ return None
+
+ def get_children(self, node_id: str) -> List[str]:
+ return list(self.children[node_id])
+
+ def update_node_stats_and_reestimate_descendants(
+ self, node_id: str, actual_stats: Any
+ ) -> None:
+ if actual_stats is None:
+ return
+
+ self.node_stats[node_id] = actual_stats
+
+ descendants = self._get_descendants_in_topo_order(node_id)
+
+ for desc_id in descendants:
+ if (
+ desc_id in self.leaves
+ or desc_id in self.ready_nodes
+ or desc_id in self.running_nodes
+ or desc_id in self.completed_nodes
+ ):
+ continue
+
+ parent_ids = list(self.parents.get(desc_id, set()))
+ if not parent_ids:
+ continue
+
+ input_stats = self.node_stats.get(parent_ids[0])
+ if input_stats is None:
+ continue
+
+ if desc_id not in self.roots:
+ operation = self.mapping[desc_id].operation(
+ params=self.mapping[desc_id].parameters
+ )
+ peak_memory = operation.estimate_peak_memory_bytes(input_stats)
+ peak_memory["cpu_peak_bytes"] += (
+ 64 * 1024 + 512 * input_stats.num_instances
+ )
+ output_stats = operation.get_output_stats(input_stats)
+
+ self.node_resources[desc_id] = (
+ int(peak_memory["cpu_peak_bytes"]),
+ int(peak_memory["gpu_peak_bytes"]),
+ )
+ self.node_stats[desc_id] = output_stats
+ else:
+ task_node = self.mapping[desc_id]
+ task = self.tasks[task_node.parameters["_task_idx"]]
+ peak_memory = task.estimate_peak_memory_bytes(input_stats)
+
+ cpu_increment = max(
+ int(peak_memory["cpu_peak_bytes"]),
+ 16 * 1024 * 1024,
+ )
+ self.node_resources[desc_id] = (
+ cpu_increment,
+ int(peak_memory["gpu_peak_bytes"]),
+ )
+ self.node_stats[desc_id] = task.get_output_stats(input_stats)
+
+ def _get_descendants_in_topo_order(self, node_id: str) -> List[str]:
+ reachable = set()
+ queue = deque(self.children.get(node_id, set()))
+ while queue:
+ cur = queue.popleft()
+ if cur in reachable:
+ continue
+ reachable.add(cur)
+ for child in self.children.get(cur, set()):
+ queue.append(child)
+
+ return [nid for nid in self.topo_order if nid in reachable]
+
+ def _is_success(self) -> bool:
+ return (
+ len(self.running_nodes) == 0
+ and self.ready_nodes == []
+ and len(self._get_pending_nodes()) == 0
+ )
+
+ def _is_deadlock(self) -> bool:
+ pending_nodes = self._get_pending_nodes()
+ blocked = len(pending_nodes) > 0
+ # if len(self.running_nodes) == 0 and len(self.ready_nodes) == 0 and blocked:
+ # return True
+
+ for node_id in pending_nodes:
+ if node_id not in self.blocked_memory_nodes_perm:
+ blocked = False
+ break
+
+ return blocked
+
+ def not_enough_memory(self) -> bool:
+ for node_id in self._get_pending_nodes():
+ cpu_mem, gpu_mem = self.node_resources[node_id]
+ if cpu_mem > self.memory_budget["cpu"] - self.memory_stats["cpu_in_use"]:
+ return True
+ if gpu_mem > 0.0 and self.n_gpu > 0:
+ gpu_id = self._gpu_with_most_free_memory(gpu_mem)
+ if gpu_id is None:
+ return True
+ return self.memory_stats["cpu_in_use"] > self.memory_budget["cpu"]
+
+ def _check_memory_constraints(self, node_id: str) -> bool:
+ cpu_mem, gpu_mem = self.node_resources[node_id]
+ gpu_id = None
+ if cpu_mem > self.memory_budget["cpu"] - self.memory_stats["cpu_in_use"]:
+ if cpu_mem > self.memory_budget["cpu"]:
+ self.blocked_memory_nodes_perm.append(node_id)
+ self.topo_order.remove(node_id)
+ return False, None
+
+ if gpu_mem > 0.0 and self.n_gpu > 0:
+ gpu_id = self._gpu_with_most_free_memory(gpu_mem)
+
+ if gpu_id is None:
+ if DEBUG:
+ print(f"Node {node_id} has no available GPU")
+ return False, None
+
+ return True, gpu_id
+
+ def _gpu_with_most_free_memory(self, memory_needed):
+ free_memory = []
+ for i in range(self.n_gpu):
+ free_memory.append(
+ self.memory_budget["gpu"][i] - self.memory_stats["gpu_in_use"][i]
+ )
+
+ if max(free_memory) < memory_needed:
+ return None
+
+ return free_memory.index(max(free_memory))
+
+ def _get_pending_nodes(self) -> List[str]:
+ return [
+ node_id
+ for node_id in self.topo_order
+ if node_id not in self.leaves and self.unresolved_parents[node_id] == 0
+ ]
+
+ def _reserve_memory(self, node_id: str, gpu_id: int) -> bool:
+ cpu_mem, gpu_mem = self.node_resources[node_id]
+ self.memory_stats["cpu_in_use"] += cpu_mem
+ if gpu_id is not None:
+ self.memory_stats["gpu_in_use"][gpu_id] += gpu_mem
+
+ def _release_memory(self, node_id: str, gpu_id: int) -> bool:
+ cpu_mem, gpu_mem = self.node_resources[node_id]
+ self.memory_stats["cpu_in_use"] -= cpu_mem
+ if gpu_id is not None:
+ self.memory_stats["gpu_in_use"][gpu_id] -= gpu_mem
+
+ def _get_nodes_from_dags(
+ self, dags: List[RepresentationDag]
+ ) -> List[RepresentationNode]:
+ nodes = []
+ for dag in dags:
+ self.roots.add(dag.root_node_id)
+ for node in dag.nodes:
+ self.mapping[node.node_id] = node
+ if not node.inputs:
+ self.leaves.add(node.node_id)
+ for parent_id in node.inputs:
+ self.parents[node.node_id].add(parent_id)
+ self.children[parent_id].add(node.node_id)
+
+ indegree = {
+ node_id: len(self.parents.get(node_id, set())) for node_id in self.mapping
+ }
+ queue = deque([node_id for node_id, deg in indegree.items() if deg == 0])
+ topo_order: List[str] = []
+ while queue:
+ node_id = queue.popleft()
+ topo_order.append(node_id)
+ for child_id in self.children.get(node_id, set()):
+ indegree[child_id] -= 1
+ if indegree[child_id] == 0:
+ queue.append(child_id)
+ if len(topo_order) != len(self.mapping):
+ raise ValueError("Node graph contains a cycle")
+ self.topo_order = topo_order
+ self.parent_refcounts = {
+ node_id: len(self.children.get(node_id, set())) for node_id in self.mapping
+ }
+
+ self.remaining_children = {
+ node_id: len(self.children.get(node_id, set())) for node_id in self.mapping
+ }
+ return nodes
+
+ def _get_unresolved_parents(self) -> Dict[str, int]:
+ unresolved_parents = {}
+ for node_id in self.mapping:
+ unresolved_parents[node_id] = sum(
+ 1
+ for parent_id in self.parents.get(node_id, set())
+ if parent_id not in self.leaves
+ )
+ return unresolved_parents
+
+ def _get_modality_from_id(self, modality_id: str) -> Modality:
+ for modality in self.modalities:
+ if modality.modality_id == modality_id:
+ return modality
+ return None
+
+ def _estimate_node_resources(self):
+ node_resources = {}
+ node_stats = {}
+ for node in self.topo_order:
+ if node in self.leaves:
+ modality = self._get_modality_from_id(self.mapping[node].modality_id)
+ peak_memory = modality.estimate_peak_memory_bytes()
+ node_resources[node] = (
+ peak_memory["cpu_peak_bytes"],
+ peak_memory["gpu_peak_bytes"],
+ )
+ node_stats[node] = modality.get_stats()
+ else:
+ parent_ids = list(self.parents.get(node, set()))
+ parent_stats = [node_stats[parent_id] for parent_id in parent_ids]
+ input_stats = parent_stats[0] if parent_stats else None
+ if node not in self.roots:
+ operation = self.mapping[node].operation(
+ params=self.mapping[node].parameters
+ )
+ peak_memory = operation.estimate_peak_memory_bytes(input_stats)
+ peak_memory["cpu_peak_bytes"] += (
+ 64 * 1024 + 512 * input_stats.num_instances
+ ) # Placeholder for transformed modality creation overhead
+ peak_memory["cpu_peak_bytes"] *= 1
+ output_stats = operation.get_output_stats(input_stats)
+ node_resources[node] = (
+ int(peak_memory["cpu_peak_bytes"]),
+ int(peak_memory["gpu_peak_bytes"]),
+ )
+ node_stats[node] = output_stats
+ else:
+ task_node = self.mapping[node]
+ task = self.tasks[task_node.parameters["_task_idx"]]
+ peak_memory = task.estimate_peak_memory_bytes(input_stats)
+
+ cpu_increment = max(
+ int(peak_memory["cpu_peak_bytes"]),
+ 16 * 1024 * 1024,
+ )
+ node_resources[node] = (
+ cpu_increment,
+ int(peak_memory["gpu_peak_bytes"]),
+ )
+
+ node_stats[node] = task.get_output_stats(input_stats)
+ self.node_stats = node_stats
+ return node_resources
+
+ @staticmethod
+ def _stats_to_bytes(stats: Optional[Any], dtype_size: int = 4) -> int:
+ if stats is None:
+ return 0
+ num_instances = int(getattr(stats, "num_instances", 0))
+ output_shape = tuple(getattr(stats, "output_shape", ()))
+ numel = 1
+ for dim in output_shape:
+ try:
+ numel *= int(dim)
+ except Exception:
+ numel *= 1
+ return max(0, int(num_instances * numel * dtype_size))
diff --git a/src/main/python/systemds/scuro/drsearch/operator_registry.py b/src/main/python/systemds/scuro/drsearch/operator_registry.py
index e9c302ba90..bc9bd406a8 100644
--- a/src/main/python/systemds/scuro/drsearch/operator_registry.py
+++ b/src/main/python/systemds/scuro/drsearch/operator_registry.py
@@ -35,8 +35,7 @@ class Registry:
_representations = {}
_context_operators = {}
_fusion_operators = []
- _text_context_operators = []
- _video_context_operators = []
+ _context_representation_operators = {}
_dimensionality_reduction_operators = {}
def __new__(cls):
@@ -86,6 +85,18 @@ class Registry:
dimensionality_reduction_operator
)
+ def add_context_representation_operator(
+ self, context_representation, modality_type
+ ):
+ if not isinstance(modality_type, list):
+ modality_type = [modality_type]
+ for m_type in modality_type:
+ if not m_type in self._context_representation_operators.keys():
+ self._context_representation_operators[m_type] = []
+ self._context_representation_operators[m_type].append(
+ context_representation
+ )
+
def get_representations(self, modality: ModalityType):
return self._representations[modality]
@@ -102,7 +113,7 @@ class Registry:
return self._context_operators[modality_type]
def get_dimensionality_reduction_operators(self, modality_type):
- return self._dimensionality_reduction_operators[modality_type]
+ return self._dimensionality_reduction_operators.get(modality_type, [])
def get_fusion_operators(self):
return self._fusion_operators
@@ -125,6 +136,9 @@ class Registry:
return None, False
+ def get_context_representations(self, modality_type):
+ return self._context_representation_operators[modality_type]
+
def register_representation(modalities: Union[ModalityType, List[ModalityType]]):
"""
@@ -157,6 +171,18 @@ def register_dimensionality_reduction_operator(modality_type):
return decorator
+def register_context_representation_operator(modality_type):
+ """
+ Decorator to register a context representation operator.
+ """
+
+ def decorator(cls):
+ Registry().add_context_representation_operator(cls, modality_type)
+ return cls
+
+ return decorator
+
+
def register_context_operator(modality_type):
"""
Decorator to register a context operator.
diff --git a/src/main/python/systemds/scuro/drsearch/ranking.py b/src/main/python/systemds/scuro/drsearch/ranking.py
index b4b8a392ea..381a6d1fb9 100644
--- a/src/main/python/systemds/scuro/drsearch/ranking.py
+++ b/src/main/python/systemds/scuro/drsearch/ranking.py
@@ -33,7 +33,7 @@ def rank_by_tradeoff(
):
entries = list(entries)
if not entries:
- return []
+ return [], []
performance_score_accessor = lambda entry: getattr(entry, "val_score")[
performance_metric_name
diff --git a/src/main/python/systemds/scuro/drsearch/representation_dag.py b/src/main/python/systemds/scuro/drsearch/representation_dag.py
index 01020546a0..b2911dbe7a 100644
--- a/src/main/python/systemds/scuro/drsearch/representation_dag.py
+++ b/src/main/python/systemds/scuro/drsearch/representation_dag.py
@@ -20,7 +20,7 @@
# -------------------------------------------------------------
import copy
from dataclasses import dataclass, field
-from typing import List, Dict, Union, Any, Hashable, Optional
+from typing import List, Dict, Set, Tuple, Union, Any, Hashable, Optional
from systemds.scuro.modality.modality import Modality
from systemds.scuro.modality.transformed import TransformedModality
from systemds.scuro.representations.representation import (
@@ -34,7 +34,7 @@ from systemds.scuro.representations.dimensionality_reduction import (
DimensionalityReduction,
)
from systemds.scuro.utils.identifier import get_op_id, get_node_id
-from collections import OrderedDict
+from collections import OrderedDict, defaultdict, deque
class LRUCache:
@@ -68,6 +68,7 @@ class RepresentationNode:
modality_id: str = None
representation_index: int = None
parameters: Dict[str, Any] = field(default_factory=dict)
+ gpu_id: int = None
@dataclass
@@ -118,6 +119,22 @@ class RepresentationDag:
children.append(node.node_id)
return children
+ def get_first_level_nodes(self) -> List[str]:
+ leaf_nodes = set(self.get_leaf_nodes())
+ first_level_nodes = []
+
+ for node in self.nodes:
+ if node.node_id in leaf_nodes:
+ continue
+
+ if node.inputs and all(input_id in leaf_nodes for input_id in node.inputs):
+ first_level_nodes.append(node.node_id)
+
+ return first_level_nodes
+
+ def get_first_level_node_set(self) -> set:
+ return set(self.get_first_level_nodes())
+
def validate(self) -> bool:
node_ids = {node.node_id for node in self.nodes}
@@ -146,6 +163,65 @@ class RepresentationDag:
return not has_cycle(self.root_node_id, set())
+ def _get_node_representation_label(self, node: RepresentationNode) -> str:
+ if not node.operation:
+ if node.representation_index is None:
+ return f"modality={node.modality_id}"
+ return (
+ f"modality={node.modality_id}, "
+ f"representation_index={node.representation_index}"
+ )
+
+ try:
+ operation_instance = node.operation(params=node.parameters)
+ if hasattr(operation_instance, "name"):
+ return str(operation_instance.name)
+ except Exception:
+ pass
+
+ if hasattr(node.operation, "__name__"):
+ return str(node.operation.__name__)
+ return str(node.operation)
+
+ def to_graphviz(self, graph_name: str = "RepresentationDag", rankdir: str = "BT"):
+ try:
+ from graphviz import Digraph
+ except ImportError as exc:
+ raise ImportError(
+ "Graphviz visualization requires the 'graphviz' Python package. "
+ "Install it with 'pip install graphviz' and ensure Graphviz "
+ "binaries are available on your PATH."
+ ) from exc
+
+ graph = Digraph(name=graph_name)
+ graph.attr(rankdir=rankdir)
+
+ for node in self.nodes:
+ node_label = f"{node.node_id}\\n{self._get_node_representation_label(node)}"
+ node_shape = "ellipse" if not node.inputs else "box"
+ node_attributes = {"shape": node_shape}
+ if node.node_id == self.root_node_id:
+ node_attributes["style"] = "bold"
+ graph.node(node.node_id, label=node_label, **node_attributes)
+
+ for node in self.nodes:
+ for input_id in node.inputs:
+ graph.edge(input_id, node.node_id)
+
+ return graph
+
+ def render_graphviz(
+ self,
+ output_path: str,
+ format: str = "png",
+ view: bool = False,
+ cleanup: bool = True,
+ ) -> str:
+ graph = self.to_graphviz()
+ return graph.render(
+ filename=output_path, format=format, view=view, cleanup=cleanup
+ )
+
def _compute_leaf_signature(self, node) -> Hashable:
return ("leaf", node.modality_id, node.representation_index)
@@ -161,40 +237,70 @@ class RepresentationDag:
external_cache: Optional[LRUCache] = None,
enable_cache=True,
rep_cache: Dict[Any, TransformedModality] = None,
+ consumer_count: Dict[str, int] = None,
+ gpu_id: int = None,
) -> Union[Dict[str, TransformedModality], TransformedModality]:
- cache: Dict[str, TransformedModality] = {}
- node_signatures: Dict[str, Hashable] = {}
def execute_node(node_id: str, task) -> TransformedModality:
- if node_id in cache:
- return cache[node_id]
+ if external_cache is not None:
+ cached = external_cache.get(node_id)
+ if cached is not None:
+ return cached
node = self.get_node_by_id(node_id)
if not node.inputs:
modality = get_modality_by_id_and_instance_id(
- modalities, node.modality_id, node.representation_index
+ modalities, int(node.modality_id), node.representation_index
)
- if enable_cache:
- cache[node_id] = modality
- node_signatures[node_id] = self._compute_leaf_signature(node)
+
return modality
input_mods = [execute_node(input_id, task) for input_id in node.inputs]
- input_signatures = tuple(
- node_signatures[input_id] for input_id in node.inputs
- )
- node_signature = self._compute_node_signature(node, input_signatures)
is_unimodal = len(input_mods) == 1
- cached_result = None
- if external_cache and is_unimodal:
- cached_result = external_cache.get(node_signature)
- if cached_result is not None:
- result = cached_result
-
+ if external_cache is not None and is_unimodal:
+ cached = external_cache.get(node_id)
+ if cached is not None:
+ result = cached
+ else:
+ node_operation = node.operation(params=node.parameters)
+ if gpu_id is not None and hasattr(node_operation, "gpu_id"):
+ node_operation.gpu_id = gpu_id
+ if len(input_mods) == 1:
+ # It's a unimodal operation
+ if isinstance(node_operation, Context):
+ result = input_mods[0].context(node_operation)
+ elif isinstance(node_operation, DimensionalityReduction):
+ result = input_mods[0].dimensionality_reduction(
+ node_operation
+ )
+ elif isinstance(node_operation, AggregatedRepresentation):
+ result = node_operation.transform(input_mods[0])
+ elif isinstance(node_operation, UnimodalRepresentation):
+ if rep_cache is not None:
+ result = rep_cache[node_operation.name]
+ else:
+ # Compute the representation
+ result = input_mods[0].apply_representation(
+ node_operation
+ )
+ else:
+ # It's a fusion operation
+ fusion_op = node_operation
+ if (
+ hasattr(fusion_op, "needs_training")
+ and fusion_op.needs_training
+ ):
+ result = input_mods[0].combine_with_training(
+ input_mods[1:], fusion_op, task
+ )
+ else:
+ result = input_mods[0].combine(input_mods[1:], fusion_op)
else:
- node_operation = copy.deepcopy(node.operation())
+ node_operation = node.operation(params=node.parameters)
+ if gpu_id is not None and hasattr(node_operation, "gpu_id"):
+ node_operation.gpu_id = gpu_id
if len(input_mods) == 1:
# It's a unimodal operation
if isinstance(node_operation, Context):
@@ -206,13 +312,6 @@ class RepresentationDag:
elif isinstance(node_operation, UnimodalRepresentation):
if rep_cache is not None:
result = rep_cache[node_operation.name]
- elif (
- isinstance(input_mods[0], TransformedModality)
- and input_mods[0].transformation[0].__class__
- == node.operation
- ):
- # Avoid duplicate transformations
- result = input_mods[0]
else:
# Compute the representation
result = input_mods[0].apply_representation(node_operation)
@@ -228,17 +327,59 @@ class RepresentationDag:
)
else:
result = input_mods[0].combine(input_mods[1:], fusion_op)
- if external_cache and is_unimodal:
- external_cache.put(node_signature, result)
- if enable_cache:
- cache[node_id] = result
- node_signatures[node_id] = node_signature
+ if (
+ enable_cache
+ and external_cache is not None
+ and is_unimodal
+ and consumer_count[node_id] > 1
+ ):
+ external_cache.put(node_id, result)
return result
result = execute_node(self.root_node_id, task)
- return cache if enable_cache else result
+ return result
+
+ def compute_full_node_signature(
+ self, node_id: str, node_signatures: Dict[str, Hashable] = None
+ ) -> Hashable:
+ if node_signatures is None:
+ node_signatures = {}
+
+ if node_id in node_signatures:
+ return node_signatures[node_id]
+
+ node = self.get_node_by_id(node_id)
+ if not node:
+ return None
+
+ if not node.inputs:
+ sig = self._compute_leaf_signature(node)
+ node_signatures[node_id] = sig
+ return sig
+
+ input_sigs = []
+ for input_id in sorted(node.inputs):
+ input_sig = self.compute_full_node_signature(input_id, node_signatures)
+ input_sigs.append(input_sig)
+
+ sig = self._compute_node_signature(node, tuple(input_sigs))
+ node_signatures[node_id] = sig
+ return sig
+
+ def find_nodes_with_same_predecessors(
+ self, target_node_id: str, other_dag: "RepresentationDag", other_node_id: str
+ ) -> bool:
+ target_sig = self.compute_full_node_signature(target_node_id)
+ other_sig = other_dag.compute_full_node_signature(other_node_id)
+ return target_sig == other_sig
+
+ def get_leaf_node_id(self) -> str:
+ for node in self.nodes:
+ if not node.inputs:
+ return node.node_id
+ return None
def get_modality_by_id_and_instance_id(
@@ -303,3 +444,287 @@ class RepresentationDAGBuilder:
if node.node_id == node_id:
return node
return None
+
+
+def get_consumer_count(dags: List[RepresentationDag]) -> Dict[str, int]:
+ consumer_count: Dict[str, int] = defaultdict(int)
+ for dag in dags:
+ for node in dag.nodes:
+ for inp in node.inputs:
+ consumer_count[inp] += 1
+ return consumer_count
+
+
+def pushdown_aggregation(dag_group: List[RepresentationDag]) -> List[RepresentationDag]:
+ consumer_count: Dict[str, int] = defaultdict(int)
+
+ for dag in dag_group:
+ for node in dag.nodes:
+ for inp in node.inputs:
+ consumer_count[inp] += 1
+
+ processed_agg_ids: Set[str] = set()
+
+ for dag in dag_group:
+ agg_nodes = [
+ n
+ for n in dag.nodes
+ if n.operation and issubclass(n.operation, AggregatedRepresentation)
+ ]
+ for agg_node in agg_nodes:
+ if agg_node.node_id in processed_agg_ids:
+ continue
+ processed_agg_ids.add(agg_node.node_id)
+
+ if len(agg_node.inputs) != 1:
+ print(
+ f"Aggregation node {agg_node.node_id} has {len(agg_node.inputs)} inputs, skipping (SHOULD NOT HAPPEN)"
+ )
+ continue
+
+ input_id = agg_node.inputs[0]
+
+ processed_agg_ids.add(input_id)
+ if consumer_count[input_id] != 1:
+ continue
+
+ input_node = None
+ for d in dag_group:
+ input_node = d.get_node_by_id(input_id)
+ if input_node is not None:
+ break
+
+ if not input_node or not input_node.operation:
+ continue
+
+ op_instance = input_node.operation(params=input_node.parameters)
+ if op_instance.__class__.__bases__[0].__name__ != "BertFamily":
+ continue
+
+ input_node.parameters["_pushdown_aggregation"] = agg_node.parameters
+
+ for d in dag_group:
+ for node in d.nodes:
+ node.inputs = [
+ input_id if inp == agg_node.node_id else inp
+ for inp in node.inputs
+ ]
+
+ if d.root_node_id == agg_node.node_id:
+ d.root_node_id = input_id
+
+ d.nodes = [n for n in d.nodes if n.node_id != agg_node.node_id]
+
+ return dag_group
+
+
+class CSEAwareDAGBuilder:
+ def __init__(self):
+ self.global_nodes: List[RepresentationNode] = []
+ self.signature_to_node: Dict[Hashable, str] = {}
+ self.node_to_signature: Dict[str, Hashable] = {}
+ self.node_counter = 0
+
+ def _compute_node_signature(
+ self, operation: Any, inputs: List[str], parameters: Dict[str, Any] = None
+ ) -> Hashable:
+ ip = [self.node_to_signature[inp] for inp in inputs]
+ input_sigs = tuple(sorted(ip)) if inputs else ()
+ op_cls = operation().name
+ params_items = tuple(sorted((parameters or {}).items()))
+ return ("op", op_cls, params_items, input_sigs)
+
+ def _compute_leaf_signature(
+ self, modality_id: str, representation_index: int = -1
+ ) -> Hashable:
+ return ("leaf", modality_id, representation_index)
+
+ def _get_or_create_node(
+ self,
+ operation: Any,
+ inputs: List[str],
+ modality_id: str = None,
+ representation_index: int = None,
+ parameters: Dict[str, Any] = None,
+ is_leaf: bool = False,
+ ) -> str:
+ if is_leaf:
+ signature = self._compute_leaf_signature(modality_id, representation_index)
+ else:
+ signature = self._compute_node_signature(operation, inputs, parameters)
+
+ try:
+ if signature in self.signature_to_node:
+ return self.signature_to_node[signature]
+ except:
+ pass
+
+ if is_leaf:
+ if representation_index != -1:
+ node_id = f"leaf_{modality_id}_{representation_index}"
+ else:
+ node_id = f"leaf_{get_node_id()}"
+ else:
+ node_id = f"op_{get_op_id()}"
+ self.node_counter += 1
+
+ node = RepresentationNode(
+ node_id=node_id,
+ inputs=inputs,
+ operation=operation,
+ modality_id=modality_id,
+ representation_index=representation_index,
+ parameters=parameters or {},
+ gpu_id=None,
+ )
+
+ self.global_nodes.append(node)
+ self.signature_to_node[signature] = node_id
+ self.node_to_signature[node_id] = signature
+
+ return node_id
+
+ def create_leaf_node(
+ self, modality_id: str, representation_index: int = -1, operation=None
+ ) -> str:
+ return self._get_or_create_node(
+ operation=operation,
+ inputs=[],
+ modality_id=modality_id,
+ representation_index=representation_index,
+ is_leaf=True,
+ )
+
+ def create_operation_node(
+ self, operation: Any, inputs: List[str], parameters: Dict[str, Any] = None
+ ):
+ return self._get_or_create_node(
+ operation=operation, inputs=inputs, parameters=parameters, is_leaf=False
+ )
+
+ def build(self, root_node_id: str) -> RepresentationDag:
+ dag = RepresentationDag(nodes=self.global_nodes, root_node_id=root_node_id)
+ if not dag.validate():
+ raise ValueError("Invalid DAG construction")
+ return dag
+
+ def get_node(self, node_id: str) -> Optional[RepresentationNode]:
+ for node in self.global_nodes:
+ if node.node_id == node_id:
+ return node
+ return None
+
+
+def group_dags_by_dependencies(
+ dags: List[RepresentationDag],
+) -> List[List[RepresentationDag]]:
+ if not dags:
+ return []
+
+ unique_dags: List[RepresentationDag] = []
+ seen_signatures: set[Hashable] = set()
+
+ for dag in dags:
+ dag_sig = dag.compute_full_node_signature(dag.root_node_id)
+ if dag_sig not in seen_signatures:
+ seen_signatures.add(dag_sig)
+ unique_dags.append(dag)
+
+ dags = unique_dags
+ dag_first_level_sets = []
+ for dag in dags:
+ first_level_nodes = dag.get_first_level_node_set()
+ dag_first_level_sets.append(first_level_nodes)
+
+ n = len(dags)
+ parent = list(range(n))
+
+ def find(x):
+ if parent[x] != x:
+ parent[x] = find(parent[x])
+ return parent[x]
+
+ def union(x, y):
+ root_x = find(x)
+ root_y = find(y)
+ if root_x != root_y:
+ parent[root_y] = root_x
+
+ for i in range(n):
+ for j in range(i + 1, n):
+ if dag_first_level_sets[i] & dag_first_level_sets[j]:
+ union(i, j)
+
+ groups: Dict[int, List[RepresentationDag]] = {}
+ for i in range(n):
+ root = find(i)
+ if root not in groups:
+ groups[root] = []
+ groups[root].append(dags[i])
+
+ return list(groups.values())
+
+
+def dags_to_graphviz(
+ dags: List[RepresentationDag],
+ graph_name: str = "RepresentationDagGroup",
+ rankdir: str = "BT",
+ show_dag_roots: bool = True,
+):
+ try:
+ from graphviz import Digraph
+ except ImportError as exc:
+ raise ImportError(
+ "Graphviz visualization requires the 'graphviz' Python package. "
+ "Install it with 'pip install graphviz' and ensure Graphviz "
+ "binaries are available on your PATH."
+ ) from exc
+
+ graph = Digraph(name=graph_name)
+ graph.attr(rankdir=rankdir)
+
+ root_ids = {dag.root_node_id for dag in dags}
+ rendered_nodes: Set[str] = set()
+ rendered_edges: Set[Tuple[str, str]] = set()
+
+ for dag_idx, dag in enumerate(dags):
+ if show_dag_roots:
+ dag_node_id = f"dag_{dag_idx}"
+ graph.node(
+ dag_node_id,
+ label=f"DAG {dag_idx}\\nroot={dag.root_node_id}",
+ shape="oval",
+ style="dashed",
+ )
+ graph.edge(dag_node_id, dag.root_node_id, style="dashed")
+
+ for node in dag.nodes:
+ if node.node_id not in rendered_nodes:
+ node_label = (
+ f"{node.node_id}\\n{dag._get_node_representation_label(node)}"
+ )
+ node_shape = "ellipse" if not node.inputs else "box"
+ node_attributes = {"shape": node_shape}
+ if node.node_id in root_ids:
+ node_attributes["style"] = "bold"
+ graph.node(node.node_id, label=node_label, **node_attributes)
+ rendered_nodes.add(node.node_id)
+
+ for input_id in node.inputs:
+ edge_key = (input_id, node.node_id)
+ if edge_key not in rendered_edges:
+ graph.edge(input_id, node.node_id)
+ rendered_edges.add(edge_key)
+
+ return graph
+
+
+def render_dags_graphviz(
+ dags: List[RepresentationDag],
+ output_path: str,
+ format: str = "png",
+ view: bool = False,
+ cleanup: bool = True,
+) -> str:
+ graph = dags_to_graphviz(dags)
+ return graph.render(filename=output_path, format=format, view=view, cleanup=cleanup)
diff --git a/src/main/python/systemds/scuro/drsearch/representation_dag_visualizer.py b/src/main/python/systemds/scuro/drsearch/representation_dag_visualizer.py
index b0a1bbe285..42b1d8defa 100644
--- a/src/main/python/systemds/scuro/drsearch/representation_dag_visualizer.py
+++ b/src/main/python/systemds/scuro/drsearch/representation_dag_visualizer.py
@@ -20,36 +20,22 @@
# -------------------------------------------------------------
from typing import Dict, Any
-from systemds.scuro.drsearch.representation_dag import RepresentationDag
+from typing import List
+from systemds.scuro.drsearch.representation_dag import (
+ RepresentationDag,
+ dags_to_graphviz,
+)
def visualize_dag(dag: RepresentationDag) -> Dict[str, Any]:
- nodes = []
- edges = []
+ graph = dag.to_graphviz()
+ return {"root": dag.root_node_id, "dot": graph.source, "graph": graph}
- for i, node in enumerate(dag.nodes):
- # Create node entry
- node_type = "operation" if node.operation else "modality"
- label = node.operation if node.operation else f"Modality: {node.modality_id}"
- nodes.append(
- {
- "id": node.node_id,
- "label": label,
- "type": node_type,
- "parameters": node.parameters,
- }
- )
-
- print(nodes[i])
-
- # Create edges
- for input_id in node.inputs:
- edges.append({"from": input_id, "to": node.node_id})
-
- for edge in edges:
- print(edge)
-
- print(f"Root Node ID: {dag.root_node_id}")
-
- return {"nodes": nodes, "edges": edges, "root": dag.root_node_id}
+def visualize_dag_group(dags: List[RepresentationDag]) -> Dict[str, Any]:
+ graph = dags_to_graphviz(dags)
+ return {
+ "roots": [dag.root_node_id for dag in dags],
+ "dot": graph.source,
+ "graph": graph,
+ }
diff --git a/src/main/python/systemds/scuro/drsearch/task.py b/src/main/python/systemds/scuro/drsearch/task.py
index fbe08bcc61..404624ec3a 100644
--- a/src/main/python/systemds/scuro/drsearch/task.py
+++ b/src/main/python/systemds/scuro/drsearch/task.py
@@ -24,6 +24,7 @@ from typing import List
from systemds.scuro.models.model import Model
import numpy as np
from sklearn.model_selection import train_test_split
+from systemds.scuro.representations.representation import RepresentationStats
class PerformanceMeasure:
@@ -97,6 +98,85 @@ class Task:
self.fusion_train_indices = None
self._create_cv_splits()
+ def get_output_stats(self, input_stats):
+ # TODO: Implement a default estimate of the output stats for the task
+ return RepresentationStats(0, (0,))
+
+ def estimate_peak_memory_bytes(self, input_stats):
+ labels_array = np.asarray(self.labels)
+ n_instances = int(input_stats.num_instances)
+ feature_dim = (
+ int(np.prod(input_stats.output_shape)) if input_stats.output_shape else 0
+ )
+ feature_dtype_bytes = 4
+
+ n_train = len(self.train_indices) if self.train_indices is not None else 0
+ n_test = len(self.test_indices) if self.test_indices is not None else 0
+ n_labels = int(labels_array.shape[1]) if labels_array.ndim > 1 else 1
+
+ max_fold_train = max((len(fold) for fold in self.cv_train_indices), default=0)
+ max_fold_val = max((len(fold) for fold in self.cv_val_indices), default=0)
+
+ base_input_bytes = n_instances * feature_dim * feature_dtype_bytes
+ test_slice_bytes = n_test * feature_dim * feature_dtype_bytes
+ fold_slice_bytes = (
+ (max_fold_train + max_fold_val) * feature_dim * feature_dtype_bytes
+ )
+ label_slice_bytes = (
+ (max_fold_train + max_fold_val + n_test) * n_labels * feature_dtype_bytes
+ )
+ label_storage_bytes = int(labels_array.nbytes) * 2
+
+ total_index_entries = (
+ n_train
+ + n_test
+ + sum(len(fold) for fold in self.cv_train_indices)
+ + sum(len(fold) for fold in self.cv_val_indices)
+ + (
+ len(self.fusion_train_indices)
+ if self.fusion_train_indices is not None
+ else 0
+ )
+ )
+ index_bytes = total_index_entries * np.dtype(np.int64).itemsize
+
+ scheduler_side_cpu_bytes = (
+ base_input_bytes
+ + test_slice_bytes * 2
+ + fold_slice_bytes * 2
+ + label_slice_bytes * 2
+ + label_storage_bytes
+ + index_bytes
+ + (192 * 1024 * 1024)
+ )
+
+ model_peak_memory_cpu = 0.0
+ model_peak_memory_gpu = 0.0
+ if hasattr(self.model, "estimate_peak_memory_bytes"):
+ try:
+ model_peak_memory_cpu, model_peak_memory_gpu = (
+ self.model.estimate_peak_memory_bytes(
+ input_dim=feature_dim,
+ n_train_samples=(
+ max_fold_train if max_fold_train > 0 else n_train
+ ),
+ n_labels=n_labels,
+ n_val_samples=max_fold_val,
+ n_test_samples=n_test,
+ )
+ )
+ except TypeError:
+ model_peak_memory_cpu, model_peak_memory_gpu = (
+ self.model.estimate_peak_memory_bytes(feature_dim, n_train)
+ )
+
+ return {
+ "cpu_peak_bytes": int(
+ model_peak_memory_cpu * 1.5 + scheduler_side_cpu_bytes * 2
+ ),
+ "gpu_peak_bytes": int(model_peak_memory_gpu) * 1.4,
+ }
+
def _create_cv_splits(self):
train_labels = [self.labels[i] for i in self.train_indices]
train_labels_array = np.array(train_labels)
@@ -166,17 +246,17 @@ class Task:
self._reset_params()
model = self.create_model()
- test_X = np.array([data[i] for i in self.test_indices])
- test_y = np.array([self.labels[i] for i in self.test_indices])
+ test_X = self._gather_by_indices(data, self.test_indices)
+ test_y = self._gather_by_indices(self.labels, self.test_indices)
for fold_idx in range(self.kfold):
fold_train_indices = self.cv_train_indices[fold_idx]
fold_val_indices = self.cv_val_indices[fold_idx]
- train_X = np.array([data[i] for i in fold_train_indices])
- train_y = np.array([self.labels[i] for i in fold_train_indices])
- val_X = np.array([data[i] for i in fold_val_indices])
- val_y = np.array([self.labels[i] for i in fold_val_indices])
+ train_X = self._gather_by_indices(data, fold_train_indices)
+ train_y = self._gather_by_indices(self.labels, fold_train_indices)
+ val_X = self._gather_by_indices(data, fold_val_indices)
+ val_y = self._gather_by_indices(self.labels, fold_val_indices)
self._run_fold(model, train_X, train_y, val_X, val_y, test_X, test_y)
@@ -201,8 +281,12 @@ class Task:
self.train_scores.add_scores(train_score[0])
val_score = model.test(val_X, val_y)
test_start = time.time()
- test_score = model.test(np.array(test_X), test_y)
+ test_score = model.test(np.asarray(test_X), test_y)
test_end = time.time()
self.inference_time.append(test_end - test_start)
self.val_scores.add_scores(val_score[0])
self.test_scores.add_scores(test_score[0])
+
+ @staticmethod
+ def _gather_by_indices(values, indices):
+ return [values[i] for i in indices]
diff --git a/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py b/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py
index 5b03147ec1..f733dc61ae 100644
--- a/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py
+++ b/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py
@@ -18,17 +18,19 @@
# under the License.
#
# -------------------------------------------------------------
+import copy
import pickle
import time
from concurrent.futures import ProcessPoolExecutor, as_completed
from dataclasses import dataclass
import multiprocessing as mp
-from typing import List, Any
+from typing import List, Any, Optional, Dict
from functools import lru_cache
+
from systemds.scuro import ModalityType
+from systemds.scuro.drsearch.node_executor import NodeExecutor, ResultEntry
from systemds.scuro.drsearch.ranking import rank_by_tradeoff
from systemds.scuro.drsearch.task import PerformanceMeasure
-from systemds.scuro.representations.fusion import Fusion
from systemds.scuro.representations.concatenation import Concatenation
from systemds.scuro.representations.hadamard import Hadamard
from systemds.scuro.representations.sum import Sum
@@ -36,16 +38,15 @@ from systemds.scuro.representations.aggregated_representation import (
AggregatedRepresentation,
)
from systemds.scuro.modality.modality import Modality
-from systemds.scuro.representations.aggregate import Aggregation
from systemds.scuro.drsearch.operator_registry import Registry
-from systemds.scuro.utils.schema_helpers import get_shape
+from systemds.scuro.utils.checkpointing import CheckpointManager
from systemds.scuro.drsearch.representation_dag import (
RepresentationDag,
RepresentationNode,
- RepresentationDAGBuilder,
+ CSEAwareDAGBuilder,
+ pushdown_aggregation,
)
-from systemds.scuro.drsearch.representation_dag_visualizer import visualize_dag
-from systemds.scuro.drsearch.representation_dag import LRUCache
+from bisect import bisect_left
class UnimodalOptimizer:
@@ -58,6 +59,9 @@ class UnimodalOptimizer:
result_path=None,
k=2,
metric_name="accuracy",
+ checkpoint_every: Optional[int] = 1,
+ resume: bool = False,
+ max_num_workers: int = -1,
):
self.modalities = modalities
self.tasks = tasks
@@ -66,15 +70,25 @@ class UnimodalOptimizer:
self.result_path = result_path
self.k = k
self.metric_name = metric_name
+ self.checkpoint_every = checkpoint_every
+ self.resume = resume
+ self.max_num_workers = max_num_workers
+ self._checkpoint_manager = CheckpointManager(
+ self.result_path or ".",
+ "unimodal_checkpoint_",
+ checkpoint_every=self.checkpoint_every,
+ resume=self.resume,
+ )
+ # TODO: check if we should make this a local variable (might keep unnecessary memory usage)
self.builders = {
- modality.modality_id: RepresentationDAGBuilder() for modality in modalities
+ modality.modality_id: CSEAwareDAGBuilder() for modality in modalities
}
self.debug = debug
self.operator_registry = Registry()
self.operator_performance = UnimodalResults(
- modalities, tasks, debug, True, k, metric_name
+ modalities, tasks, debug, True, k, self.metric_name
)
self._tasks_require_same_dims = True
@@ -148,103 +162,201 @@ class UnimodalOptimizer:
task.model.name
].append(c)
+ def _count_results(self, results) -> int:
+ count = 0
+ for modality_id in results:
+ for task_name in results[modality_id]:
+ count += len(results[modality_id][task_name])
+ return count
+
+ def _count_results_by_modality(self, results) -> Dict[Any, int]:
+ counts = {}
+ for modality_id in results:
+ counts[modality_id] = len(
+ results[modality_id][list(results[modality_id].keys())[0]]
+ )
+
+ return counts
+
+ def resume_from_checkpoint(self):
+ loaded = self._checkpoint_manager.resume_from_checkpoint(
+ "eval_count_by_modality", self._count_results_by_modality
+ )
+ if loaded:
+ results, _, _ = loaded
+ self.operator_performance.results = results
+
def optimize_parallel(self, n_workers=None):
+ if self.resume:
+ self.resume_from_checkpoint()
+ # TODO: check which modalities have been processed and skip the ones that have been processed
+
if n_workers is None:
n_workers = min(len(self.modalities), mp.cpu_count())
- ctx = mp.get_context("spawn")
- with ProcessPoolExecutor(max_workers=n_workers, mp_context=ctx) as executor:
- future_to_modality = {
- executor.submit(self._process_modality, modality, True): modality
- for modality in self.modalities
- }
-
- for future in as_completed(future_to_modality):
- modality = future_to_modality[future]
- try:
- results = future.result()
- self._merge_results(results)
- except Exception as e:
- print(f"Error processing modality {modality.modality_id}: {e}")
- import traceback
-
- traceback.print_exc()
- continue
+
+ with mp.Manager() as manager:
+
+ ctx = mp.get_context("spawn")
+ with ProcessPoolExecutor(max_workers=n_workers, mp_context=ctx) as executor:
+ future_to_modality = {
+ executor.submit(
+ self._process_modality,
+ modality,
+ self._checkpoint_manager.skip_remaining_by_key.get(
+ modality.modality_id, 0
+ )
+ / len(self.tasks),
+ scheduler=None,
+ ): modality
+ for modality in self.modalities
+ }
+
+ for future in as_completed(future_to_modality):
+ modality = future_to_modality[future]
+ try:
+ results = future.result()
+ self._merge_results(results)
+ new_count = self._count_results(results.results)
+ self._checkpoint_manager.increment(
+ modality.modality_id, new_count
+ )
+ self._checkpoint_manager.checkpoint_if_due(
+ self.operator_performance.results,
+ )
+ except Exception as e:
+ print(f"Error processing modality {modality.modality_id}: {e}")
+ import traceback
+
+ traceback.print_exc()
+ self._checkpoint_manager.save_checkpoint(
+ self.operator_performance.results,
+ {},
+ )
+ continue
def optimize(self):
- """Optimize representations for each modality"""
+ if self.resume:
+ self.resume_from_checkpoint()
+
for modality in self.modalities:
- # try:
- local_result = self._process_modality(modality, False)
- if self.save_all_results:
- self.store_results(f"{modality.modality_id}_unimodal_results.pkl")
- # except Exception as e:
- # print(f"Error processing modality {modality.modality_id}: {e}")
- # if self.save_all_results:
- # self.store_results(f"{modality.modality_id}_unimodal_results.pkl")
- # continue
-
- def _process_modality(self, modality, parallel):
- if parallel:
- local_results = UnimodalResults(
- [modality], self.tasks, debug=False, store_cache=False
- )
- else:
- local_results = self.operator_performance
+ try:
+ local_result = self._process_modality(
+ modality,
+ (
+ int(
+ round(
+ self._checkpoint_manager.skip_remaining_by_key.get(
+ modality.modality_id, 0
+ )
+ / len(self.tasks)
+ )
+ )
+ if self.resume
+ else 0
+ ),
+ )
+ self._merge_results(local_result)
+ new_count = self._count_results(local_result.results)
+ self._checkpoint_manager.increment(modality.modality_id, new_count)
+ self._checkpoint_manager.checkpoint_if_due(
+ self.operator_performance.results
+ )
+ if self.save_all_results:
+ self.store_results(f"{modality.modality_id}_unimodal_results.pkl")
+ except Exception as e:
+ print(f"Error processing modality {modality.modality_id}: {e}")
+ import traceback
+
+ traceback.print_exc()
+ self._checkpoint_manager.save_checkpoint(
+ self.operator_performance.results, {}
+ )
+ raise
+
+ def _expand_dags_with_task_roots(
+ self, dags: List[RepresentationDag]
+ ) -> List[RepresentationDag]:
+ expanded_dags: List[RepresentationDag] = []
+
+ for dag in dags:
+ root_id = dag.root_node_id
+ for task_idx, _ in enumerate(self.tasks):
+ task_node_id = f"task_{root_id}_{task_idx}"
+
+ task_node = RepresentationNode(
+ node_id=task_node_id,
+ operation=None,
+ inputs=[root_id],
+ parameters={
+ "_node_kind": "task",
+ "_task_idx": task_idx,
+ "_dag_root_id": root_id,
+ },
+ )
+
+ task_root_dag = RepresentationDag(
+ nodes=[*dag.nodes, task_node], root_node_id=task_node_id
+ )
+ expanded_dags.append(task_root_dag)
+ return expanded_dags
+
+ def _process_modality(self, modality, skip_remaining: int = 0, scheduler=None):
+ local_results = UnimodalResults(
+ [modality],
+ self.tasks,
+ debug=False,
+ store_cache=False,
+ metric_name=self.metric_name,
+ )
modality_specific_operators = self._get_modality_operators(
modality.modality_type
)
dags = []
- operators = []
for operator in modality_specific_operators:
dags.extend(self._build_modality_dag(modality, operator()))
- operators.append(operator())
- external_cache = LRUCache(max_size=32)
- rep_cache = None
- if hasattr(modality, "data_loader") and modality.data_loader.chunk_size:
- rep_cache = modality.apply_representations(operators)
-
- for dag in dags:
- representations = dag.execute(
- [modality], external_cache=external_cache, rep_cache=rep_cache
- )
- node_id = list(representations.keys())[-1]
- node = dag.get_node_by_id(node_id)
- if node.operation is None:
- continue
-
- reps = self._get_representation_chain(node, dag)
- combination = next((op for op in reps if isinstance(op, Fusion)), None)
- self._evaluate_local(
- representations[node_id], local_results, dag, combination
+ if (
+ modality.modality_type == ModalityType.TIMESERIES
+ or modality.modality_type == ModalityType.AUDIO
+ ):
+ dags.extend(
+ self.temporal_context_operators(
+ modality,
+ self.builders[modality.modality_id],
+ dags[0].get_leaf_node_id(),
+ )
)
- if self.debug:
- visualize_dag(dag)
+
+ dags = self.add_aggregation_operator(self.builders[modality.modality_id], dags)
+ dags = pushdown_aggregation(dags)
+
+ if skip_remaining > 0:
+ dags = dags[skip_remaining:]
+
+ expanded_dags = self._expand_dags_with_task_roots(dags)
+
+ node_executor = NodeExecutor(
+ expanded_dags,
+ [modality],
+ self.tasks,
+ self._checkpoint_manager,
+ self.max_num_workers,
+ self.result_path,
+ )
+ task_results = node_executor.run()
+
+ for task_result in task_results:
+ local_results.add_task_result(task_result)
if self.save_all_results:
timestr = time.strftime("%Y%m%d-%H%M%S")
file_name = f"{modality.modality_id}_unimodal_results_{timestr}.pkl"
-
with open(file_name, "wb") as f:
pickle.dump(local_results.results, f)
return local_results
- def _get_representation_chain(
- self, node: "RepresentationNode", dag: RepresentationDag
- ) -> List[Any]:
- representations = []
- if node.operation:
- representations.append(node.operation)
-
- for input_id in node.inputs:
- input_node = dag.get_node_by_id(input_id)
- if input_node.operation:
- representations.extend(self._get_representation_chain(input_node, dag))
-
- return representations
-
def _merge_results(self, local_results):
for modality_id in local_results.results:
for task_name in local_results.results[modality_id]:
@@ -252,78 +364,6 @@ class UnimodalOptimizer:
local_results.results[modality_id][task_name]
)
- # for modality in self.modalities:
- # for task_name in local_results.cache[modality]:
- # for key, value in local_results.cache[modality][task_name].items():
- # self.operator_performance.cache[modality][task_name][key] = value
-
- def _evaluate_local(self, modality, local_results, dag, combination=None):
- if self._tasks_require_same_dims:
- if self.expected_dimensions == 1 and get_shape(modality.metadata) > 1:
- builder = self.builders[modality.modality_id]
- agg_operator = AggregatedRepresentation()
- rep_node_id = builder.create_operation_node(
- agg_operator.__class__,
- [dag.root_node_id],
- agg_operator.get_current_parameters(),
- )
- dag = builder.build(rep_node_id)
-
- aggregated_modality = agg_operator.transform(modality)
-
- for task in self.tasks:
- start = time.perf_counter()
- scores = task.run(aggregated_modality.data)
- end = time.perf_counter()
-
- local_results.add_result(
- scores,
- aggregated_modality,
- task.model.name,
- end - start,
- combination,
- dag,
- )
- else:
- modality.pad()
- for task in self.tasks:
- start = time.perf_counter()
- scores = task.run(modality.data)
- end = time.perf_counter()
- local_results.add_result(
- scores, modality, task.model.name, end - start, combination, dag
- )
- else:
- for task in self.tasks:
- if task.expected_dim == 1 and get_shape(modality.metadata) > 1:
- builder = self.builders[modality.modality_id]
- agg_operator = AggregatedRepresentation(Aggregation())
- rep_node_id = builder.create_operation_node(
- agg_operator.__class__,
- [dag.root_node_id],
- agg_operator.get_current_parameters(),
- )
- dag = builder.build(rep_node_id)
- start_rep = time.perf_counter()
- representations = dag.execute([modality])
- end_rep = time.perf_counter()
- modality.transform_time += end_rep - start_rep
- node_id = list(representations.keys())[-1]
-
- start = time.perf_counter()
- scores = task.run(representations[node_id].data)
- end = time.perf_counter()
- local_results.add_result(
- scores, modality, task.model.name, end - start, combination, dag
- )
- else:
- start = time.perf_counter()
- scores = task.run(modality.data)
- end = time.perf_counter()
- local_results.add_result(
- scores, modality, task.model.name, end - start, combination, dag
- )
-
def add_dimensionality_reduction_operators(self, builder, current_node_id):
dags = []
modality_type = (
@@ -389,7 +429,7 @@ class UnimodalOptimizer:
operator.get_current_parameters(),
)
- agg_operator = AggregatedRepresentation()
+ agg_operator = AggregatedRepresentation(target_dimensions=1)
context_agg_node_id = builder.create_operation_node(
agg_operator.__class__,
[context_rep_node_id],
@@ -409,7 +449,7 @@ class UnimodalOptimizer:
for rep in not_self_contained_reps:
other_rep_id = builder.create_operation_node(
- rep, [leaf_id], rep().parameters
+ rep, [leaf_id], rep().get_current_parameters()
)
for combination in self._combination_operators:
combine_id = builder.create_operation_node(
@@ -448,13 +488,41 @@ class UnimodalOptimizer:
)
)
- if (
- modality.modality_type == ModalityType.TIMESERIES
- or modality.modality_type == ModalityType.AUDIO
- ):
- dags.extend(self.temporal_context_operators(modality, builder, leaf_id))
return dags
+ def _aggregation_needed(self, dag: RepresentationDag) -> bool:
+ for modality in self.modalities:
+ if modality.modality_id == dag.nodes[0].modality_id:
+ last_stats = modality.stats
+ break
+ for node in dag.nodes[1:]:
+ last_stats = node.operation(params=node.parameters).get_output_stats(
+ last_stats
+ )
+ return len(last_stats.output_shape) > 1
+
+ def add_aggregation_operator(self, builder, dags):
+ new_dags = []
+ if self._tasks_require_same_dims and self.expected_dimensions == 1:
+ aggregated_dags = []
+ for dag in dags:
+ if self._aggregation_needed(dag):
+ agg_op = AggregatedRepresentation(
+ target_dimensions=self.expected_dimensions
+ )
+ agg_node_id = builder.create_operation_node(
+ agg_op.__class__,
+ [dag.root_node_id],
+ agg_op.get_current_parameters(),
+ )
+ aggregated_dags.append(builder.build(agg_node_id))
+ else:
+ aggregated_dags.append(dag)
+ new_dags = aggregated_dags
+ else:
+ new_dags = dags
+ return new_dags
+
def default_context_operators(
self, modality, builder, leaf_id, rep_dag, apply_context_to_leaf=False
):
@@ -463,6 +531,7 @@ class UnimodalOptimizer:
if (
modality.modality_type != ModalityType.TEXT
and modality.modality_type != ModalityType.VIDEO
+ and modality.modality_type != ModalityType.IMAGE
):
context_operators = self._get_context_operators(modality.modality_type)
for context_op in context_operators:
@@ -487,7 +556,9 @@ class UnimodalOptimizer:
return dags
def temporal_context_operators(self, modality, builder, leaf_id):
- aggregators = self.operator_registry.get_representations(modality.modality_type)
+ aggregators = self.operator_registry.get_context_representations(
+ modality.modality_type
+ )
context_operators = self._get_context_operators(modality.modality_type)
dags = []
@@ -525,33 +596,65 @@ class UnimodalResults:
self.results[modality] = {task_name: [] for task_name in self.task_names}
self.cache[modality] = {task_name: [] for task_name in self.task_names}
- def add_result(self, scores, modality, task_name, task_time, combination, dag):
+ def add_task_result(self, task_result: ResultEntry):
+ task_name = self.task_names[
+ task_result.dag.nodes[-1].parameters.get("_task_idx", 0)
+ ]
+ self.results[task_result.dag.nodes[0].modality_id][task_name].append(
+ task_result
+ )
+ # TODO: Take care of modality cache in executor
+ if (
+ self.store_cache
+ and task_result.val_score[self.metric_name]
+ > self.cache[task_result.dag.nodes[0].modality_id][task_name][-1].val_score[
+ self.metric_name
+ ]
+ ):
+ self.cache[task_result.dag.nodes[0].modality_id][task_name].append(
+ task_result
+ )
+ self.cache[task_result.dag.nodes[0].modality_id][task_name] = self.cache[
+ task_result.dag.nodes[0].modality_id
+ ][task_name][: self.k]
+ if self.debug:
+ print(f"{task_result.dag.nodes[0].modality_id}_{task_name}: {task_result}")
+
+ def add_result(
+ self,
+ scores,
+ transform_time,
+ task_name,
+ task_time,
+ dag,
+ modality_id,
+ modality=None,
+ ):
entry = ResultEntry(
train_score=scores[0].average_scores,
val_score=scores[1].average_scores,
test_score=scores[2].average_scores,
- representation_time=modality.transform_time,
+ representation_time=transform_time,
task_time=task_time,
- combination=combination.name if combination else "",
dag=dag,
)
- self.results[modality.modality_id][task_name].append(entry)
- if self.store_cache:
- self.cache[modality.modality_id][task_name].append(modality)
+ scores = [
+ -item.val_score[self.metric_name]
+ for item in self.results[modality_id][task_name]
+ ]
+ pos = (
+ bisect_left(scores, -entry.val_score[self.metric_name])
+ if len(scores) > 0
+ else 0
+ )
+ self.results[modality_id][task_name].insert(pos, entry)
- results = self.results[modality.modality_id][task_name]
- if self.k != -1 and len(results) > self.k:
- ranked, sorted_indices = rank_by_tradeoff(
- results, performance_metric_name=self.metric_name
- )
- keep = set(sorted_indices[: self.k])
-
- self.cache[modality.modality_id][task_name] = [
- m
- for i, m in enumerate(self.cache[modality.modality_id][task_name])
- if i in keep
- ]
+ if self.store_cache and pos < self.k and modality is not None:
+ self.cache[modality.modality_id][task_name].insert(pos, modality)
+ self.cache[modality.modality_id][task_name] = self.cache[
+ modality.modality_id
+ ][task_name][: self.k]
if self.debug:
print(f"{modality.modality_id}_{task_name}: {entry}")
@@ -583,10 +686,14 @@ class UnimodalResults:
sorted_indices = sorted_indices[: self.k]
task_cache = self.cache.get(modality.modality_id, {}).get(task.model.name, None)
if not task_cache:
- cache = [
- list(results[i].dag.execute([modality]).values())[-1]
- for i in range(len(results))
- ]
+ cache = []
+ for result in results:
+ if result.dag.nodes[-1].parameters.get("_node_kind", False) == "task":
+ dag = copy.deepcopy(result.dag)
+ dag.nodes = dag.nodes[:-1]
+ dag.root_node_id = dag.nodes[-1].node_id
+ cache.append(dag.execute([modality]))
+
elif isinstance(task_cache, list):
cache = task_cache
else:
@@ -609,15 +716,3 @@ class UnimodalResults:
] = cache
return results, cache
-
-
-@dataclass
-class ResultEntry:
- val_score: PerformanceMeasure = None
- train_score: PerformanceMeasure = None
- test_score: PerformanceMeasure = None
- representation_time: float = 0.0
- task_time: float = 0.0
- combination: str = ""
- dag: RepresentationDag = None
- tradeoff_score: float = 0.0
diff --git a/src/main/python/systemds/scuro/modality/joined.py b/src/main/python/systemds/scuro/modality/joined.py
index 1a58df9256..335d1959fd 100644
--- a/src/main/python/systemds/scuro/modality/joined.py
+++ b/src/main/python/systemds/scuro/modality/joined.py
@@ -231,18 +231,14 @@ class JoinedModality(Modality):
new_left = Modality(left_modality.modality_type, {})
new_right = Modality(right_modality.modality_type, {})
- while (
- left_modality.data_loader.next_chunk < left_modality.data_loader.num_chunks
- ):
+ for _ in left_modality.iter_raw_data_chunks(reset=True):
if chunk_right:
right_modality.extract_raw_data()
starting_idx = 0
else:
starting_idx = (
- left_modality.data_loader.next_chunk
- * left_modality.data_loader.chunk_size
- )
- left_modality.extract_raw_data()
+ left_modality.data_loader.next_chunk - 1
+ ) * left_modality.data_loader.chunk_size
self.execute(starting_idx)
diff --git a/src/main/python/systemds/scuro/modality/modality.py b/src/main/python/systemds/scuro/modality/modality.py
index 8162662b7f..1bc8180e19 100644
--- a/src/main/python/systemds/scuro/modality/modality.py
+++ b/src/main/python/systemds/scuro/modality/modality.py
@@ -24,6 +24,7 @@ import numpy as np
from systemds.scuro.modality.type import ModalityType
from systemds.scuro.representations import utils
+from systemds.scuro.representations.representation import RepresentationStats
class Modality:
@@ -45,10 +46,9 @@ class Modality:
self.metadata = metadata
self.data = []
self.data_type = data_type
- self.cost = None
- self.shape = None
self.modality_id = modality_id
self.transform_time = transform_time if transform_time else 0
+ self.stats = RepresentationStats(0, ())
@property
def data(self):
diff --git a/src/main/python/systemds/scuro/modality/transformed.py b/src/main/python/systemds/scuro/modality/transformed.py
index a443f5a313..e185deb7c9 100644
--- a/src/main/python/systemds/scuro/modality/transformed.py
+++ b/src/main/python/systemds/scuro/modality/transformed.py
@@ -25,7 +25,7 @@ from systemds.scuro.modality.joined import JoinedModality
from systemds.scuro.modality.modality import Modality
from systemds.scuro.representations.window_aggregation import WindowAggregation
import time
-import copy
+import sys
class TransformedModality(Modality):
@@ -37,6 +37,7 @@ class TransformedModality(Modality):
new_modality_type=None,
self_contained=True,
set_data=False,
+ aggregate_dim=(0,),
):
"""
Parent class of the different Modalities (unimodal & multimodal)
@@ -56,36 +57,47 @@ class TransformedModality(Modality):
)
if set_data:
self.data = modality.data
- self.transformation = None
+
self.self_contained = (
self_contained and transformation.self_contained
if isinstance(transformation, TransformedModality)
else True
)
- self.add_transformation(transformation, modality)
+ self.aggregate_dim = aggregate_dim
if modality.__class__.__name__ == "UnimodalModality":
for k, v in self.metadata.items():
if "attention_masks" in v:
del self.metadata[k]["attention_masks"]
- def add_transformation(self, transformation, modality):
- if (
- transformation.__class__.__bases__[0].__name__ == "Fusion"
- and type(modality).__name__ == "TransformedModality"
- and modality.transformation[0].__class__.__bases__[0].__name__ != "Fusion"
- ):
- self.transformation = []
- else:
- self.transformation = (
- []
- if type(modality).__name__ != "TransformedModality"
- else copy.deepcopy(modality.transformation)
- )
- self.transformation.append(transformation)
-
def copy_from_instance(self):
- return type(self)(self, self.transformation)
+ """
+ Create a copy of the transformed modality instance
+ """
+ return type(self)(self, None, self.modality_type)
+
+ def calculate_memory_usage(self):
+ data_bytes = 0
+ for instance in self.data:
+ data_bytes += self._estimate_data_bytes(instance)
+
+ md_bytes = 0
+ for key, value in self.metadata.items():
+ md_bytes += self._estimate_data_bytes(key)
+ md_bytes += self._estimate_data_bytes(value)
+
+ total_bytes = (
+ data_bytes
+ + md_bytes
+ + sys.getsizeof(self.data_type)
+ + sys.getsizeof(self.modality_id)
+ + sys.getsizeof(self.schema)
+ + sys.getsizeof(self.stats)
+ + sys.getsizeof(self.self_contained)
+ + sys.getsizeof(self.transform_time)
+ + sys.getsizeof(self.modality_type)
+ )
+ return total_bytes
def join(self, right, join_condition):
chunked_execution = False
@@ -152,9 +164,9 @@ class TransformedModality(Modality):
transformed_modality.transform_time += time.time() - start
return transformed_modality
- def apply_representation(self, representation):
+ def apply_representation(self, representation, aggregation=None):
start = time.time()
- new_modality = representation.transform(self)
+ new_modality = representation.transform(self, aggregation=aggregation)
new_modality.update_metadata()
new_modality.transform_time += time.time() - start
new_modality.self_contained = representation.self_contained
@@ -225,3 +237,13 @@ class TransformedModality(Modality):
all_outputs.append(out)
start = end
return np.concatenate(all_outputs, axis=0)
+
+ def _estimate_data_bytes(self, instance):
+ if isinstance(instance, np.ndarray):
+ return instance.nbytes
+ elif isinstance(instance, list):
+ return sum(self._estimate_data_bytes(item) for item in instance)
+ elif isinstance(instance, dict):
+ return sum(self._estimate_data_bytes(item) for item in instance.values())
+ else:
+ return sys.getsizeof(instance)
diff --git a/src/main/python/systemds/scuro/modality/type.py b/src/main/python/systemds/scuro/modality/type.py
index 85f4d04e9b..9c883efec3 100644
--- a/src/main/python/systemds/scuro/modality/type.py
+++ b/src/main/python/systemds/scuro/modality/type.py
@@ -205,6 +205,7 @@ class ModalityType(Flag):
"VIDEO": "create_video_metadata",
"IMAGE": "create_image_metadata",
"TIMESERIES": "create_ts_metadata",
+ "EMBEDDING": "create_embedding_metadata",
}
def get_schema(self):
@@ -295,6 +296,14 @@ class ModalityType(Flag):
md["is_multivariate"] = len(signal_names) > 1
return md
+ def create_embedding_metadata(self, data, is_single_instance=True):
+ md = deepcopy(self.get_schema())
+ md = ModalitySchemas.update_base_metadata(md, data, is_single_instance)
+ md["data_layout"]["representation"] = DataLayout.SINGLE_LEVEL
+ md["data_layout"]["type"] = np.float32
+ md["data_layout"]["shape"] = data.shape
+ return md
+
def create_video_metadata(self, frequency, length, width, height, num_channels):
md = deepcopy(self.get_schema())
md["frequency"] = frequency
diff --git a/src/main/python/systemds/scuro/modality/unimodal_modality.py b/src/main/python/systemds/scuro/modality/unimodal_modality.py
index 89d95810e0..319599d680 100644
--- a/src/main/python/systemds/scuro/modality/unimodal_modality.py
+++ b/src/main/python/systemds/scuro/modality/unimodal_modality.py
@@ -26,6 +26,7 @@ from systemds.scuro.dataloader.base_loader import BaseLoader
from systemds.scuro.modality.modality import Modality
from systemds.scuro.modality.joined import JoinedModality
from systemds.scuro.modality.transformed import TransformedModality
+from systemds.scuro.representations.representation import RepresentationStats
from systemds.scuro.utils.identifier import Identifier
@@ -44,6 +45,7 @@ class UnimodalModality(Modality):
data_loader.data_type,
)
self.data_loader = data_loader
+ self.stats = data_loader.stats
def copy_from_instance(self):
new_instance = type(self)(self.data_loader)
@@ -59,6 +61,24 @@ class UnimodalModality(Modality):
return self.metadata[self.dataIndex][position]
+ def get_stats(self):
+ return self.stats
+
+ def get_output_stats(self):
+ return RepresentationStats(self.stats.num_instances, self.stats.output_shape)
+
+ def estimate_memory_bytes(self):
+ memory_bytes = 1
+ for i in self.stats.output_shape:
+ memory_bytes *= i
+
+ return (
+ self.stats.num_instances * memory_bytes * 4
+ ) # TODO: check how to meausure str size
+
+ def estimate_peak_memory_bytes(self):
+ return {"cpu_peak_bytes": self.estimate_memory_bytes(), "gpu_peak_bytes": 0.0}
+
def extract_raw_data(self):
"""
Uses the data loader to read the raw data from a specified location
@@ -66,6 +86,14 @@ class UnimodalModality(Modality):
"""
self.data, self.metadata = self.data_loader.load()
+ def iter_raw_data_chunks(self, reset: bool = True):
+ for data, metadata, chunk_indices in self.data_loader.iter_loaded_chunks(
+ reset=reset
+ ):
+ self.data = data
+ self.metadata = metadata
+ yield chunk_indices
+
def join(self, other, join_condition):
if isinstance(other, UnimodalModality):
self.data_loader.update_chunk_sizes(other.data_loader)
@@ -106,7 +134,7 @@ class UnimodalModality(Modality):
if self.data is None:
raise Exception("Data is None")
- def apply_representations(self, representations):
+ def apply_representations(self, representations, aggregation=None):
"""
Applies a list of representations to the modality. Specifically, it applies the representations to the modality in a chunked manner.
:param representations: List of representations to apply
@@ -116,7 +144,6 @@ class UnimodalModality(Modality):
padding_per_representation = {}
original_lengths_per_representation = {}
- # Initialize dictionaries for each representation
for representation in representations:
transformed_modality = TransformedModality(self, representation.name)
transformed_modality.data = []
@@ -130,9 +157,7 @@ class UnimodalModality(Modality):
time.time()
) # TODO: should be repalced in unimodal_representation.transform
if self.data_loader.chunk_size:
- self.data_loader.reset()
- while self.data_loader.next_chunk < self.data_loader.num_chunks:
- self.extract_raw_data()
+ for _ in self.iter_raw_data_chunks(reset=True):
for representation in representations:
transformed_chunk = representation.transform(self)
transformed_modalities_per_representation[
@@ -150,25 +175,6 @@ class UnimodalModality(Modality):
if not self.has_data():
self.extract_raw_data()
new_modality = representation.transform(self)
-
- for i, d in enumerate(new_modality.data):
- output = np.array(d)
- if np.isnan(output).any():
- new_modality.data[i] = np.where(np.isnan(output), 0, output)
-
- if not all(
- "attention_masks" in entry for entry in new_modality.metadata.values()
- ):
- for d in new_modality.data:
- if d.shape[0] == 1 and d.ndim == 2:
- padding_per_representation[representation.name] = True
- original_lengths_per_representation[representation.name].append(
- d.shape[1]
- )
- else:
- original_lengths_per_representation[representation.name].append(
- d.shape[0]
- )
transformed_modalities_per_representation[representation.name] = (
new_modality
)
@@ -188,8 +194,10 @@ class UnimodalModality(Modality):
gc.collect()
return transformed_modalities_per_representation
- def apply_representation(self, representation):
- return self.apply_representations([representation])[representation.name]
+ def apply_representation(self, representation, aggregation=None):
+ return self.apply_representations([representation], aggregation=aggregation)[
+ representation.name
+ ]
def _apply_padding(self, modality, original_lengths, pad_dim_one):
if len(original_lengths) > 0 and min(original_lengths) < max(original_lengths):
@@ -217,6 +225,7 @@ class UnimodalModality(Modality):
mode="constant",
constant_values=0,
)
+ padded_embeddings.append(padded)
elif len(embeddings.shape) == 2:
padded = np.pad(
embeddings,
@@ -224,6 +233,7 @@ class UnimodalModality(Modality):
mode="constant",
constant_values=0,
)
+ padded_embeddings.append(padded)
elif len(embeddings.shape) == 3:
padded = np.pad(
embeddings,
diff --git a/src/main/python/systemds/scuro/representations/aggregate.py b/src/main/python/systemds/scuro/representations/aggregate.py
index 9503a48587..8389fadbd0 100644
--- a/src/main/python/systemds/scuro/representations/aggregate.py
+++ b/src/main/python/systemds/scuro/representations/aggregate.py
@@ -26,20 +26,20 @@ from systemds.scuro.representations import utils
class Aggregation:
@staticmethod
- def _mean_agg(data):
- return np.mean(data, axis=0)
+ def _mean_agg(data, aggregate_dim=0):
+ return np.mean(data, axis=aggregate_dim)
@staticmethod
- def _max_agg(data):
- return np.max(data, axis=0)
+ def _max_agg(data, aggregate_dim=0):
+ return np.max(data, axis=aggregate_dim)
@staticmethod
- def _min_agg(data):
- return np.min(data, axis=0)
+ def _min_agg(data, aggregate_dim=0):
+ return np.min(data, axis=aggregate_dim)
@staticmethod
- def _sum_agg(data):
- return np.sum(data, axis=0)
+ def _sum_agg(data, aggregate_dim=0):
+ return np.sum(data, axis=aggregate_dim)
_aggregation_function = {
"mean": _mean_agg.__func__,
@@ -48,7 +48,7 @@ class Aggregation:
"sum": _sum_agg.__func__,
}
- def __init__(self, aggregation_function="mean", pad_modality=False, params=None):
+ def __init__(self, aggregation_function="mean", pad_modality=True, params=None):
if params is not None:
aggregation_function = params["aggregation_function"]
pad_modality = params["pad_modality"]
@@ -62,11 +62,16 @@ class Aggregation:
self.aggregation_function_name = aggregation_function
self.parameters = {
- "aggregation_function": aggregation_function,
- "pad_modality": pad_modality,
+ "aggregation_function": self._aggregation_function.keys(),
}
- def execute(self, modality):
+ def get_current_parameters(self):
+ return {
+ "aggregation_function": self.aggregation_function_name,
+ "pad_modality": self.pad_modality,
+ }
+
+ def execute(self, modality, aggregate_dim=(0,)):
data = []
max_len = 0
for i, instance in enumerate(modality.data):
@@ -77,33 +82,51 @@ class Aggregation:
or modality.modality_type == ModalityType.VIDEO
) and instance.ndim > 2:
aggregated_data = instance.flatten()
+ elif (
+ isinstance(instance, np.ndarray)
+ and instance.ndim == 2
+ and instance.shape[1] == 1
+ ):
+ aggregated_data = instance.flatten()
else:
- aggregated_data = self._aggregation_func(instance)
+ aggregated_data = self._aggregation_func(instance, aggregate_dim)
else:
aggregated_data = []
for entry in instance:
- aggregated_data.append(self._aggregation_func(entry))
+ aggregated_data.append(self._aggregation_func(entry, aggregate_dim))
+
+ if isinstance(aggregated_data, np.generic):
+ aggregated_data = np.array(
+ [aggregated_data], dtype=aggregated_data.dtype
+ )
+ elif isinstance(aggregated_data, np.ndarray) and aggregated_data.ndim == 0:
+ aggregated_data = aggregated_data.reshape(1)
+ elif not isinstance(aggregated_data, (list, np.ndarray)):
+ aggregated_data = np.array([aggregated_data])
+
max_len = max(max_len, len(aggregated_data))
data[i] = aggregated_data
- for i, instance in enumerate(data):
- if isinstance(instance, np.ndarray):
- if len(instance) < max_len:
- padded_data = np.zeros(max_len, dtype=instance.dtype)
- padded_data[: len(instance)] = instance
+ if self.pad_modality:
+ for i, instance in enumerate(data):
+ if isinstance(instance, np.ndarray):
+ if len(instance) < max_len:
+ padded_data = np.zeros(max_len, dtype=instance.dtype)
+ padded_data[: len(instance)] = instance
+ data[i] = padded_data
+ else:
+ padded_data = []
+ for entry in instance:
+ padded_data.append(utils.pad_sequences(entry, max_len))
data[i] = padded_data
- else:
- padded_data = []
- for entry in instance:
- padded_data.append(utils.pad_sequences(entry, max_len))
- data[i] = padded_data
+ data = np.asarray(data)
- return np.array(data)
+ return data
def transform(self, modality):
return self.execute(modality)
- def aggregate_instance(self, instance):
+ def compute_feature(self, instance):
return self._aggregation_func(instance)
def get_aggregation_functions(self):
diff --git a/src/main/python/systemds/scuro/representations/aggregated_representation.py b/src/main/python/systemds/scuro/representations/aggregated_representation.py
index cad1a4a448..4caa22762b 100644
--- a/src/main/python/systemds/scuro/representations/aggregated_representation.py
+++ b/src/main/python/systemds/scuro/representations/aggregated_representation.py
@@ -18,27 +18,120 @@
# under the License.
#
# -------------------------------------------------------------
+import copy
from systemds.scuro.modality.transformed import TransformedModality
-from systemds.scuro.representations.representation import Representation
+from systemds.scuro.representations.representation import (
+ Representation,
+ RepresentationStats,
+)
from systemds.scuro.representations.aggregate import Aggregation
import time
+import numpy as np
class AggregatedRepresentation(Representation):
- def __init__(self, aggregation="mean"):
+ def __init__(self, aggregation="mean", target_dimensions=None, params=None):
+ if params is not None:
+ if "aggregation_function_aggregation_function" in params:
+ aggregation = params["aggregation_function_aggregation_function"]
+ elif "aggregation_function" in params:
+ aggregation = params["aggregation_function"]
+ else:
+ aggregation = params["aggregation"]
+ if "target_dimensions" in params:
+ target_dimensions = params["target_dimensions"]
parameters = {
"aggregation": list(Aggregation().get_aggregation_functions()),
}
super().__init__("AggregatedRepresentation", parameters)
+ self.aggregation_function = aggregation
self.aggregation = Aggregation(aggregation)
self.self_contained = True
+ self.target_dimensions = target_dimensions
+ self.data_type = np.float32
+
+ def get_output_stats(self, input_stats: RepresentationStats) -> RepresentationStats:
+ input_shape = list(copy.deepcopy(input_stats.output_shape))
+ input_aggregate_dim = copy.deepcopy(input_stats.aggregate_dim)
+ for input_aggregate_dim in reversed(input_aggregate_dim):
+ input_shape.pop(input_aggregate_dim)
+ out_shape = tuple(input_shape)
+ self.stats = RepresentationStats(
+ input_stats.num_instances,
+ out_shape,
+ output_shape_is_known=input_stats.output_shape_is_known,
+ aggregate_dim=input_stats.aggregate_dim,
+ )
+ return self.stats
+
+ def estimate_output_memory_bytes(self, input_stats: RepresentationStats) -> int:
+ out_shape = self.get_output_stats(input_stats).output_shape
+ out_numel = int(np.prod(out_shape)) if len(out_shape) > 0 else 1
+ dtype_size = np.dtype(self.data_type).itemsize
+ return int(input_stats.num_instances * out_numel * dtype_size)
+
+ def estimate_peak_memory_bytes(self, input_stats: RepresentationStats) -> dict:
+ dtype_size = np.dtype(self.data_type).itemsize
+ in_shape = tuple(input_stats.output_shape)
+ in_numel = int(np.prod(in_shape)) if len(in_shape) > 0 else 1
+ input_bytes = int(input_stats.num_instances * in_numel * dtype_size)
+ output_bytes = self.estimate_output_memory_bytes(input_stats)
+ safety = 1.2
+ cpu_peak = input_bytes * 2 + output_bytes * 2
+
+ return {
+ "cpu_peak_bytes": int(cpu_peak * safety),
+ "gpu_peak_bytes": 0,
+ }
+
+ def execute(self, data):
+ return self.aggregation.compute_feature(data)
def transform(self, modality):
start = time.perf_counter()
aggregated_modality = TransformedModality(
modality, self, self_contained=modality.self_contained
)
- aggregated_modality.data = self.aggregation.execute(modality)
+
+ aggregate_dim = (0,)
+ if self.target_dimensions is not None:
+ input_dimensions = self._get_input_dimensions(modality.data)
+
+ if len(input_dimensions) == self.target_dimensions:
+ return modality
+ else:
+
+ i = 1
+ while len(input_dimensions) - 1 > self.target_dimensions:
+ aggregate_dim = aggregate_dim + (i,)
+ i += 1
+ input_dimensions = input_dimensions[:-1]
+
+ aggregated_data = self.aggregation.execute(modality, aggregate_dim)
+
+ aggregated_modality.data = aggregated_data
end = time.perf_counter()
aggregated_modality.transform_time += end - start
+
+ self.assert_output_stats(aggregated_data)
return aggregated_modality
+
+ def get_current_parameters(self):
+ current_params = {}
+ for key, value in self.aggregation.get_current_parameters().items():
+ current_params[f"aggregation_function_{key}"] = value
+ current_params["self_contained"] = self.self_contained
+ current_params["target_dimensions"] = self.target_dimensions
+ return current_params
+
+ def assert_output_stats(self, aggregated_data):
+ if self.stats:
+ assert len(aggregated_data) == self.stats.num_instances
+ assert aggregated_data[0].shape == self.stats.output_shape
+
+ def _get_input_dimensions(self, data):
+ if isinstance(data[0], list):
+ input_dimensions = (len(data[0]), *data[0][0].shape)
+ else:
+ input_dimensions = data[0].shape
+ return input_dimensions
diff --git a/src/main/python/systemds/scuro/representations/bert.py b/src/main/python/systemds/scuro/representations/bert.py
index 6f4d3705a1..ab23f46345 100644
--- a/src/main/python/systemds/scuro/representations/bert.py
+++ b/src/main/python/systemds/scuro/representations/bert.py
@@ -18,15 +18,19 @@
# under the License.
#
# -------------------------------------------------------------
+from dataclasses import dataclass
import numpy as np
from systemds.scuro.modality.transformed import TransformedModality
+from systemds.scuro.representations.representation import RepresentationStats
from systemds.scuro.representations.unimodal import UnimodalRepresentation
import torch
from transformers import AutoTokenizer, AutoModel
from systemds.scuro.representations.utils import save_embeddings
from systemds.scuro.modality.type import ModalityType
from systemds.scuro.drsearch.operator_registry import register_representation
-from systemds.scuro.utils.static_variables import get_device
+from systemds.scuro.utils.memory_utility import (
+ get_device,
+)
import os
from torch.utils.data import DataLoader
from systemds.scuro.utils.torch_dataset import TextDataset, TextSpanDataset
@@ -43,7 +47,11 @@ class BertFamily(UnimodalRepresentation):
parameters={},
output_file=None,
max_seq_length=512,
+ batch_size=32,
+ aggregation=None,
+ params=None,
):
+ parameters = {"batch_size": [1, 2, 4, 8, 16, 32, 64, 128]}
self.model_name = model_name
super().__init__(representation_name, ModalityType.EMBEDDING, parameters)
@@ -52,13 +60,126 @@ class BertFamily(UnimodalRepresentation):
self.max_seq_length = max_seq_length
self.needs_context = True
self.initial_context_length = 350
+ self.device = None
+ self.batch_size = batch_size
+ self.gpu_id = None
+ self.device = get_device()
+ self.data_type = torch.float32
+ self.aggregation = aggregation
+ self.params = params
- def transform(self, modality):
+ @property
+ def gpu_id(self):
+ return self._gpu_id
+
+ @gpu_id.setter
+ def gpu_id(self, gpu_id):
+ self._gpu_id = gpu_id
+ self.device = get_device(gpu_id)
+
+ def set_parameters(
+ self, params, max_seq_length, batch_size, layer, output_file, aggregation=None
+ ):
+ if params is not None:
+ self.max_seq_length = int(params.get("max_seq_length", max_seq_length))
+ self.batch_size = int(params.get("batch_size", batch_size))
+ self.layer_name = params.get("layer_name", layer)
+ self.output_file = params.get("output_file", output_file)
+ else:
+ self.max_seq_length = max_seq_length
+ self.batch_size = batch_size
+ self.layer_name = layer
+ self.output_file = output_file
+
+ def get_output_stats(self, input_stats) -> RepresentationStats:
+ if not isinstance(input_stats, RepresentationStats):
+ self.stats = RepresentationStats(
+ input_stats.num_instances,
+ (self.max_seq_length, 768),
+ aggregate_dim=(0,),
+ )
+ else:
+ self.stats = RepresentationStats(
+ input_stats.num_instances,
+ (input_stats.output_shape[0], self.max_seq_length, 768),
+ aggregate_dim=(
+ 0,
+ 1,
+ ),
+ )
+ if self.params and "_pushdown_aggregation" in self.params:
+ output_shape = (768,)
+ self.stats.output_shape = output_shape
+ self.stats.aggregate_dim = None
+ return self.stats
+
+ def estimate_output_memory_bytes(self, input_stats):
+ output_stats = self.get_output_stats(input_stats).output_shape
+ return int(
+ input_stats.num_instances * np.prod(output_stats) * self.data_type.itemsize
+ )
+
+ def estimate_peak_memory_bytes(self, input_stats):
+ model = AutoModel.from_pretrained(self.model_name)
+ params_bytes = model.get_memory_footprint()
+ tokenizer_bytes = 60 * 1024 * 1024 # rough upper bound
+
+ output_bytes = self.estimate_output_memory_bytes(input_stats)
+ if isinstance(input_stats, RepresentationStats):
+ per_instance_input_bytes = (
+ int(np.prod(input_stats.output_shape)) * self.data_type.itemsize
+ )
+ input_bytes_all_instances = per_instance_input_bytes
+ else:
+ per_instance_input_bytes = (
+ int(np.prod(input_stats.output_shape)) * self.data_type.itemsize
+ )
+ input_bytes_all_instances = self.batch_size * per_instance_input_bytes
+ safety_margin_bytes = 100 * 1024 * 1024
+
+ batch_output_bytes = output_bytes / input_stats.num_instances * self.batch_size
+
+ model_specific_margin_bytes = 0
+ if "albert" in self.model_name.lower():
+ model_specific_margin_bytes = 300 * 1024 * 1024
+
+ cpu_peak = (
+ params_bytes
+ + tokenizer_bytes
+ + input_bytes_all_instances
+ + output_bytes
+ + batch_output_bytes
+ + safety_margin_bytes
+ + model_specific_margin_bytes
+ + 128 * 1024 * 1024
+ )
+
+ cfg = model.config
+ hidden_size = cfg.hidden_size
+ num_layers = cfg.num_hidden_layers
+ intermediate_size = getattr(cfg, "intermediate_size", 4 * hidden_size)
+
+ batch_tokens = self.batch_size * self.max_seq_length
+ activations_per_token_per_layer = (
+ self.C * (hidden_size + intermediate_size) * self.data_type.itemsize
+ )
+ activations_bytes = batch_tokens * num_layers * activations_per_token_per_layer
+
+ gpu_peak = params_bytes + per_instance_input_bytes + activations_bytes
+
+ return {
+ "cpu_peak_bytes": int(cpu_peak),
+ "gpu_peak_bytes": int(gpu_peak),
+ }
+
+ def transform(self, modality, aggregation=None):
transformed_modality = TransformedModality(modality, self)
tokenizer = AutoTokenizer.from_pretrained(
self.model_name, clean_up_tokenization_spaces=True
)
- self.model = AutoModel.from_pretrained(self.model_name).to(get_device())
+ self.model = AutoModel.from_pretrained(self.model_name)
+
+ self.model = self.model.to(self.device)
self.bert_output = None
def get_activation(name):
@@ -67,31 +188,65 @@ class BertFamily(UnimodalRepresentation):
return hook
+ aggregate_dim = (0,)
if self.layer_name != "cls":
for name, layer in self.model.named_modules():
if name == self.layer_name:
layer.register_forward_hook(get_activation(name))
break
-
if ModalityType.TEXT.has_field(modality.metadata, "text_spans"):
dataset = TextSpanDataset(modality.data, modality.metadata)
embeddings = []
+ aggregate_dim = (0, 1)
for text in dataset:
- embedding = self.create_embeddings(text, self.model, tokenizer)
- embeddings.append(embedding)
+ embedding = self.create_embeddings(
+ text, self.model, tokenizer, aggregation
+ )
+ embeddings.append(
+ aggregation.execute(embedding)
+ if aggregation is not None
+ else embedding
+ )
else:
- embeddings = self.create_embeddings(modality.data, self.model, tokenizer)
+ embeddings = self.create_embeddings(
+ modality.data, self.model, tokenizer, aggregation
+ )
if self.output_file is not None:
save_embeddings(embeddings, self.output_file)
transformed_modality.data_type = np.float32
+ transformed_modality.aggregate_dim = aggregate_dim
transformed_modality.data = embeddings
+ self.assert_output_stats(transformed_modality)
return transformed_modality
- def create_embeddings(self, data, model, tokenizer):
+ def assert_output_stats(self, transformed_modality):
+ if self.stats:
+ assert len(transformed_modality.data) == self.stats.num_instances
+ if len(self.stats.output_shape) == 3:
+ assert (
+ transformed_modality.data[0].shape[0] <= self.stats.output_shape[0]
+ ), f"Output shape: {transformed_modality.data[0].shape}, Expected shape: {self.stats.output_shape}"
+ assert (
+ transformed_modality.data[0].shape[1] == self.stats.output_shape[1]
+ ), f"Output shape: {transformed_modality.data[0].shape}, Expected shape: {self.stats.output_shape}"
+ assert (
+ transformed_modality.data[0].shape[2] == self.stats.output_shape[2]
+ ), f"Output shape: {transformed_modality.data[0].shape}, Expected shape: {self.stats.output_shape}"
+ else:
+ assert (
+ transformed_modality.data[0].shape[0] == self.stats.output_shape[0]
+ ), f"Output shape: {transformed_modality.data[0].shape}, Expected shape: {self.stats.output_shape}"
+ assert (
+ transformed_modality.data[0].shape[1] == self.stats.output_shape[1]
+ ), f"Output shape: {transformed_modality.data[0].shape}, Expected shape: {self.stats.output_shape}"
+
+ def create_embeddings(self, data, model, tokenizer, aggregation=None):
dataset = TextDataset(data)
- dataloader = DataLoader(dataset, batch_size=32, shuffle=False, collate_fn=None)
+ dataloader = DataLoader(
+ dataset, batch_size=self.batch_size, shuffle=False, collate_fn=None
+ )
cls_embeddings = []
for batch in dataloader:
inputs = tokenizer(
@@ -101,10 +256,9 @@ class BertFamily(UnimodalRepresentation):
padding="max_length",
return_attention_mask=True,
truncation=True,
- max_length=512, # TODO: make this dynamic with parameter to tune
+ max_length=self.max_seq_length, # TODO: make this dynamic with parameter to tune
)
-
- inputs.to(get_device())
+ inputs.to(self.device)
# ModalityType.TEXT.add_field_for_instances(
# modality.metadata,
# "token_to_character_mapping",
@@ -123,15 +277,26 @@ class BertFamily(UnimodalRepresentation):
if self.layer_name == "cls":
cls_embedding = outputs.last_hidden_state.detach().cpu().numpy()
else:
- cls_embedding = self.bert_output
+ cls_embedding = self.bert_output.cpu().numpy()
+ if aggregation is not None:
+ cls_embedding = aggregation.execute(cls_embedding)
cls_embeddings.extend(cls_embedding)
- return np.array(cls_embeddings)
+ return cls_embeddings
@register_representation(ModalityType.TEXT)
class Bert(BertFamily):
- def __init__(self, layer="cls", output_file=None, max_seq_length=512):
+ def __init__(
+ self,
+ layer="cls",
+ output_file=None,
+ max_seq_length=512,
+ batch_size=32,
+ params=None,
+ ):
+ self.set_parameters(params, max_seq_length, batch_size, layer, output_file)
+
parameters = {
"layer_name": [
"cls",
@@ -152,13 +317,30 @@ class Bert(BertFamily):
]
}
super().__init__(
- "Bert", "bert-base-uncased", layer, parameters, output_file, max_seq_length
+ "Bert",
+ "bert-base-uncased",
+ layer,
+ parameters,
+ output_file,
+ max_seq_length,
+ batch_size,
+ params=params,
)
+ self.C = 0.3
-@register_representation(ModalityType.TEXT)
+# @register_representation(ModalityType.TEXT)
class RoBERTa(BertFamily):
- def __init__(self, layer="cls", output_file=None, max_seq_length=512):
+ def __init__(
+ self,
+ layer="cls",
+ output_file=None,
+ max_seq_length=512,
+ batch_size=32,
+ params=None,
+ ):
+ self.set_parameters(params, max_seq_length, batch_size, layer, output_file)
+
parameters = {
"layer_name": [
"cls",
@@ -179,13 +361,30 @@ class RoBERTa(BertFamily):
]
}
super().__init__(
- "RoBERTa", "roberta-base", layer, parameters, output_file, max_seq_length
+ "RoBERTa",
+ "roberta-base",
+ layer,
+ parameters,
+ output_file,
+ max_seq_length,
+ batch_size,
+ params=params,
)
+ self.C = 0.3
-@register_representation(ModalityType.TEXT)
+# @register_representation(ModalityType.TEXT)
class DistillBERT(BertFamily):
- def __init__(self, layer="cls", output_file=None, max_seq_length=512):
+ def __init__(
+ self,
+ layer="cls",
+ output_file=None,
+ max_seq_length=512,
+ batch_size=32,
+ params=None,
+ ):
+ self.set_parameters(params, max_seq_length, batch_size, layer, output_file)
+
parameters = {
"layer_name": [
"cls",
@@ -204,21 +403,47 @@ class DistillBERT(BertFamily):
parameters,
output_file,
max_seq_length,
+ batch_size,
+ params=params,
)
+ self.C = 0.5
-@register_representation(ModalityType.TEXT)
+# @register_representation(ModalityType.TEXT)
class ALBERT(BertFamily):
- def __init__(self, layer="cls", output_file=None, max_seq_length=512):
+ def __init__(
+ self,
+ layer="cls",
+ output_file=None,
+ max_seq_length=512,
+ batch_size=32,
+ params=None,
+ ):
+ self.set_parameters(params, max_seq_length, batch_size, layer, output_file)
parameters = {"layer_name": ["cls", "encoder.albert_layer_groups.0", "pooler"]}
super().__init__(
- "ALBERT", "albert-base-v2", layer, parameters, output_file, max_seq_length
+ "ALBERT",
+ "albert-base-v2",
+ layer,
+ parameters,
+ output_file,
+ max_seq_length,
+ batch_size,
+ params=params,
)
+ self.C = 0.4
-@register_representation(ModalityType.TEXT)
+# @register_representation(ModalityType.TEXT)
class ELECTRA(BertFamily):
- def __init__(self, layer="cls", output_file=None, max_seq_length=512):
+ def __init__(
+ self,
+ layer="cls",
+ output_file=None,
+ max_seq_length=512,
+ batch_size=32,
+ params=None,
+ ):
parameters = {
"layer_name": [
"cls",
@@ -236,6 +461,7 @@ class ELECTRA(BertFamily):
"encoder.layer.11",
]
}
+ self.set_parameters(params, max_seq_length, batch_size, layer, output_file)
super().__init__(
"ELECTRA",
"google/electra-base-discriminator",
@@ -243,4 +469,7 @@ class ELECTRA(BertFamily):
parameters,
output_file,
max_seq_length,
+ batch_size,
+ params=params,
)
+ self.C = 0.5
diff --git a/src/main/python/systemds/scuro/representations/bow.py b/src/main/python/systemds/scuro/representations/bow.py
index 9d1d82a6be..65e490972a 100644
--- a/src/main/python/systemds/scuro/representations/bow.py
+++ b/src/main/python/systemds/scuro/representations/bow.py
@@ -22,34 +22,67 @@ import numpy as np
from sklearn.feature_extraction.text import CountVectorizer
from systemds.scuro.modality.transformed import TransformedModality
+from systemds.scuro.representations.representation import RepresentationStats
from systemds.scuro.representations.unimodal import UnimodalRepresentation
from systemds.scuro.representations.utils import save_embeddings
from systemds.scuro.modality.type import ModalityType
from systemds.scuro.drsearch.operator_registry import register_representation
+from systemds.scuro.dataloader.text_loader import TextStats
@register_representation(ModalityType.TEXT)
class BoW(UnimodalRepresentation):
- def __init__(self, ngram_range=2, min_df=2, output_file=None):
+ def __init__(self, ngram_range=2, min_df=2, output_file=None, params=None):
parameters = {"ngram_range": [2, 3, 5, 10], "min_df": [1, 2, 4, 8]}
super().__init__("BoW", ModalityType.EMBEDDING, parameters)
self.ngram_range = int(ngram_range)
self.min_df = int(min_df)
self.output_file = output_file
+ self.data_type = np.float32
- def transform(self, modality):
+ def get_output_stats(self, input_stats: TextStats) -> RepresentationStats:
+ vocab_estimate = min(
+ 100_000,
+ max(
+ 1000,
+ input_stats.num_instances * input_stats.max_length * self.ngram_range,
+ ),
+ )
+ return RepresentationStats(
+ input_stats.num_instances, (vocab_estimate,), output_shape_is_known=False
+ )
+
+ def estimate_output_memory_bytes(self, input_stats: TextStats) -> int:
+ output_bytes = 1
+ output_shape = self.get_output_stats(input_stats).output_shape
+ for dim in output_shape:
+ output_bytes *= dim
+ return (
+ input_stats.num_instances * output_bytes * np.dtype(self.data_type).itemsize
+ )
+
+ def estimate_peak_memory_bytes(self, input_stats: TextStats) -> dict:
+ output_bytes = self.estimate_output_memory_bytes(input_stats)
+ vectorizer_overhead = 640 * 1024 # 640 KB
+ return {
+ "cpu_peak_bytes": output_bytes + vectorizer_overhead,
+ "gpu_peak_bytes": 0,
+ }
+
+ def transform(self, modality, aggregation=None):
transformed_modality = TransformedModality(modality, self)
vectorizer = CountVectorizer(
ngram_range=(1, self.ngram_range), min_df=self.min_df
)
X = vectorizer.fit_transform(modality.data).toarray()
- X = [X[i : i + 1] for i in range(X.shape[0])]
if self.output_file is not None:
save_embeddings(X, self.output_file)
- transformed_modality.data_type = np.float32
+ transformed_modality.data_type = self.data_type
transformed_modality.data = X
+ if aggregation is not None:
+ transformed_modality.data = aggregation.execute(transformed_modality.data)
return transformed_modality
diff --git a/src/main/python/systemds/scuro/representations/clip.py b/src/main/python/systemds/scuro/representations/clip.py
index a431e52761..b13d2dfeb8 100644
--- a/src/main/python/systemds/scuro/representations/clip.py
+++ b/src/main/python/systemds/scuro/representations/clip.py
@@ -22,6 +22,7 @@ import numpy as np
from torchvision import transforms
from systemds.scuro.modality.transformed import TransformedModality
+from systemds.scuro.representations.representation import RepresentationStats
from systemds.scuro.representations.unimodal import UnimodalRepresentation
import torch
from systemds.scuro.representations.utils import save_embeddings
@@ -31,21 +32,143 @@ from transformers import CLIPProcessor, CLIPModel
from systemds.scuro.utils.converter import numpy_dtype_to_torch_dtype
from systemds.scuro.utils.static_variables import get_device
-from systemds.scuro.utils.torch_dataset import CustomDataset
+from systemds.scuro.utils.torch_dataset import (
+ CustomDataset,
+ TextDataset,
+ TextSpanDataset,
+)
+from systemds.scuro.utils.static_variables import (
+ get_device,
+ PY_LIST_HEADER_BYTES,
+ PY_LIST_SLOT_BYTES,
+ NP_ARRAY_HEADER_BYTES,
+)
+from torch.utils.data import DataLoader
@register_representation([ModalityType.VIDEO, ModalityType.IMAGE])
class CLIPVisual(UnimodalRepresentation):
- def __init__(self, output_file=None):
+ def __init__(self, output_file=None, batch_size=32, params=None):
parameters = {}
super().__init__("CLIPVisual", ModalityType.EMBEDDING, parameters)
- self.model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(
- get_device()
- )
+ self.model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
self.processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
self.output_file = output_file
+ self.data_type = torch.float32
+ self.batch_size = batch_size
+ self.gpu_id = None
+ self.device = get_device()
+
+ @property
+ def gpu_id(self):
+ return self._gpu_id
+
+ @gpu_id.setter
+ def gpu_id(self, gpu_id):
+ self._gpu_id = gpu_id
+ self.device = get_device(gpu_id)
+
+ def estimate_output_memory_bytes(self, input_stats) -> int:
+ return input_stats.num_instances * 512 * self.data_type.itemsize
+
+ def get_output_stats(self, input_stats) -> RepresentationStats:
+ if not isinstance(input_stats, RepresentationStats):
+ return RepresentationStats(input_stats.num_instances, (512,))
+ else:
+ return RepresentationStats(
+ input_stats.num_instances,
+ (input_stats.output_shape[0], 512),
+ )
+
+ def estimate_peak_memory_bytes(self, input_stats) -> dict:
+ CPU_RUNTIME_OVERHEAD = 100 * 1024 * 1024
+ GPU_RUNTIME_OVERHEAD = 80 * 1024 * 1024
+
+ EMB_DIM = 512
+ out_dtype = np.float32
+ out_dtype_size = np.dtype(out_dtype).itemsize
+
+ batch_size = int(self.batch_size)
+
+ n = int(getattr(input_stats, "num_instances", 1))
+ max_h = int(getattr(input_stats, "max_height", 224))
+ max_w = int(getattr(input_stats, "max_width", 224))
+ max_c = int(
+ getattr(
+ input_stats, "max_channels", getattr(input_stats, "max_num_channels", 3)
+ )
+ )
+ max_frames = int(getattr(input_stats, "max_length", 1))
+
+ model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
+ model_bytes = int(model.get_memory_footprint())
+
+ per_image_payload = EMB_DIM * out_dtype_size
+ per_image_item = per_image_payload + NP_ARRAY_HEADER_BYTES + PY_LIST_SLOT_BYTES
+ image_outputs_retained = PY_LIST_HEADER_BYTES + n * per_image_item
+
+ per_video_payload = max_frames * EMB_DIM * out_dtype_size
+ per_video_item = per_video_payload + NP_ARRAY_HEADER_BYTES + PY_LIST_SLOT_BYTES
+ video_outputs_retained = PY_LIST_HEADER_BYTES + n * per_video_item
+
+ is_video_like = hasattr(input_stats, "max_length")
+ outputs_retained = (
+ video_outputs_retained if is_video_like else image_outputs_retained
+ )
+
+ batch_pixels_cpu = batch_size * 3 * 224 * 224 * out_dtype_size
+ cpu_processor_workspace = int(2.5 * batch_pixels_cpu)
+
+ cpu_raw_batch = batch_size * max_h * max_w * max_c * out_dtype_size
+
+ cpu_batch_output = batch_size * EMB_DIM * out_dtype_size
+ cpu_batch_output_path = int(2.0 * cpu_batch_output)
- def transform(self, modality):
+ cpu_video_instance_tmp = 0
+ if is_video_like:
+ cpu_video_instance_tmp = int(
+ 2.0 * per_video_payload
+ + PY_LIST_HEADER_BYTES
+ + max_frames * PY_LIST_SLOT_BYTES
+ )
+
+ cpu_transient = (
+ cpu_raw_batch
+ + cpu_processor_workspace
+ + cpu_batch_output_path
+ + cpu_video_instance_tmp
+ + CPU_RUNTIME_OVERHEAD
+ )
+
+ cpu_peak = model_bytes + outputs_retained + cpu_transient
+
+ cfg = model.vision_model.config
+ hidden_size = int(cfg.hidden_size)
+ num_layers = int(cfg.num_hidden_layers)
+ patch = int(cfg.patch_size)
+ seq_len = (224 // patch) ** 2 + 1
+
+ gpu_input = batch_size * 3 * 224 * 224 * out_dtype_size
+
+ per_layer_act = batch_size * seq_len * hidden_size * out_dtype_size
+ gpu_activations = int(num_layers * per_layer_act * 2)
+
+ gpu_output = batch_size * EMB_DIM * out_dtype_size
+
+ gpu_peak = (
+ model_bytes
+ + gpu_input
+ + gpu_activations
+ + gpu_output
+ + GPU_RUNTIME_OVERHEAD
+ )
+
+ return {
+ "cpu_peak_bytes": int(cpu_peak),
+ "gpu_peak_bytes": int(gpu_peak),
+ }
+
+ def transform(self, modality, aggregation=None):
transformed_modality = TransformedModality(
modality, self, self.output_modality_type
)
@@ -53,12 +176,14 @@ class CLIPVisual(UnimodalRepresentation):
if next(self.model.parameters()).dtype != self.data_type:
self.model = self.model.to(self.data_type)
+ self.model = self.model.to(self.device)
+
embeddings = self.create_visual_embeddings(modality)
if self.output_file is not None:
save_embeddings(embeddings, self.output_file)
- transformed_modality.data = list(embeddings.values())
+ transformed_modality.data = embeddings
return transformed_modality
def create_visual_embeddings(self, modality):
@@ -72,24 +197,49 @@ class CLIPVisual(UnimodalRepresentation):
transforms.ConvertImageDtype(dtype=self.data_type),
]
)
- dataset = CustomDataset(
- modality.data, self.data_type, get_device(), tf=clip_transform
- )
+ dataset = CustomDataset(modality.data, self.data_type, "cpu", tf=clip_transform)
embeddings = {}
+ if modality.modality_type == ModalityType.IMAGE:
+ embeddings = []
+ for batch in torch.utils.data.DataLoader(
+ dataset, batch_size=self.batch_size
+ ):
+ images = batch["data"]
+ inputs = self.processor(
+ images=images, return_tensors="pt", do_rescale=False
+ )
+ inputs.to(self.device)
+
+ with torch.no_grad():
+ output = self.model.get_image_features(**inputs)
+ if len(output.shape) > 2:
+ output = torch.nn.functional.adaptive_avg_pool2d(output, (1, 1))
+ embeddings.extend(
+ torch.flatten(output, 1)
+ .detach()
+ .cpu()
+ .float()
+ .numpy()
+ .astype(np.float32)
+ )
+ return embeddings
+
for instance in torch.utils.data.DataLoader(dataset):
id = int(instance["id"][0])
frames = instance["data"][0]
embeddings[id] = []
- batch_size = 64
+ batch_size = self.batch_size
for start_index in range(0, len(frames), batch_size):
end_index = min(start_index + batch_size, len(frames))
frame_ids_range = range(start_index, end_index)
frame_batch = frames[frame_ids_range]
- inputs = self.processor(images=frame_batch, return_tensors="pt")
- inputs.to(get_device())
+ inputs = self.processor(
+ images=frame_batch, return_tensors="pt", do_rescale=False
+ )
+ inputs.to(self.device)
with torch.no_grad():
output = self.model.get_image_features(**inputs)
@@ -106,33 +256,139 @@ class CLIPVisual(UnimodalRepresentation):
)
embeddings[id] = np.array(embeddings[id])
- return embeddings
+ return list(embeddings.values())
@register_representation(ModalityType.TEXT)
class CLIPText(UnimodalRepresentation):
- def __init__(self, output_file=None):
- parameters = {}
+ def __init__(self, output_file=None, batch_size=32, params=None):
+ self.batch_size = batch_size
+ self.max_seq_length = 77
+ parameters = {"batch_size": [1, 2, 4, 8, 16, 32, 64, 128]}
+
super().__init__("CLIPText", ModalityType.EMBEDDING, parameters)
- self.model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(
- get_device()
- )
- self.processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
+ self.model = None
+ self.processor = None
self.output_file = output_file
self.needs_context = True
self.initial_context_length = 55
+ self.data_type = torch.float32
+ self.gpu_id = None
+ self.device = get_device()
+ self.params = params
+
+ @property
+ def gpu_id(self):
+ return self._gpu_id
+
+ @gpu_id.setter
+ def gpu_id(self, gpu_id):
+ self._gpu_id = gpu_id
+ self.device = get_device(gpu_id)
- def transform(self, modality):
+ def estimate_output_memory_bytes(self, input_stats) -> int:
+ output_stats = self.get_output_stats(input_stats).output_shape
+ return int(
+ input_stats.num_instances * np.prod(output_stats) * self.data_type.itemsize
+ )
+
+ def get_output_stats(self, input_stats) -> RepresentationStats:
+ if not isinstance(input_stats, RepresentationStats):
+ self.stats = RepresentationStats(
+ input_stats.num_instances, (512,), aggregate_dim=(0,)
+ )
+ else:
+ self.stats = RepresentationStats(
+ input_stats.num_instances,
+ (input_stats.output_shape[0], 512),
+ aggregate_dim=(
+ 0,
+ 1,
+ ),
+ )
+ if self.params and "_pushdown_aggregation" in self.params:
+ output_shape = (512,)
+ self.stats.output_shape = output_shape
+ self.stats.aggregate_dim = None
+ return self.stats
+
+ def estimate_peak_memory_bytes(self, input_stats) -> dict:
+ output_bytes = self.estimate_output_memory_bytes(input_stats)
+
+ model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
+
+ cfg = model.text_model.config
+ hidden_size = cfg.hidden_size
+ num_layers = cfg.num_hidden_layers
+ intermediate_size = getattr(cfg, "intermediate_size", 4 * hidden_size)
+ num_heads = cfg.num_attention_heads
+ dtype_size = self.data_type.itemsize
+ batch_tokens = self.batch_size * self.max_seq_length
+ hidden_ffn_bytes = (
+ batch_tokens * (hidden_size + intermediate_size) * dtype_size * num_layers
+ )
+ attn_matrix_bytes = (
+ self.batch_size
+ * num_heads
+ * self.max_seq_length
+ * self.max_seq_length
+ * dtype_size
+ * num_layers
+ )
+ activation_scale = 0.6
+ activations_bytes = int(
+ (hidden_ffn_bytes + attn_matrix_bytes) * activation_scale
+ )
+
+ batch_peak_bytes = self.batch_size * self.max_seq_length * 8 * 3
+
+ if isinstance(input_stats, RepresentationStats):
+ per_instance_input_bytes = (
+ int(np.prod(input_stats.output_shape)) * self.data_type.itemsize
+ )
+ input_bytes_all_instances = per_instance_input_bytes
+ else:
+ per_instance_input_bytes = (
+ int(np.prod(input_stats.output_shape)) * self.data_type.itemsize
+ )
+ input_bytes_all_instances = self.batch_size * per_instance_input_bytes
+ batch_output_bytes = self.batch_size * 512 * np.dtype(np.float32).itemsize
+ cpu_peak = (
+ model.get_memory_footprint()
+ + 100 * 1024 * 1024
+ + output_bytes
+ + batch_peak_bytes
+ + batch_output_bytes
+ + input_bytes_all_instances
+ )
+ gpu_peak = (
+ model.get_memory_footprint()
+ + batch_peak_bytes
+ + activations_bytes
+ + batch_output_bytes
+ )
+ return {"cpu_peak_bytes": cpu_peak, "gpu_peak_bytes": gpu_peak}
+
+ def transform(self, modality, aggregation=None):
transformed_modality = TransformedModality(
modality, self, self.output_modality_type
)
+ self.processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
+ self.model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
+ self.model = self.model.to(self.device)
- if isinstance(modality.data[0], list):
+ if ModalityType.TEXT.has_field(modality.metadata, "text_spans"):
+ dataset = TextSpanDataset(modality.data, modality.metadata)
embeddings = []
- for d in modality.data:
- embeddings.append(self.create_text_embeddings(d, self.model))
+ for text_chunks in dataset:
+ embedding = self.create_text_embeddings(
+ text_chunks, self.model, aggregation
+ )
+ embeddings.append(embedding)
else:
- embeddings = self.create_text_embeddings(modality.data, self.model)
+ embeddings = self.create_text_embeddings(
+ modality.data, self.model, aggregation
+ )
if self.output_file is not None:
save_embeddings(embeddings, self.output_file)
@@ -140,21 +396,28 @@ class CLIPText(UnimodalRepresentation):
transformed_modality.data = embeddings
return transformed_modality
- def create_text_embeddings(self, data, model):
+ def create_text_embeddings(self, data, model, aggregation=None):
+ dataset = TextDataset(data)
+ dataloader = DataLoader(
+ dataset, batch_size=self.batch_size, shuffle=False, collate_fn=None
+ )
embeddings = []
- for d in data:
+ for batch in dataloader:
inputs = self.processor(
- text=d,
+ text=batch,
return_tensors="pt",
padding=True,
truncation=True,
max_length=77,
)
- inputs.to(get_device())
+ inputs.to(self.device)
with torch.no_grad():
- text_embedding = model.get_text_features(**inputs)
- embeddings.append(
- text_embedding.squeeze().detach().cpu().numpy().reshape(1, -1)
- )
+ text_features = model.get_text_features(**inputs)
+
+ batch_np = text_features.detach().cpu().float().numpy()
+ if aggregation is not None:
+ batch_np = aggregation.execute(batch_np)
+
+ embeddings.extend(batch_np)
return embeddings
diff --git a/src/main/python/systemds/scuro/representations/color_histogram.py b/src/main/python/systemds/scuro/representations/color_histogram.py
index 2d780939e1..d1c7175b16 100644
--- a/src/main/python/systemds/scuro/representations/color_histogram.py
+++ b/src/main/python/systemds/scuro/representations/color_histogram.py
@@ -22,10 +22,17 @@
import numpy as np
import cv2
+from systemds.scuro.dataloader.image_loader import ImageStats
from systemds.scuro.drsearch.operator_registry import register_representation
from systemds.scuro.modality.type import ModalityType
+from systemds.scuro.representations.representation import RepresentationStats
from systemds.scuro.representations.unimodal import UnimodalRepresentation
from systemds.scuro.modality.transformed import TransformedModality
+from systemds.scuro.utils.static_variables import (
+ PY_LIST_HEADER_BYTES,
+ PY_LIST_SLOT_BYTES,
+ NP_ARRAY_HEADER_BYTES,
+)
@register_representation(ModalityType.IMAGE)
@@ -37,6 +44,7 @@ class ColorHistogram(UnimodalRepresentation):
normalize=False,
aggregation="mean",
output_file=None,
+ params=None,
):
super().__init__(
"ColorHistogram", ModalityType.EMBEDDING, self._get_parameters()
@@ -46,6 +54,7 @@ class ColorHistogram(UnimodalRepresentation):
self.normalize = normalize
self.aggregation = aggregation
self.output_file = output_file
+ self.data_type = np.float32
def _get_parameters(self):
return {
@@ -84,7 +93,56 @@ class ColorHistogram(UnimodalRepresentation):
hist /= hist_sum
return hist.astype(np.float32)
- def transform(self, modality):
+ def estimate_output_memory_bytes(self, input_stats: ImageStats) -> int:
+ return (
+ input_stats.num_instances * self.bins**3 * np.dtype(self.data_type).itemsize
+ )
+
+ def calculate_hist_dim(self):
+ num_channels = 1 if self.color_space == "GRAY" else 3
+ if isinstance(self.bins, (tuple, list)):
+ hist_dim = 1
+ for b in self.bins:
+ hist_dim *= int(b)
+ return hist_dim * num_channels
+ else:
+ return int(self.bins) ** num_channels
+
+ def get_output_stats(self, input_stats) -> RepresentationStats:
+ return RepresentationStats(
+ input_stats.num_instances, (self.calculate_hist_dim(),)
+ )
+
+ def estimate_peak_memory_bytes(self, input_stats: ImageStats) -> dict:
+ elem_size = np.dtype(self.data_type).itemsize
+ n = int(input_stats.num_instances)
+
+ hist_payload_bytes = self.calculate_hist_dim() * elem_size
+ per_instance_retained = (
+ hist_payload_bytes + NP_ARRAY_HEADER_BYTES + PY_LIST_SLOT_BYTES
+ )
+ retained_output_bytes = PY_LIST_HEADER_BYTES + n * per_instance_retained
+ transient_hist_bytes = 3 * hist_payload_bytes
+ max_h = int(input_stats.max_height)
+ max_w = int(input_stats.max_width)
+ if self.color_space == "HSV":
+ cvt_tmp_bytes = max_h * max_w * 3 * np.dtype(np.uint8).itemsize
+ elif self.color_space == "GRAY":
+ cvt_tmp_bytes = max_h * max_w * 1 * np.dtype(np.uint8).itemsize
+ else:
+ cvt_tmp_bytes = 0
+
+ opencv_workspace_bytes = max(1 * 1024 * 1024, int(hist_payload_bytes))
+ transient_one_instance_bytes = (
+ transient_hist_bytes + cvt_tmp_bytes + opencv_workspace_bytes
+ )
+ cpu_peak_bytes = retained_output_bytes + transient_one_instance_bytes
+ return {
+ "cpu_peak_bytes": int(cpu_peak_bytes * 1.05),
+ "gpu_peak_bytes": 0,
+ }
+
+ def transform(self, modality, aggregation=None):
if modality.modality_type == ModalityType.IMAGE:
images = modality.data
hist_list = [self.compute_histogram(img) for img in images]
diff --git a/src/main/python/systemds/scuro/representations/concatenation.py b/src/main/python/systemds/scuro/representations/concatenation.py
index ea199d5827..a7ca905f47 100644
--- a/src/main/python/systemds/scuro/representations/concatenation.py
+++ b/src/main/python/systemds/scuro/representations/concatenation.py
@@ -27,13 +27,14 @@ from systemds.scuro.modality.modality import Modality
from systemds.scuro.representations.utils import pad_sequences
from systemds.scuro.representations.fusion import Fusion
+from systemds.scuro.representations.representation import RepresentationStats
from systemds.scuro.drsearch.operator_registry import register_fusion_operator
@register_fusion_operator()
class Concatenation(Fusion):
- def __init__(self):
+ def __init__(self, params=None):
"""
Combines modalities using concatenation
"""
@@ -72,3 +73,33 @@ class Concatenation(Fusion):
)
return np.array(data)
+
+ def get_output_stats(self, input_stats_list) -> RepresentationStats:
+ if isinstance(input_stats_list, RepresentationStats):
+ return input_stats_list
+
+ stats_list = list(input_stats_list)
+ if not stats_list:
+ return RepresentationStats(0, (0,))
+
+ num_instances = stats_list[0].num_instances
+ rank = len(stats_list[0].output_shape)
+
+ if rank == 1:
+ total_dim = sum(s.output_shape[0] for s in stats_list)
+ output_shape = (total_dim,)
+ elif rank == 2:
+ time_dim = stats_list[0].output_shape[0]
+ total_dim = sum(s.output_shape[1] for s in stats_list)
+ output_shape = (time_dim, total_dim)
+ else:
+ output_shape = stats_list[0].output_shape
+
+ return RepresentationStats(num_instances, output_shape)
+
+ def estimate_peak_memory_bytes(self, input_stats) -> dict:
+ # TODO
+ return {
+ "cpu_peak_bytes": 0,
+ "gpu_peak_bytes": 0,
+ }
diff --git a/src/main/python/systemds/scuro/representations/context.py b/src/main/python/systemds/scuro/representations/context.py
index 27c118bcba..8438c9d1ef 100644
--- a/src/main/python/systemds/scuro/representations/context.py
+++ b/src/main/python/systemds/scuro/representations/context.py
@@ -49,9 +49,18 @@ class Context(Representation):
for parameter in list(self.parameters.keys()):
if self.is_ts_rep:
if parameter == "agg_params":
- current_params[parameter] = (
- self.aggregation_function.get_current_parameters()
- )
+ for (
+ agg_param,
+ agg_value,
+ ) in self.aggregation_function.get_current_parameters().items():
+ if isinstance(agg_value, dict):
+ for (
+ key,
+ value,
+ ) in agg_value.get_current_parameters().items():
+ current_params[f"agg_params_{agg_param}_{key}"] = value
+ else:
+ current_params[f"agg_params_{agg_param}"] = agg_value
continue
current_params[parameter] = getattr(self, parameter)
return current_params
diff --git a/src/main/python/systemds/scuro/representations/covarep_audio_features.py b/src/main/python/systemds/scuro/representations/covarep_audio_features.py
index 3b4398cb11..28dd6fc297 100644
--- a/src/main/python/systemds/scuro/representations/covarep_audio_features.py
+++ b/src/main/python/systemds/scuro/representations/covarep_audio_features.py
@@ -24,20 +24,24 @@ import numpy as np
from systemds.scuro.modality.type import ModalityType
from systemds.scuro.modality.transformed import TransformedModality
+from systemds.scuro.representations.representation import RepresentationStats
from systemds.scuro.representations.unimodal import UnimodalRepresentation
-from systemds.scuro.drsearch.operator_registry import register_representation
+from systemds.scuro.drsearch.operator_registry import (
+ register_representation,
+ register_context_representation_operator,
+)
@register_representation(ModalityType.AUDIO)
class Spectral(UnimodalRepresentation):
- def __init__(self, hop_length=512):
+ def __init__(self, hop_length=512, params=None):
parameters = {
"hop_length": [256, 512, 1024, 2048],
} # TODO
super().__init__("Spectral", ModalityType.EMBEDDING, parameters, False)
self.hop_length = hop_length
- def transform(self, modality):
+ def transform(self, modality, aggregation=None):
transformed_modality = TransformedModality(
modality, self, self.output_modality_type
)
@@ -73,17 +77,43 @@ class Spectral(UnimodalRepresentation):
return transformed_modality
+ def get_output_stats(self, input_stats) -> RepresentationStats:
+ num_instances = getattr(input_stats, "num_instances", 0)
+
+ # Try to infer signal length from stats
+ if hasattr(input_stats, "max_length"):
+ signal_length = input_stats.max_length
+ elif hasattr(input_stats, "output_shape") and input_stats.output_shape:
+ signal_length = input_stats.output_shape[0]
+ else:
+ signal_length = 0
+
+ if signal_length <= 0:
+ num_frames = 1
+ else:
+ num_frames = 1 + max(int((signal_length - 1) // self.hop_length), 0)
+ num_frames = max(int(num_frames), 1)
+
+ return RepresentationStats(num_instances, (num_frames, 4))
+
+ def estimate_peak_memory_bytes(self, input_stats) -> dict:
+ # TODO
+ return {
+ "cpu_peak_bytes": 0,
+ "gpu_peak_bytes": 0,
+ }
+
@register_representation(ModalityType.AUDIO)
class ZeroCrossing(UnimodalRepresentation):
- def __init__(self, hop_length=512):
+ def __init__(self, hop_length=512, params=None):
parameters = {
"hop_length": [256, 512, 1024, 2048],
} # TODO
super().__init__("ZeroCrossing", ModalityType.EMBEDDING, parameters, False)
self.hop_length = hop_length
- def transform(self, modality):
+ def transform(self, modality, aggregation=None):
transformed_modality = TransformedModality(
modality, self, self.output_modality_type
)
@@ -99,10 +129,35 @@ class ZeroCrossing(UnimodalRepresentation):
return transformed_modality
+ def get_output_stats(self, input_stats) -> RepresentationStats:
+ num_instances = getattr(input_stats, "num_instances", 0)
+
+ if hasattr(input_stats, "max_length"):
+ signal_length = input_stats.max_length
+ elif hasattr(input_stats, "output_shape") and input_stats.output_shape:
+ signal_length = input_stats.output_shape[0]
+ else:
+ signal_length = 0
+
+ if signal_length <= 0:
+ num_frames = 1
+ else:
+ num_frames = 1 + max(int((signal_length - 1) // self.hop_length), 0)
+ num_frames = max(int(num_frames), 1)
+
+ return RepresentationStats(num_instances, (num_frames, 1))
+
+ def estimate_peak_memory_bytes(self, input_stats) -> dict:
+ # TODO
+ return {
+ "cpu_peak_bytes": 0,
+ "gpu_peak_bytes": 0,
+ }
+
@register_representation(ModalityType.AUDIO)
class RMSE(UnimodalRepresentation):
- def __init__(self, frame_length=1024, hop_length=512):
+ def __init__(self, frame_length=1024, hop_length=512, params=None):
parameters = {
"frame_length": [1024, 2048, 4096],
"hop_length": [256, 512, 1024, 2048],
@@ -111,7 +166,7 @@ class RMSE(UnimodalRepresentation):
self.hop_length = hop_length
self.frame_length = frame_length
- def transform(self, modality):
+ def transform(self, modality, aggregation=None):
transformed_modality = TransformedModality(
modality, self, self.output_modality_type
)
@@ -126,17 +181,42 @@ class RMSE(UnimodalRepresentation):
return transformed_modality
+ def get_output_stats(self, input_stats) -> RepresentationStats:
+ num_instances = getattr(input_stats, "num_instances", 0)
+
+ if hasattr(input_stats, "max_length"):
+ signal_length = input_stats.max_length
+ elif hasattr(input_stats, "output_shape") and input_stats.output_shape:
+ signal_length = input_stats.output_shape[0]
+ else:
+ signal_length = 0
+
+ if signal_length <= 0:
+ num_frames = 1
+ else:
+ num_frames = 1 + max(int((signal_length - 1) // self.hop_length), 0)
+ num_frames = max(int(num_frames), 1)
+
+ return RepresentationStats(num_instances, (num_frames, 1))
+
+ def estimate_peak_memory_bytes(self, input_stats) -> dict:
+ # TODO
+ return {
+ "cpu_peak_bytes": 0,
+ "gpu_peak_bytes": 0,
+ }
+
@register_representation(ModalityType.AUDIO)
class Pitch(UnimodalRepresentation):
- def __init__(self, hop_length=512):
+ def __init__(self, hop_length=512, params=None):
parameters = {
"hop_length": [256, 512, 1024, 2048],
} # TODO
super().__init__("Pitch", ModalityType.EMBEDDING, parameters, False)
self.hop_length = hop_length
- def transform(self, modality):
+ def transform(self, modality, aggregation=None):
transformed_modality = TransformedModality(
modality, self, self.output_modality_type
)
@@ -154,3 +234,28 @@ class Pitch(UnimodalRepresentation):
transformed_modality.data = result
return transformed_modality
+
+ def get_output_stats(self, input_stats) -> RepresentationStats:
+ num_instances = getattr(input_stats, "num_instances", 0)
+
+ if hasattr(input_stats, "max_length"):
+ signal_length = input_stats.max_length
+ elif hasattr(input_stats, "output_shape") and input_stats.output_shape:
+ signal_length = input_stats.output_shape[0]
+ else:
+ signal_length = 0
+
+ if signal_length <= 0:
+ num_frames = 1
+ else:
+ num_frames = 1 + max(int((signal_length - 1) // self.hop_length), 0)
+ num_frames = max(int(num_frames), 1)
+
+ return RepresentationStats(num_instances, (num_frames, 1))
+
+ def estimate_peak_memory_bytes(self, input_stats) -> dict:
+ # TODO
+ return {
+ "cpu_peak_bytes": 0,
+ "gpu_peak_bytes": 0,
+ }
diff --git a/src/main/python/systemds/scuro/representations/dimensionality_reduction.py b/src/main/python/systemds/scuro/representations/dimensionality_reduction.py
index 71138b3641..731b3b9523 100644
--- a/src/main/python/systemds/scuro/representations/dimensionality_reduction.py
+++ b/src/main/python/systemds/scuro/representations/dimensionality_reduction.py
@@ -27,12 +27,12 @@ from systemds.scuro.representations.representation import Representation
class DimensionalityReduction(Representation):
- def __init__(self, name, parameters=None):
+ def __init__(self, name, params=None):
"""
Parent class for different dimensionality reduction operations
:param name: Name of the dimensionality reduction operator
"""
- super().__init__(name, parameters)
+ super().__init__(name, params)
self.needs_training = False
@abc.abstractmethod
diff --git a/src/main/python/systemds/scuro/representations/elmo.py b/src/main/python/systemds/scuro/representations/elmo.py
index 33e4f74141..85b04183d5 100644
--- a/src/main/python/systemds/scuro/representations/elmo.py
+++ b/src/main/python/systemds/scuro/representations/elmo.py
@@ -26,7 +26,7 @@ import torch.utils.data
import torch
import numpy as np
from systemds.scuro.modality.type import ModalityType
-from systemds.scuro.utils.static_variables import get_device
+from systemds.scuro.utils.static_variables import get_device_for_model
from flair.embeddings import ELMoEmbeddings
from flair.data import Sentence
from systemds.scuro.utils.torch_dataset import TextDataset
@@ -36,7 +36,12 @@ from torch.utils.data import DataLoader
# @register_representation([ModalityType.TEXT])
class ELMoRepresentation(UnimodalRepresentation):
def __init__(
- self, model_name="elmo-original", layer="mix", pooling="mean", output_file=None
+ self,
+ model_name="elmo-original",
+ layer="mix",
+ pooling="mean",
+ output_file=None,
+ params=None,
):
self.data_type = torch.float32
self.model_name = model_name
@@ -67,7 +72,8 @@ class ELMoRepresentation(UnimodalRepresentation):
else:
raise NotImplementedError(f"Model {model_name} not supported")
- self.model = self.model.to(get_device())
+ self.device = get_device_for_model(self.model, memory_factor=1.5)
+ self.model = self.model.to(self.device)
def _get_parameters(self):
parameters = {
@@ -82,7 +88,7 @@ class ELMoRepresentation(UnimodalRepresentation):
}
return parameters
- def transform(self, modality):
+ def transform(self, modality, aggregation=None):
transformed_modality = TransformedModality(
modality, self, ModalityType.EMBEDDING
)
diff --git a/src/main/python/systemds/scuro/representations/frequency_representation.py b/src/main/python/systemds/scuro/representations/frequency_representation.py
new file mode 100644
index 0000000000..52038cfb46
--- /dev/null
+++ b/src/main/python/systemds/scuro/representations/frequency_representation.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 numpy as np
+from scipy import signal
+
+from systemds.scuro.modality.type import ModalityType
+from systemds.scuro.representations.unimodal import UnimodalRepresentation
+from systemds.scuro.modality.transformed import TransformedModality
+
+
+class FrequencyRepresentation(UnimodalRepresentation):
+ def __init__(
+ self, fft_top_k=10, use_spectrogram=False, wavelet=None, output_file=None
+ ):
+ super().__init__(
+ "FrequencyRepresentation", ModalityType.EMBEDDING, self._get_parameters()
+ )
+ self.fft_top_k = fft_top_k
+ self.use_spectrogram = use_spectrogram
+ self.wavelet = wavelet
+ self.output_file = output_file
+
+ def _get_parameters(self):
+ return {
+ "fft_top_k": [5, 10, 20],
+ "use_spectrogram": [True, False],
+ "wavelet": [None, "db4", "haar", "coif1"],
+ }
+
+ def transform(self, modality, aggregation=None):
+ x = modality.data
+ output = []
+
+ fft_vals = np.fft.fft(x)
+ fft_freqs = np.fft.fftfreq(len(fft_vals))
+ idxs = np.argsort(np.abs(fft_vals))[-self.fft_top_k :]
+ fft_feats = np.stack(
+ [np.abs(fft_vals[idxs]), fft_freqs[idxs], np.angle(fft_vals[idxs])], axis=1
+ )
+ output.append(fft_feats.flatten())
+
+ if self.use_spectrogram:
+ f, t, Sxx = signal.spectrogram(x)
+ output.append(np.mean(Sxx, axis=1))
+ output.append(np.std(Sxx, axis=1))
+
+ if self.wavelet is not None:
+ import pywt
+
+ coeffs = pywt.wavedec(x, self.wavelet)
+ output.append(np.concatenate([c.flatten() for c in coeffs[:3]]))
+
+ transformed_modality = TransformedModality(
+ modality, self, ModalityType.EMBEDDING
+ )
+ transformed_modality.data = np.array(output)
+ return transformed_modality
diff --git a/src/main/python/systemds/scuro/representations/glove.py b/src/main/python/systemds/scuro/representations/glove.py
index 8f9a73d0d5..b45213ae19 100644
--- a/src/main/python/systemds/scuro/representations/glove.py
+++ b/src/main/python/systemds/scuro/representations/glove.py
@@ -23,7 +23,9 @@ import numpy as np
from gensim.utils import tokenize
from huggingface_hub import hf_hub_download
+from systemds.scuro.dataloader.text_loader import TextStats
from systemds.scuro.modality.transformed import TransformedModality
+from systemds.scuro.representations.representation import RepresentationStats
from systemds.scuro.representations.unimodal import UnimodalRepresentation
from systemds.scuro.representations.utils import save_embeddings
from systemds.scuro.modality.type import ModalityType
@@ -43,7 +45,7 @@ def load_glove_embeddings(file_path):
@register_representation(ModalityType.TEXT)
class GloVe(UnimodalRepresentation):
- def __init__(self, output_file=None):
+ def __init__(self, output_file=None, params=None):
super().__init__("GloVe", ModalityType.TEXT)
file_path = hf_hub_download(
repo_id="stanfordnlp/glove", filename="glove.6B.zip"
@@ -53,8 +55,31 @@ class GloVe(UnimodalRepresentation):
self.glove_path = "./glove_extracted/glove.6B.100d.txt"
self.output_file = output_file
+ self.data_type = np.float32
+ self.embedding_dim = 100
- def transform(self, modality):
+ def get_output_stats(self, input_stats: TextStats) -> RepresentationStats:
+ return RepresentationStats(input_stats.num_instances, (self.embedding_dim,))
+
+ def estimate_output_memory_bytes(self, input_stats: TextStats) -> int:
+ output_bytes = 1
+ output_shape = self.get_output_stats(input_stats).output_shape
+ for dim in output_shape:
+ output_bytes *= dim
+ return (
+ input_stats.num_instances * output_bytes * np.dtype(self.data_type).itemsize
+ )
+
+ def estimate_peak_memory_bytes(self, input_stats: TextStats) -> dict:
+ output_bytes = self.estimate_output_memory_bytes(input_stats)
+
+ glove_dict_bytes = 256 * 1024 * 1024
+ return {
+ "cpu_peak_bytes": glove_dict_bytes + output_bytes * 1.2,
+ "gpu_peak_bytes": 0,
+ }
+
+ def transform(self, modality, aggregation=None):
transformed_modality = TransformedModality(modality, self)
glove_embeddings = load_glove_embeddings(self.glove_path)
@@ -72,7 +97,7 @@ class GloVe(UnimodalRepresentation):
if len(token_embeddings) > 0:
embeddings.append(np.mean(token_embeddings, axis=0))
else:
- embeddings.append(np.zeros(embedding_dim, dtype=np.float32))
+ embeddings.append(np.zeros(embedding_dim, dtype=self.data_type))
if self.output_file is not None:
save_embeddings(np.array(embeddings), self.output_file)
diff --git a/src/main/python/systemds/scuro/representations/hadamard.py b/src/main/python/systemds/scuro/representations/hadamard.py
index a777768ff6..a29f63e2b2 100644
--- a/src/main/python/systemds/scuro/representations/hadamard.py
+++ b/src/main/python/systemds/scuro/representations/hadamard.py
@@ -26,13 +26,14 @@ import numpy as np
from systemds.scuro.modality.modality import Modality
from systemds.scuro.representations.fusion import Fusion
+from systemds.scuro.representations.representation import RepresentationStats
from systemds.scuro.drsearch.operator_registry import register_fusion_operator
@register_fusion_operator()
class Hadamard(Fusion):
- def __init__(self):
+ def __init__(self, params=None):
"""
Combines modalities using elementwise multiply (Hadamard product)
"""
@@ -45,3 +46,27 @@ class Hadamard(Fusion):
fused_data = np.prod([m.data for m in modalities], axis=0)
return fused_data
+
+ def get_output_stats(self, input_stats_list) -> RepresentationStats:
+ if isinstance(input_stats_list, RepresentationStats):
+ return input_stats_list
+
+ stats_list = list(input_stats_list)
+ if not stats_list:
+ return RepresentationStats(0, (0,))
+
+ def num_elements(stats: RepresentationStats) -> int:
+ n = 1
+ for d in stats.output_shape:
+ n *= d
+ return n
+
+ largest = max(stats_list, key=num_elements)
+ return RepresentationStats(largest.num_instances, largest.output_shape)
+
+ def estimate_peak_memory_bytes(self, input_stats) -> dict:
+ # TODO
+ return {
+ "cpu_peak_bytes": 0,
+ "gpu_peak_bytes": 0,
+ }
diff --git a/src/main/python/systemds/scuro/representations/image_bind.py b/src/main/python/systemds/scuro/representations/image_bind.py
index e934d521af..a96cffeb6f 100644
--- a/src/main/python/systemds/scuro/representations/image_bind.py
+++ b/src/main/python/systemds/scuro/representations/image_bind.py
@@ -50,7 +50,7 @@ class ImageBind(UnimodalRepresentation):
self.model.eval()
self.model.to(DEVICE)
- def transform(self, modality):
+ def transform(self, modality, aggregation=None):
transformed_modality = TransformedModality(
modality, self, ModalityType.EMBEDDING
)
diff --git a/src/main/python/systemds/scuro/representations/lstm.py b/src/main/python/systemds/scuro/representations/lstm.py
index 58b878820e..e067e2ccf9 100644
--- a/src/main/python/systemds/scuro/representations/lstm.py
+++ b/src/main/python/systemds/scuro/representations/lstm.py
@@ -25,7 +25,7 @@ import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset
from typing import List, Dict, Any
-from systemds.scuro.utils.static_variables import get_device
+from systemds.scuro.utils.static_variables import get_device, get_device_for_model
import numpy as np
from systemds.scuro.modality.modality import Modality
@@ -44,6 +44,7 @@ class LSTM(Fusion):
learning_rate=0.001,
epochs=20,
batch_size=32,
+ params=None,
):
parameters = {
"width": [128, 256, 512],
@@ -179,8 +180,8 @@ class LSTM(Fusion):
self.input_dim = X.shape[2]
self.model = self._build_model(self.input_dim, self.num_classes)
- device = get_device()
- self.model.to(device)
+ self.device = get_device_for_model(self.model, memory_factor=1.5)
+ self.model = self.model.to(self.device)
if self.is_multilabel:
criterion = nn.BCEWithLogitsLoss()
@@ -201,8 +202,8 @@ class LSTM(Fusion):
for epoch in range(self.epochs):
total_loss = 0
for batch_X, batch_y in dataloader:
- batch_X = batch_X.to(device)
- batch_y = batch_y.to(device)
+ batch_X = batch_X.to(self.device)
+ batch_y = batch_y.to(self.device)
optimizer.zero_grad()
features, predictions = self.model(batch_X)
@@ -244,8 +245,8 @@ class LSTM(Fusion):
X = self._prepare_data(modalities)
- device = get_device()
- self.model.to(device)
+ self.device = get_device_for_model(self.model, memory_factor=1.5)
+ self.model = self.model.to(self.device)
X_tensor = torch.FloatTensor(X)
all_features = []
@@ -255,7 +256,7 @@ class LSTM(Fusion):
TensorDataset(X_tensor), batch_size=self.batch_size, shuffle=False
)
for (batch_X,) in inference_dataloader:
- batch_X = batch_X.to(device)
+ batch_X = batch_X.to(self.device)
features, _ = self.model(batch_X)
all_features.append(features.cpu())
diff --git a/src/main/python/systemds/scuro/representations/mel_spectrogram.py b/src/main/python/systemds/scuro/representations/mel_spectrogram.py
index 6ea9061901..67ff955a98 100644
--- a/src/main/python/systemds/scuro/representations/mel_spectrogram.py
+++ b/src/main/python/systemds/scuro/representations/mel_spectrogram.py
@@ -24,40 +24,78 @@ import numpy as np
from systemds.scuro.modality.type import ModalityType
from systemds.scuro.modality.transformed import TransformedModality
+from systemds.scuro.representations.representation import RepresentationStats
from systemds.scuro.representations.unimodal import UnimodalRepresentation
-from systemds.scuro.drsearch.operator_registry import register_representation
+from systemds.scuro.drsearch.operator_registry import (
+ register_representation,
+ register_context_representation_operator,
+)
@register_representation(ModalityType.AUDIO)
+@register_context_representation_operator(ModalityType.AUDIO)
class MelSpectrogram(UnimodalRepresentation):
- def __init__(self, n_mels=128, hop_length=512, n_fft=2048):
+ def __init__(self, n_mels=128, hop_length=512, n_fft=2048, params=None):
parameters = {
"n_mels": [20, 32, 64, 128],
"hop_length": [256, 512, 1024, 2048],
"n_fft": [1024, 2048, 4096],
}
super().__init__("MelSpectrogram", ModalityType.TIMESERIES, parameters, False)
+
+ if params is not None:
+ n_mels = params.get("n_mels", n_mels)
+ hop_length = params.get("hop_length", hop_length)
+ n_fft = params.get("n_fft", n_fft)
+
self.n_mels = int(n_mels)
self.hop_length = int(hop_length)
self.n_fft = int(n_fft)
- def transform(self, modality):
+ def transform(self, modality, aggregation=None):
transformed_modality = TransformedModality(
modality, self, self.output_modality_type
)
result = []
+ metadata_values = list(modality.metadata.values())
for i, sample in enumerate(modality.data):
- sr = list(modality.metadata.values())[i]["frequency"]
- S = librosa.feature.melspectrogram(
- y=np.array(sample),
- sr=sr,
- n_mels=self.n_mels,
- hop_length=self.hop_length,
- n_fft=self.n_fft,
- ).astype(modality.data_type)
-
- result.append(S.T)
+ sr = metadata_values[i]["frequency"]
+ computed_feature = self.compute_feature(sample, sr)
+ result.append(computed_feature)
transformed_modality.data = result
return transformed_modality
+
+ def compute_feature(self, instance, sr=None):
+ if sr is None:
+ sr = 22050
+ S = librosa.feature.melspectrogram(
+ y=np.asarray(instance),
+ sr=sr,
+ n_mels=self.n_mels,
+ hop_length=self.hop_length,
+ n_fft=self.n_fft,
+ )
+ return S.T
+
+ def get_output_stats(self, input_stats) -> RepresentationStats:
+ num_instances = getattr(input_stats, "num_instances", 0)
+
+ if hasattr(input_stats, "max_length"):
+ signal_length = input_stats.max_length
+ elif hasattr(input_stats, "output_shape") and input_stats.output_shape:
+ signal_length = input_stats.output_shape[0]
+ else:
+ signal_length = 0
+
+ if signal_length <= 0:
+ num_frames = 1
+ else:
+ if signal_length < self.n_fft:
+ num_frames = 1
+ else:
+ num_frames = 1 + (signal_length - self.n_fft) // self.hop_length
+ num_frames = max(int(num_frames), 1)
+
+ return RepresentationStats(num_instances, (num_frames, self.n_mels))
diff --git a/src/main/python/systemds/scuro/representations/mfcc.py b/src/main/python/systemds/scuro/representations/mfcc.py
index 4d9989add9..5e5b54b079 100644
--- a/src/main/python/systemds/scuro/representations/mfcc.py
+++ b/src/main/python/systemds/scuro/representations/mfcc.py
@@ -24,13 +24,18 @@ import numpy as np
from systemds.scuro.modality.type import ModalityType
from systemds.scuro.modality.transformed import TransformedModality
+from systemds.scuro.representations.representation import RepresentationStats
from systemds.scuro.representations.unimodal import UnimodalRepresentation
-from systemds.scuro.drsearch.operator_registry import register_representation
+from systemds.scuro.drsearch.operator_registry import (
+ register_representation,
+ register_context_representation_operator,
+)
@register_representation(ModalityType.AUDIO)
+@register_context_representation_operator(ModalityType.AUDIO)
class MFCC(UnimodalRepresentation):
- def __init__(self, n_mfcc=12, dct_type=2, n_mels=128, hop_length=512):
+ def __init__(self, n_mfcc=12, dct_type=2, n_mels=128, hop_length=512, params=None):
parameters = {
"n_mfcc": [x for x in range(10, 26)],
"dct_type": [1, 2, 3],
@@ -38,12 +43,20 @@ class MFCC(UnimodalRepresentation):
"n_mels": [20, 32, 64, 128],
} # TODO
super().__init__("MFCC", ModalityType.TIMESERIES, parameters, False)
+
+ # Allow construction from a parameter dict (used by optimizer)
+ if params is not None:
+ n_mfcc = params.get("n_mfcc", n_mfcc)
+ dct_type = params.get("dct_type", dct_type)
+ n_mels = params.get("n_mels", n_mels)
+ hop_length = params.get("hop_length", hop_length)
+
self.n_mfcc = int(n_mfcc)
self.dct_type = int(dct_type)
self.n_mels = int(n_mels)
self.hop_length = int(hop_length)
- def transform(self, modality):
+ def transform(self, modality, aggregation=None):
transformed_modality = TransformedModality(
modality, self, self.output_modality_type
)
@@ -51,17 +64,40 @@ class MFCC(UnimodalRepresentation):
for i, sample in enumerate(modality.data):
sr = list(modality.metadata.values())[i]["frequency"]
- mfcc = librosa.feature.mfcc(
- y=np.array(sample),
- sr=sr,
- n_mfcc=self.n_mfcc,
- dct_type=self.dct_type,
- hop_length=self.hop_length,
- n_mels=self.n_mels,
- ).astype(modality.data_type)
- mfcc = (mfcc - np.mean(mfcc)) / np.std(mfcc)
-
- result.append(mfcc.T)
+ computed_feature = self.compute_feature(sample, sr)
+ result.append(computed_feature)
transformed_modality.data = result
return transformed_modality
+
+ def compute_feature(self, instance, sr=None):
+ if sr is None:
+ sr = 22050
+ mfcc = librosa.feature.mfcc(
+ y=np.array(instance),
+ sr=sr,
+ n_mfcc=self.n_mfcc,
+ dct_type=self.dct_type,
+ hop_length=self.hop_length,
+ n_mels=self.n_mels,
+ )
+ mfcc = (mfcc - np.mean(mfcc)) / np.std(mfcc)
+ return mfcc.T
+
+ def get_output_stats(self, input_stats) -> RepresentationStats:
+ num_instances = getattr(input_stats, "num_instances", 0)
+
+ if hasattr(input_stats, "max_length"):
+ signal_length = input_stats.max_length
+ elif hasattr(input_stats, "output_shape") and input_stats.output_shape:
+ signal_length = input_stats.output_shape[0]
+ else:
+ signal_length = 0
+
+ if signal_length <= 0:
+ num_frames = 1
+ else:
+ num_frames = 1 + max(int((signal_length - 1) // self.hop_length), 0)
+ num_frames = max(int(num_frames), 1)
+
+ return RepresentationStats(num_instances, (num_frames, self.n_mfcc))
diff --git a/src/main/python/systemds/scuro/representations/mlp_averaging.py b/src/main/python/systemds/scuro/representations/mlp_averaging.py
index 6935ab3721..bedfaf415b 100644
--- a/src/main/python/systemds/scuro/representations/mlp_averaging.py
+++ b/src/main/python/systemds/scuro/representations/mlp_averaging.py
@@ -26,7 +26,12 @@ import numpy as np
import warnings
from systemds.scuro.modality.type import ModalityType
-from systemds.scuro.utils.static_variables import get_device
+from systemds.scuro.representations.representation import RepresentationStats
+from systemds.scuro.utils.static_variables import (
+ compute_batch_size,
+ get_device,
+ get_device_for_model,
+)
from systemds.scuro.utils.utils import set_random_seeds
from systemds.scuro.drsearch.operator_registry import (
register_dimensionality_reduction_operator,
@@ -43,7 +48,7 @@ class MLPAveraging(DimensionalityReduction):
This operator is used to reduce the dimensionality of a representation using a simple average pooling operation.
"""
- def __init__(self, output_dim=512, batch_size=32):
+ def __init__(self, output_dim=512, batch_size=32, params=None):
parameters = {
"output_dim": [64, 128, 256, 512, 1024, 2048, 4096],
"batch_size": [8, 16, 32, 64, 128],
@@ -51,31 +56,112 @@ class MLPAveraging(DimensionalityReduction):
super().__init__("MLPAveraging", parameters)
self.output_dim = output_dim
self.batch_size = batch_size
+ self.device = None
+ self.data_type = np.float32
+ self.gpu_id = None
+
+ @property
+ def gpu_id(self):
+ return self._gpu_id
+
+ @gpu_id.setter
+ def gpu_id(self, gpu_id):
+ self._gpu_id = gpu_id
+ self.device = get_device(gpu_id)
+
+ def get_output_stats(self, input_stats: RepresentationStats) -> RepresentationStats:
+ if len(input_stats.output_shape) > 1:
+ return RepresentationStats(
+ input_stats.num_instances,
+ (self.output_dim,),
+ output_shape_is_known=True,
+ )
+ if (
+ len(input_stats.output_shape) == 1
+ and input_stats.output_shape[0] <= self.output_dim
+ ):
+ return RepresentationStats(
+ input_stats.num_instances,
+ (input_stats.output_shape[0],),
+ output_shape_is_known=input_stats.output_shape_is_known,
+ )
+ return RepresentationStats(
+ input_stats.num_instances,
+ (self.output_dim,),
+ output_shape_is_known=input_stats.output_shape_is_known,
+ )
+
+ def estimate_output_memory_bytes(self, input_stats: RepresentationStats) -> int:
+ output_bytes = 1
+ for dim in input_stats.output_shape:
+ output_bytes *= dim
+ return (
+ input_stats.num_instances * output_bytes * np.dtype(self.data_type).itemsize
+ )
+
+ def estimate_peak_memory_bytes(self, input_stats: RepresentationStats) -> dict:
+ n = int(input_stats.num_instances)
+ input_dim = int(np.prod(input_stats.output_shape))
+ elem_size = np.dtype(self.data_type).itemsize
+
+ if input_dim < self.output_dim or n == 0 or input_dim == 0:
+ input_bytes = n * input_dim * elem_size
+ cpu_peak = int(input_bytes * 1.05 + 8 * 1024**2) # small safety margin
+ return {"cpu_peak_bytes": cpu_peak, "gpu_peak_bytes": 0}
+
+ out_dim = int(self.output_dim)
+ batch = int(max(1, min(self.batch_size, n)))
+
+ input_bytes = n * input_dim * elem_size
+ output_bytes = n * out_dim * elem_size
+ weight_bytes = out_dim * input_dim * elem_size
+
+ batch_input_bytes = batch * input_dim * elem_size
+ batch_output_bytes = batch * out_dim * elem_size
+
+ num_batches = (n + batch - 1) // batch
+ python_overhead = num_batches * 1024
+
+ cpu_working = (
+ input_bytes
+ + 2 * output_bytes
+ + weight_bytes
+ + batch_input_bytes
+ + batch_output_bytes
+ + python_overhead
+ )
+ cpu_peak = int(cpu_working * 1.20 + 64 * 1024**2)
+
+ gpu_working = weight_bytes + batch_input_bytes + batch_output_bytes
+ gpu_peak = int(gpu_working * 1.35 + 560 * 1024**2)
+
+ return {"cpu_peak_bytes": cpu_peak, "gpu_peak_bytes": gpu_peak}
def execute(self, data):
set_random_seeds(42)
input_dim = data.shape[1]
- if input_dim < self.output_dim:
+ if input_dim <= self.output_dim:
warnings.warn(
f"Input dimension {input_dim} is smaller than output dimension {self.output_dim}. Returning original data."
) # TODO: this should be pruned as possible representation, could add output_dim as parameter to reps if possible
return data
- dim_reduction_model = AggregationMLP(input_dim, self.output_dim)
- dim_reduction_model.to(get_device())
+ dim_reduction_model = AggregationMLP(input_dim, self.output_dim).to(self.device)
dim_reduction_model.eval()
tensor_data = torch.from_numpy(data).float()
- dataset = TensorDataset(tensor_data)
+ dataset = TensorDataset(
+ tensor_data,
+ )
dataloader = DataLoader(dataset, batch_size=self.batch_size, shuffle=False)
all_features = []
with torch.no_grad():
for (batch,) in dataloader:
- batch_features = dim_reduction_model(batch.to(get_device()))
+ batch_features = dim_reduction_model(batch.to(self.device))
all_features.append(batch_features.cpu())
all_features = torch.cat(all_features, dim=0)
@@ -87,7 +173,7 @@ class AggregationMLP(nn.Module):
super(AggregationMLP, self).__init__()
agg_size = input_dim // output_dim
remainder = input_dim % output_dim
- weight = torch.zeros(output_dim, input_dim).to(get_device())
+ weight = torch.zeros(output_dim, input_dim)
start_idx = 0
for i in range(output_dim):
@@ -96,7 +182,7 @@ class AggregationMLP(nn.Module):
weight[i, start_idx:end_idx] = 1.0 / current_agg_size
start_idx = end_idx
- self.register_buffer("weight", weight)
+ self.register_buffer("weight", weight.T)
def forward(self, x):
- return torch.matmul(x, self.weight.T)
+ return torch.matmul(x, self.weight)
diff --git a/src/main/python/systemds/scuro/representations/mlp_learned_dim_reduction.py b/src/main/python/systemds/scuro/representations/mlp_learned_dim_reduction.py
index 5ea15c64d7..0cdc2b8fad 100644
--- a/src/main/python/systemds/scuro/representations/mlp_learned_dim_reduction.py
+++ b/src/main/python/systemds/scuro/representations/mlp_learned_dim_reduction.py
@@ -22,7 +22,11 @@ from torch.utils.data import DataLoader, TensorDataset
import numpy as np
import torch
import torch.nn as nn
-from systemds.scuro.utils.static_variables import get_device
+from systemds.scuro.utils.static_variables import (
+ compute_batch_size,
+ get_device,
+ get_device_for_model,
+)
from systemds.scuro.drsearch.operator_registry import (
register_dimensionality_reduction_operator,
@@ -82,12 +86,12 @@ class MLPLearnedDimReduction(DimensionalityReduction):
self.num_classes = len(np.unique(y))
input_dim = X.shape[1]
- device = get_device()
- self.model = None
+ self.device = get_device_for_model(self.model, memory_factor=1.5)
+ self.model = self.model.to(self.device)
self.is_trained = False
self.model = self._build_model(input_dim, self.output_dim, self.num_classes).to(
- device
+ self.device
)
if self.is_multilabel:
criterion = nn.BCEWithLogitsLoss()
@@ -101,6 +105,16 @@ class MLPLearnedDimReduction(DimensionalityReduction):
else:
y_tensor = torch.LongTensor(y)
+ sample = data[0] if data else ""
+ self.batch_size = compute_batch_size(
+ model=self.model,
+ device=self.device,
+ sample_data=sample,
+ tokenizer=None,
+ max_seq_length=None,
+ max_batch_size=128,
+ )
+
dataset = TensorDataset(X_tensor, y_tensor)
dataloader = DataLoader(dataset, batch_size=self.batch_size, shuffle=True)
@@ -108,8 +122,8 @@ class MLPLearnedDimReduction(DimensionalityReduction):
for epoch in range(self.epochs):
total_loss = 0
for batch_X, batch_y in dataloader:
- batch_X = batch_X.to(device)
- batch_y = batch_y.to(device)
+ batch_X = batch_X.to(self.device)
+ batch_y = batch_y.to(self.device)
optimizer.zero_grad()
features, predictions = self.model(batch_X)
@@ -128,7 +142,7 @@ class MLPLearnedDimReduction(DimensionalityReduction):
TensorDataset(X_tensor), batch_size=self.batch_size, shuffle=False
)
for (batch_X,) in inference_dataloader:
- batch_X = batch_X.to(device)
+ batch_X = batch_X.to(self.device)
features, _ = self.model(batch_X)
all_features.append(features.cpu())
@@ -138,8 +152,8 @@ class MLPLearnedDimReduction(DimensionalityReduction):
if not self.is_trained or self.model is None:
raise ValueError("Model must be trained before applying representation")
- device = get_device()
- self.model.to(device)
+ self.device = get_device_for_model(self.model, memory_factor=1.5)
+ self.model = self.model.to(self.device)
X = np.array(data)
X_tensor = torch.FloatTensor(X)
all_features = []
@@ -149,7 +163,7 @@ class MLPLearnedDimReduction(DimensionalityReduction):
TensorDataset(X_tensor), batch_size=self.batch_size, shuffle=False
)
for (batch_X,) in inference_dataloader:
- batch_X = batch_X.to(device)
+ batch_X = batch_X.to(self.device)
features, _ = self.model(batch_X)
all_features.append(features.cpu())
diff --git a/src/main/python/systemds/scuro/representations/optical_flow.py b/src/main/python/systemds/scuro/representations/optical_flow.py
index 27817302d4..943cd2ab49 100644
--- a/src/main/python/systemds/scuro/representations/optical_flow.py
+++ b/src/main/python/systemds/scuro/representations/optical_flow.py
@@ -42,11 +42,11 @@ else:
# @register_representation([ModalityType.VIDEO])
class OpticalFlow(UnimodalRepresentation):
- def __init__(self):
+ def __init__(self, params=None):
parameters = {}
super().__init__("OpticalFlow", ModalityType.TIMESERIES, parameters)
- def transform(self, modality):
+ def transform(self, modality, aggregation=None):
transformed_modality = TransformedModality(
modality, self, self.output_modality_type
)
diff --git a/src/main/python/systemds/scuro/representations/representation.py b/src/main/python/systemds/scuro/representations/representation.py
index 9f1a91a0ea..d83553ec6e 100644
--- a/src/main/python/systemds/scuro/representations/representation.py
+++ b/src/main/python/systemds/scuro/representations/representation.py
@@ -18,20 +18,34 @@
# under the License.
#
# -------------------------------------------------------------
+import abc
+from dataclasses import dataclass
from systemds.scuro.utils.identifier import Identifier
+@dataclass
+class RepresentationStats:
+ num_instances: int
+ output_shape: tuple
+ output_shape_is_known: bool = True
+ aggregate_dim: tuple = (0,)
+
+
class Representation:
def __init__(self, name, parameters):
self.name = name
self._parameters = parameters
self.self_contained = True
self.representation_id = Identifier().new_id()
+ self.stats = None
@property
def parameters(self):
return self._parameters
+ def get_stats(self):
+ return self.stats
+
def get_current_parameters(self):
current_params = {}
if not self.parameters:
@@ -44,3 +58,13 @@ class Representation:
def set_parameters(self, parameters):
for parameter in parameters:
setattr(self, parameter, parameters[parameter])
+
+ def estimate_memory_bytes(self, input_stats):
+ output_memory_bytes = self.estimate_output_memory_bytes(input_stats)
+ return output_memory_bytes
+
+ @abc.abstractmethod
+ def get_output_stats(self, input_stats):
+ raise NotImplementedError(
+ f"get_output_stats is not implemented for {self.name}"
+ )
diff --git a/src/main/python/systemds/scuro/representations/resnet.py b/src/main/python/systemds/scuro/representations/resnet.py
index 50fe084b9f..299ccad368 100644
--- a/src/main/python/systemds/scuro/representations/resnet.py
+++ b/src/main/python/systemds/scuro/representations/resnet.py
@@ -18,7 +18,8 @@
# under the License.
#
# -------------------------------------------------------------
-from systemds.scuro.utils.converter import numpy_dtype_to_torch_dtype
+from systemds.scuro.dataloader.image_loader import ImageStats
+from systemds.scuro.representations.representation import RepresentationStats
from systemds.scuro.utils.torch_dataset import CustomDataset
from systemds.scuro.modality.transformed import TransformedModality
from systemds.scuro.representations.unimodal import UnimodalRepresentation
@@ -34,9 +35,20 @@ from systemds.scuro.utils.static_variables import get_device
@register_representation([ModalityType.IMAGE, ModalityType.VIDEO])
class ResNet(UnimodalRepresentation):
- def __init__(self, model_name="ResNet18", layer="avgpool", output_file=None):
- self.data_type = torch.bfloat16
+ def __init__(
+ self,
+ model_name="ResNet18",
+ layer="avgpool",
+ output_file=None,
+ batch_size=32,
+ params=None,
+ ):
+ self.data_type = torch.float32
+ self.model = None
+ self.gpu_id = None
+ self.device = get_device()
self.model_name = model_name
+ self.batch_size = batch_size
parameters = self._get_parameters()
super().__init__("ResNet", ModalityType.EMBEDDING, parameters)
@@ -52,6 +64,17 @@ class ResNet(UnimodalRepresentation):
self.model.fc = Identity()
+ @property
+ def gpu_id(self):
+ return self._gpu_id
+
+ @gpu_id.setter
+ def gpu_id(self, gpu_id):
+ self._gpu_id = gpu_id
+ self.device = get_device(gpu_id)
+ if self.model is not None:
+ self.model = self.model.to(self.device)
+
@property
def model_name(self):
return self._model_name
@@ -60,40 +83,75 @@ class ResNet(UnimodalRepresentation):
def model_name(self, model_name):
self._model_name = model_name
if model_name == "ResNet18":
- self.model = (
- models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
- .to(get_device())
- .to(self.data_type)
- )
+ model = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
+ self.model = model.to(self.device)
+ self.model = self.model.to(self.data_type)
elif model_name == "ResNet34":
- self.model = models.resnet34(weights=models.ResNet34_Weights.DEFAULT).to(
- get_device()
- )
+ model = models.resnet34(weights=models.ResNet34_Weights.DEFAULT)
+ self.model = model.to(self.device)
self.model = self.model.to(self.data_type)
elif model_name == "ResNet50":
- self.model = (
- models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
- .to(get_device())
- .to(self.data_type)
- )
+ model = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
+ self.model = model.to(self.device)
+ self.model = self.model.to(self.data_type)
elif model_name == "ResNet101":
- self.model = (
- models.resnet101(weights=models.ResNet101_Weights.DEFAULT)
- .to(get_device())
- .to(self.data_type)
- )
+ model = models.resnet101(weights=models.ResNet101_Weights.DEFAULT)
+ self.model = model.to(self.device)
+ self.model = self.model.to(self.data_type)
elif model_name == "ResNet152":
- self.model = (
- models.resnet152(weights=models.ResNet152_Weights.DEFAULT)
- .to(get_device())
- .to(self.data_type)
- )
+ model = models.resnet152(weights=models.ResNet152_Weights.DEFAULT)
+ self.model = model.to(self.device)
+ self.model = self.model.to(self.data_type)
else:
raise NotImplementedError
+ def estimate_output_memory_bytes(self, input_stats: ImageStats) -> int:
+ return input_stats.num_instances * 512 * self.data_type.itemsize
+
+ def get_output_stats(self, input_stats) -> RepresentationStats:
+ return RepresentationStats(input_stats.num_instances, (512,))
+
+ def estimate_peak_memory_bytes(self, input_stats: ImageStats) -> dict:
+ input_bytes = (
+ self.batch_size
+ * input_stats.max_width
+ * input_stats.max_height
+ * input_stats.max_channels
+ * self.data_type.itemsize
+ )
+ output_bytes = self.estimate_output_memory_bytes(input_stats)
+ output_bytes_batch = output_bytes / input_stats.num_instances * self.batch_size
+
+ batch_peak_bytes = (
+ self.batch_size * 512 * self.data_type.itemsize
+ + self.batch_size * 224 * 224 * 3 * self.data_type.itemsize
+ ) * 2
+
+ safety_margin_bytes = 100 * 1024 * 1024
+
+ param_size = 0
+ for param in self.model.parameters():
+ param_size += param.nelement() * param.element_size()
+
+ buffer_size = 0
+ for buffer in self.model.buffers():
+ buffer_size += buffer.nelement() * buffer.element_size()
+
+ size_all_bytes = param_size + buffer_size
+
+ cpu_peak = (
+ size_all_bytes * 2 * self.data_type.itemsize
+ + output_bytes_batch
+ + output_bytes
+ + input_bytes
+ + safety_margin_bytes
+ )
+ gpu_peak = (size_all_bytes * self.data_type.itemsize + batch_peak_bytes) * 6
+ return {"cpu_peak_bytes": cpu_peak, "gpu_peak_bytes": gpu_peak}
+
def _get_parameters(self, high_level=True):
parameters = {"model_name": [], "layer_name": []}
for m in ["ResNet18", "ResNet34", "ResNet50", "ResNet101", "ResNet152"]:
@@ -113,14 +171,12 @@ class ResNet(UnimodalRepresentation):
parameters["layer_name"].append(name)
return parameters
- def transform(self, modality):
- self.data_type = torch.float32
+ def transform(self, modality, aggregation=None):
if next(self.model.parameters()).dtype != self.data_type:
self.model = self.model.to(self.data_type)
- dataset = CustomDataset(modality.data, self.data_type, get_device())
embeddings = {}
-
+ dataset = CustomDataset(modality.data, self.data_type, self.device)
res5c_output = None
def get_features(name_):
@@ -138,40 +194,56 @@ class ResNet(UnimodalRepresentation):
layer.register_forward_hook(get_features(name))
break
- for instance in torch.utils.data.DataLoader(dataset):
- video_id = instance["id"][0]
- frames = instance["data"][0]
- embeddings[video_id] = []
- batch_size = 64
-
- if modality.modality_type == ModalityType.IMAGE:
- frames = frames.unsqueeze(0)
-
- for start_index in range(0, len(frames), batch_size):
- end_index = min(start_index + batch_size, len(frames))
- frame_ids_range = range(start_index, end_index)
- frame_batch = frames[frame_ids_range]
-
- _ = self.model(frame_batch)
+ if modality.modality_type == ModalityType.IMAGE:
+ embeddings = []
+ for batch in torch.utils.data.DataLoader(
+ dataset, batch_size=self.batch_size
+ ):
+ image_batch = batch["data"]
+ _ = self.model(image_batch)
output = res5c_output
- if len(output.shape) > 2:
- output = torch.nn.functional.adaptive_avg_pool2d(output, (1, 1))
-
- embeddings[video_id].extend(
- torch.flatten(output, 1)
- .detach()
- .cpu()
- .float()
- .numpy()
- .astype(np.float32)
+ embeddings.extend(
+ output.squeeze().detach().cpu().numpy().astype(modality.data_type)
)
-
- embeddings[video_id] = np.array(embeddings[video_id])
+ torch.cuda.empty_cache()
+ else:
+ for instance in torch.utils.data.DataLoader(dataset):
+ video_id = instance["id"][0]
+ frames = instance["data"][0]
+ embeddings[video_id] = []
+ batch_size = 64
+
+ if modality.modality_type == ModalityType.IMAGE:
+ frames = frames.unsqueeze(0)
+
+ for start_index in range(0, len(frames), batch_size):
+ end_index = min(start_index + batch_size, len(frames))
+ frame_ids_range = range(start_index, end_index)
+ frame_batch = frames[frame_ids_range]
+
+ _ = self.model(frame_batch)
+ output = res5c_output
+ if len(output.shape) > 2:
+ output = torch.nn.functional.adaptive_avg_pool2d(output, (1, 1))
+ # TODO: check if the dimensions are correct here
+ embeddings[video_id].extend(
+ torch.flatten(output, 1)
+ .detach()
+ .cpu()
+ .float()
+ .numpy()
+ .astype(np.float32)
+ )
+
+ embeddings[video_id] = np.array(embeddings[video_id])
transformed_modality = TransformedModality(
modality, self, self.output_modality_type
)
- transformed_modality.data = list(embeddings.values())
+ if isinstance(embeddings, dict):
+ transformed_modality.data = list(embeddings.values())
+ else:
+ transformed_modality.data = embeddings
return transformed_modality
diff --git a/src/main/python/systemds/scuro/representations/spectrogram.py b/src/main/python/systemds/scuro/representations/spectrogram.py
index 662fb0d627..ba455cb980 100644
--- a/src/main/python/systemds/scuro/representations/spectrogram.py
+++ b/src/main/python/systemds/scuro/representations/spectrogram.py
@@ -24,31 +24,105 @@ import numpy as np
from systemds.scuro.modality.type import ModalityType
from systemds.scuro.modality.transformed import TransformedModality
+from systemds.scuro.representations.representation import RepresentationStats
from systemds.scuro.representations.unimodal import UnimodalRepresentation
-from systemds.scuro.drsearch.operator_registry import register_representation
+from systemds.scuro.drsearch.operator_registry import (
+ register_representation,
+ register_context_representation_operator,
+)
+from systemds.scuro.utils.static_variables import (
+ NP_ARRAY_HEADER_BYTES,
+ PY_LIST_HEADER_BYTES,
+ PY_LIST_SLOT_BYTES,
+)
@register_representation(ModalityType.AUDIO)
+@register_context_representation_operator(ModalityType.AUDIO)
class Spectrogram(UnimodalRepresentation):
- def __init__(self, hop_length=512, n_fft=2048):
+ def __init__(self, hop_length=512, n_fft=2048, params=None):
parameters = {"hop_length": [256, 512, 1024, 2048], "n_fft": [1024, 2048, 4096]}
super().__init__("Spectrogram", ModalityType.TIMESERIES, parameters, False)
self.hop_length = int(hop_length)
self.n_fft = int(n_fft)
- def transform(self, modality):
+ def transform(self, modality, aggregation=None):
transformed_modality = TransformedModality(
modality, self, self.output_modality_type
)
result = []
for i, sample in enumerate(modality.data):
- spectrogram = librosa.stft(
- y=np.array(np.abs(sample)), hop_length=self.hop_length, n_fft=self.n_fft
- )
- S_dB = librosa.amplitude_to_db(np.abs(spectrogram))
-
- result.append(S_dB.T)
+ computed_feature = self.compute_feature(sample)
+ result.append(computed_feature)
transformed_modality.data = result
return transformed_modality
+
+ def compute_feature(self, instance):
+ spectrogram = librosa.stft(
+ y=np.array(np.abs(instance)), hop_length=self.hop_length, n_fft=self.n_fft
+ )
+ return librosa.amplitude_to_db(np.abs(spectrogram)).T
+
+ def estimate_peak_memory_bytes(self, input_stats) -> dict:
+ # TODO: validate this function
+ n = int(input_stats.num_instances)
+ output_shape = input_stats.output_shape
+
+ signal_length = output_shape[0]
+ if signal_length < self.n_fft:
+ num_frames = 1
+ else:
+ num_frames = 1 + (signal_length - self.n_fft) // self.hop_length
+ num_frames = max(int(num_frames), 1)
+ num_freq_bins = 1 + self.n_fft // 2
+
+ output_payload_bytes = (
+ num_frames * num_freq_bins * np.dtype(np.float32).itemsize
+ )
+ per_instance_retained = (
+ output_payload_bytes + NP_ARRAY_HEADER_BYTES + PY_LIST_SLOT_BYTES
+ )
+ retained_output_bytes = PY_LIST_HEADER_BYTES + n * per_instance_retained
+ input_copy_bytes = max(signal_length, 1) * np.dtype(np.float32).itemsize
+ stft_bytes = num_frames * num_freq_bins * np.dtype(np.complex64).itemsize
+
+ magnitude_bytes = num_frames * num_freq_bins * np.dtype(np.float32).itemsize
+ db_bytes = output_payload_bytes
+ fft_workspace_bytes = max(2 * 1024 * 1024, stft_bytes // 2)
+ transient_one_instance_bytes = (
+ input_copy_bytes
+ + stft_bytes
+ + magnitude_bytes
+ + db_bytes
+ + fft_workspace_bytes
+ )
+ cpu_peak = retained_output_bytes + transient_one_instance_bytes
+
+ return {
+ "cpu_peak_bytes": int(cpu_peak),
+ "gpu_peak_bytes": 0,
+ }
+
+ def get_output_stats(self, input_stats) -> RepresentationStats:
+ num_instances = getattr(input_stats, "num_instances", 0)
+
+ if hasattr(input_stats, "max_length"):
+ signal_length = input_stats.max_length
+ elif hasattr(input_stats, "output_shape") and input_stats.output_shape:
+ signal_length = input_stats.output_shape[0]
+ else:
+ signal_length = 0
+
+ if signal_length <= 0:
+ num_frames = 1
+ else:
+ if signal_length < self.n_fft:
+ num_frames = 1
+ else:
+ num_frames = 1 + (signal_length - self.n_fft) // self.hop_length
+ num_frames = max(int(num_frames), 1)
+
+ num_freq_bins = 1 + self.n_fft // 2
+ return RepresentationStats(num_instances, (num_frames, num_freq_bins))
diff --git a/src/main/python/systemds/scuro/representations/sum.py b/src/main/python/systemds/scuro/representations/sum.py
index 79655a4b66..e6187b4050 100644
--- a/src/main/python/systemds/scuro/representations/sum.py
+++ b/src/main/python/systemds/scuro/representations/sum.py
@@ -26,13 +26,14 @@ from systemds.scuro.modality.modality import Modality
from systemds.scuro.representations.utils import pad_sequences
from systemds.scuro.representations.fusion import Fusion
+from systemds.scuro.representations.representation import RepresentationStats
from systemds.scuro.drsearch.operator_registry import register_fusion_operator
@register_fusion_operator()
class Sum(Fusion):
- def __init__(self):
+ def __init__(self, params=None):
"""
Combines modalities using colum-wise sum
"""
@@ -55,3 +56,27 @@ class Sum(Fusion):
]["type"],
)
return data
+
+ def get_output_stats(self, input_stats_list) -> RepresentationStats:
+ if isinstance(input_stats_list, RepresentationStats):
+ return input_stats_list
+
+ stats_list = list(input_stats_list)
+ if not stats_list:
+ return RepresentationStats(0, (0,))
+
+ def num_elements(stats: RepresentationStats) -> int:
+ n = 1
+ for d in stats.output_shape:
+ n *= d
+ return n
+
+ largest = max(stats_list, key=num_elements)
+ return RepresentationStats(largest.num_instances, largest.output_shape)
+
+ def estimate_peak_memory_bytes(self, input_stats) -> dict:
+ # TODO
+ return {
+ "cpu_peak_bytes": 0,
+ "gpu_peak_bytes": 0,
+ }
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 e8511dd0cf..c46d12bcb5 100644
--- a/src/main/python/systemds/scuro/representations/swin_video_transformer.py
+++ b/src/main/python/systemds/scuro/representations/swin_video_transformer.py
@@ -22,6 +22,7 @@ from torchvision.models.video.swin_transformer import swin3d_t
from systemds.scuro.modality.transformed import TransformedModality
from systemds.scuro.representations.unimodal import UnimodalRepresentation
+from systemds.scuro.representations.representation import RepresentationStats
from typing import Callable, Dict, Tuple, Any
import torch.utils.data
import torch
@@ -31,12 +32,16 @@ from systemds.scuro.modality.type import ModalityType
from systemds.scuro.drsearch.operator_registry import register_representation
from systemds.scuro.utils.torch_dataset import CustomDataset
-from systemds.scuro.utils.static_variables import get_device
+from systemds.scuro.utils.static_variables import (
+ compute_batch_size,
+ get_device,
+ get_device_for_model,
+)
@register_representation([ModalityType.VIDEO])
class SwinVideoTransformer(UnimodalRepresentation):
- def __init__(self, layer_name="avgpool"):
+ def __init__(self, layer_name="avgpool", params=None):
parameters = {
"layer_name": [
"features",
@@ -52,14 +57,17 @@ class SwinVideoTransformer(UnimodalRepresentation):
self.data_type = torch.float
super().__init__("SwinVideoTransformer", ModalityType.EMBEDDING, parameters)
self.layer_name = layer_name
- self.model = swin3d_t(weights=models.video.Swin3D_T_Weights.KINETICS400_V1).to(
- get_device()
- )
+ self.model = swin3d_t(weights=models.video.Swin3D_T_Weights.KINETICS400_V1)
+ self.device = get_device_for_model(self.model, memory_factor=1.5)
+ self.model = self.model.to(self.device)
self.model.eval()
for param in self.model.parameters():
param.requires_grad = False
- def transform(self, modality):
+ def get_output_stats(self, input_stats) -> RepresentationStats:
+ return RepresentationStats(input_stats.num_instances, (768,))
+
+ def transform(self, modality, aggregation=None):
embeddings = {}
swin_output = None
@@ -72,12 +80,22 @@ class SwinVideoTransformer(UnimodalRepresentation):
return hook
+ sample = modality.data[0] if modality.data else ""
+ self.batch_size = compute_batch_size(
+ model=self.model,
+ device=self.device,
+ sample_data=sample,
+ tokenizer=None,
+ max_seq_length=None,
+ max_batch_size=128,
+ )
+
if self.layer_name:
for name, layer in self.model.named_modules():
if name == self.layer_name:
layer.register_forward_hook(get_features(name))
break
- dataset = CustomDataset(modality.data, self.data_type, get_device())
+ dataset = CustomDataset(modality.data, self.data_type, self.device)
for instance in torch.utils.data.DataLoader(dataset):
video_id = instance["id"][0]
diff --git a/src/main/python/systemds/scuro/representations/text_context_with_indices.py b/src/main/python/systemds/scuro/representations/text_context_with_indices.py
index 5a3c3b34e0..095e9f8106 100644
--- a/src/main/python/systemds/scuro/representations/text_context_with_indices.py
+++ b/src/main/python/systemds/scuro/representations/text_context_with_indices.py
@@ -18,12 +18,15 @@
# under the License.
#
# -------------------------------------------------------------
+import math
import re
from typing import List, Any
-
+import numpy as np
+from systemds.scuro.dataloader.text_loader import TextStats
from systemds.scuro.drsearch.operator_registry import register_context_operator
from systemds.scuro.representations.context import Context
from systemds.scuro.modality.type import ModalityType
+from systemds.scuro.representations.representation import RepresentationStats
# TODO: Use this to get indices for text chunks based on different splitting strategies
# To use this approach a differnt extration of text chunks is needed in either the TextModality or the Representations
@@ -85,7 +88,7 @@ class WordCountSplitIndices(Context):
overlap (int): Number of overlapping words between chunks (default: 0)
"""
- def __init__(self, max_words=55, overlap=0):
+ def __init__(self, max_words=55, overlap=0, params=None):
parameters = {
"max_words": [40, 50, 55, 60, 70, 250, 300, 350, 400, 450],
"overlap": [0, 10, 20, 30],
@@ -144,16 +147,58 @@ class SentenceBoundarySplitIndices(Context):
min_words (int): Minimum number of words per chunk before splitting (default: 10)
"""
- def __init__(self, max_words=55, min_words=10, overlap=0.1):
+ def __init__(self, max_words=55, min_words=10, overlap=0.1, params=None):
parameters = {
"max_words": [40, 50, 55, 60, 70, 250, 300, 350, 400, 450],
"min_words": [10, 20, 30],
}
super().__init__("SentenceBoundarySplit", parameters)
- self.max_words = int(max_words)
- self.min_words = max(1, int(min_words))
- self.overlap = overlap
- self.stride = max(1, int(max_words * (1 - overlap)))
+ if params is not None:
+ self.max_words = int(params.get("max_words", max_words))
+ self.min_words = int(params.get("min_words", min_words))
+ self.overlap = float(params.get("overlap", overlap))
+ self.stride = int(params.get("stride", max_words * (1 - overlap)))
+ else:
+ self.max_words = int(max_words)
+ self.min_words = max(1, int(min_words))
+ self.overlap = overlap
+ self.stride = max(1, int(max_words * (1 - overlap)))
+ self.data_type = np.int32
+
+ def get_output_stats(self, input_stats: TextStats) -> RepresentationStats:
+ return RepresentationStats(
+ input_stats.num_instances,
+ (
+ max(
+ math.ceil(
+ input_stats.avg_words / (self.max_words - self.overlap) - 1
+ ),
+ 1,
+ ),
+ self.max_words,
+ ),
+ )
+
+ def estimate_output_memory_bytes(self, input_stats: TextStats) -> int:
+ output_memory_bytes = 1
+ output_shape = self.get_output_stats(input_stats).output_shape
+ for dim in output_shape:
+ output_memory_bytes *= dim
+
+ return (
+ input_stats.num_instances
+ * output_memory_bytes
+ * np.dtype(self.data_type).itemsize
+ )
+
+ def estimate_peak_memory_bytes(self, input_stats: TextStats) -> dict:
+ per_instance_temp = int(input_stats.max_words * 64) * input_stats.num_instances
+
+ return {
+ "cpu_peak_bytes": per_instance_temp
+ + self.estimate_output_memory_bytes(input_stats) * 2,
+ "gpu_peak_bytes": 0,
+ }
def execute(self, modality):
"""
@@ -240,7 +285,7 @@ class OverlappingSplitIndices(Context):
stride (int, optional): Step size in words. If None, stride = max_words - overlap_words
"""
- def __init__(self, max_words=55, overlap=0.5, stride=None):
+ def __init__(self, max_words=55, overlap=0.5, stride=None, params=None):
overlap_words = int(max_words * overlap)
if stride is None:
stride = max_words - overlap_words
@@ -254,6 +299,42 @@ class OverlappingSplitIndices(Context):
self.max_words = max_words
self.overlap = overlap
self.stride = stride
+ self.data_type = np.int32
+
+ def get_output_stats(self, input_stats: TextStats) -> RepresentationStats:
+ return RepresentationStats(
+ input_stats.num_instances,
+ (
+ max(
+ math.ceil(
+ input_stats.avg_words / (self.max_words - self.overlap) - 1
+ ),
+ 1,
+ ),
+ self.max_words,
+ ),
+ )
+
+ def estimate_output_memory_bytes(self, input_stats: TextStats) -> int:
+ output_memory_bytes = 1
+ output_shape = self.get_output_stats(input_stats).output_shape
+ for dim in output_shape:
+ output_memory_bytes *= dim
+
+ return (
+ input_stats.num_instances
+ * output_memory_bytes
+ * np.dtype(self.data_type).itemsize
+ )
+
+ def estimate_peak_memory_bytes(self, input_stats: TextStats) -> dict:
+ per_instance_temp = int(input_stats.max_words * 64) * input_stats.num_instances
+
+ return {
+ "cpu_peak_bytes": per_instance_temp
+ + self.estimate_output_memory_bytes(input_stats) * 2,
+ "gpu_peak_bytes": 0,
+ }
def execute(self, modality):
"""
diff --git a/src/main/python/systemds/scuro/representations/tfidf.py b/src/main/python/systemds/scuro/representations/tfidf.py
index 95fa6c111f..3fd61099f1 100644
--- a/src/main/python/systemds/scuro/representations/tfidf.py
+++ b/src/main/python/systemds/scuro/representations/tfidf.py
@@ -21,7 +21,9 @@
import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
+from systemds.scuro.dataloader.text_loader import TextStats
from systemds.scuro.modality.transformed import TransformedModality
+from systemds.scuro.representations.representation import RepresentationStats
from systemds.scuro.representations.unimodal import UnimodalRepresentation
from systemds.scuro.representations.utils import save_embeddings
@@ -31,22 +33,45 @@ from systemds.scuro.drsearch.operator_registry import register_representation
@register_representation(ModalityType.TEXT)
class TfIdf(UnimodalRepresentation):
- def __init__(self, min_df=2, output_file=None):
+ def __init__(self, min_df=2, output_file=None, params=None):
parameters = {"min_df": [min_df, 4, 8]}
super().__init__("TF-IDF", ModalityType.EMBEDDING, parameters)
self.min_df = int(min_df)
self.output_file = output_file
+ self.data_type = np.float32
- def transform(self, modality):
+ def get_output_stats(self, input_stats: TextStats) -> RepresentationStats:
+ vocab_estimate = min(
+ 100_000, max(1000, input_stats.num_instances * input_stats.max_length)
+ )
+ return RepresentationStats(
+ input_stats.num_instances, (vocab_estimate,), output_shape_is_known=False
+ )
+
+ def estimate_output_memory_bytes(self, input_stats: TextStats) -> int:
+ return (
+ input_stats.num_instances
+ * self.get_output_stats(input_stats).output_shape[0]
+ * np.dtype(self.data_type).itemsize
+ )
+
+ def estimate_peak_memory_bytes(self, input_stats: TextStats) -> dict:
+ output_bytes = self.estimate_output_memory_bytes(input_stats)
+ vectorizer_overhead = 640 * 1024
+ return {
+ "cpu_peak_bytes": output_bytes + vectorizer_overhead,
+ "gpu_peak_bytes": 0,
+ }
+
+ def transform(self, modality, aggregation=None):
transformed_modality = TransformedModality(modality, self)
vectorizer = TfidfVectorizer(min_df=self.min_df)
X = vectorizer.fit_transform(modality.data)
- X = [np.array(x).astype(np.float32).reshape(1, -1) for x in X.toarray()]
if self.output_file is not None:
save_embeddings(X, self.output_file)
- transformed_modality.data_type = np.float32
- transformed_modality.data = X
+ transformed_modality.data_type = self.data_type
+ transformed_modality.data = X.toarray()
return transformed_modality
diff --git a/src/main/python/systemds/scuro/representations/timeseries_representations.py b/src/main/python/systemds/scuro/representations/timeseries_representations.py
index 3270992a97..bb4ea4f49c 100644
--- a/src/main/python/systemds/scuro/representations/timeseries_representations.py
+++ b/src/main/python/systemds/scuro/representations/timeseries_representations.py
@@ -23,20 +23,25 @@ from scipy import stats
from systemds.scuro.modality.type import ModalityType
from systemds.scuro.modality.transformed import TransformedModality
+from systemds.scuro.representations.representation import RepresentationStats
from systemds.scuro.representations.unimodal import UnimodalRepresentation
-from systemds.scuro.drsearch.operator_registry import register_representation
+from systemds.scuro.drsearch.operator_registry import (
+ register_representation,
+ register_context_representation_operator,
+)
class TimeSeriesRepresentation(UnimodalRepresentation):
- def __init__(self, name, parameters=None):
- if parameters is None:
- parameters = {}
+ def __init__(self, name, parameters=None, params=None):
+ if params is None:
+ params = {}
+
super().__init__(name, ModalityType.EMBEDDING, parameters, False)
def compute_feature(self, signal):
raise NotImplementedError("Subclasses should implement this method.")
- def transform(self, modality):
+ def transform(self, modality, aggregation=None):
transformed_modality = TransformedModality(
modality, self, self.output_modality_type
)
@@ -51,10 +56,25 @@ class TimeSeriesRepresentation(UnimodalRepresentation):
)
return transformed_modality
+ def get_output_stats(self, input_stats):
+ return RepresentationStats(input_stats.num_instances, (1,))
+
+ def estimate_output_memory_bytes(self, input_stats):
+ # TODO: adapt this to the actual output shapes and transformations
+ return input_stats.num_instances * 4
+
+ def estimate_peak_memory_bytes(self, input_stats):
+ # TODO: adapt this to the actual output shapes and transformations
+ return {
+ "cpu_peak_bytes": self.estimate_output_memory_bytes(input_stats),
+ "gpu_peak_bytes": 0,
+ }
+
@register_representation([ModalityType.TIMESERIES])
+@register_context_representation_operator(ModalityType.TIMESERIES)
class Mean(TimeSeriesRepresentation):
- def __init__(self):
+ def __init__(self, params=None):
super().__init__("Mean")
def compute_feature(self, signal):
@@ -62,8 +82,9 @@ class Mean(TimeSeriesRepresentation):
@register_representation([ModalityType.TIMESERIES])
+@register_context_representation_operator(ModalityType.TIMESERIES)
class Min(TimeSeriesRepresentation):
- def __init__(self):
+ def __init__(self, params=None):
super().__init__("Min")
def compute_feature(self, signal):
@@ -71,8 +92,9 @@ class Min(TimeSeriesRepresentation):
@register_representation([ModalityType.TIMESERIES])
+@register_context_representation_operator(ModalityType.TIMESERIES)
class Max(TimeSeriesRepresentation):
- def __init__(self):
+ def __init__(self, params=None):
super().__init__("Max")
def compute_feature(self, signal):
@@ -80,8 +102,9 @@ class Max(TimeSeriesRepresentation):
@register_representation([ModalityType.TIMESERIES])
+@register_context_representation_operator(ModalityType.TIMESERIES)
class Sum(TimeSeriesRepresentation):
- def __init__(self):
+ def __init__(self, params=None):
super().__init__("Sum")
def compute_feature(self, signal):
@@ -89,8 +112,9 @@ class Sum(TimeSeriesRepresentation):
@register_representation([ModalityType.TIMESERIES])
+@register_context_representation_operator(ModalityType.TIMESERIES)
class Std(TimeSeriesRepresentation):
- def __init__(self):
+ def __init__(self, params=None):
super().__init__("Std")
def compute_feature(self, signal):
@@ -98,8 +122,9 @@ class Std(TimeSeriesRepresentation):
@register_representation([ModalityType.TIMESERIES])
+@register_context_representation_operator(ModalityType.TIMESERIES)
class Skew(TimeSeriesRepresentation):
- def __init__(self):
+ def __init__(self, params=None):
super().__init__("Skew")
def compute_feature(self, signal):
@@ -107,8 +132,9 @@ class Skew(TimeSeriesRepresentation):
@register_representation([ModalityType.TIMESERIES])
+@register_context_representation_operator(ModalityType.TIMESERIES)
class Quantile(TimeSeriesRepresentation):
- def __init__(self, quantile=0.9):
+ def __init__(self, quantile=0.9, params=None):
super().__init__(
"Qunatile", {"quantile": [0.1, 0.25, 0.5, 0.75, 0.9, 0.95, 0.99]}
)
@@ -119,8 +145,9 @@ class Quantile(TimeSeriesRepresentation):
@register_representation([ModalityType.TIMESERIES])
+@register_context_representation_operator(ModalityType.TIMESERIES)
class Kurtosis(TimeSeriesRepresentation):
- def __init__(self):
+ def __init__(self, params=None):
super().__init__("Kurtosis")
def compute_feature(self, signal):
@@ -128,8 +155,9 @@ class Kurtosis(TimeSeriesRepresentation):
@register_representation([ModalityType.TIMESERIES])
+@register_context_representation_operator(ModalityType.TIMESERIES)
class RMS(TimeSeriesRepresentation):
- def __init__(self):
+ def __init__(self, params=None):
super().__init__("RMS")
def compute_feature(self, signal):
@@ -137,8 +165,9 @@ class RMS(TimeSeriesRepresentation):
@register_representation([ModalityType.TIMESERIES])
+@register_context_representation_operator(ModalityType.TIMESERIES)
class ZeroCrossingRate(TimeSeriesRepresentation):
- def __init__(self):
+ def __init__(self, params=None):
super().__init__("ZeroCrossingRate")
def compute_feature(self, signal):
@@ -146,8 +175,9 @@ class ZeroCrossingRate(TimeSeriesRepresentation):
@register_representation([ModalityType.TIMESERIES])
+@register_context_representation_operator(ModalityType.TIMESERIES)
class ACF(TimeSeriesRepresentation):
- def __init__(self, k=1):
+ def __init__(self, k=1, params=None):
super().__init__("ACF", {"k": [1, 2, 5, 10, 20, 25, 50, 100, 200, 500]})
self.k = k
@@ -173,6 +203,7 @@ class ACF(TimeSeriesRepresentation):
@register_representation([ModalityType.TIMESERIES])
+@register_context_representation_operator(ModalityType.TIMESERIES)
class FrequencyMagnitude(TimeSeriesRepresentation):
def __init__(self):
super().__init__("FrequencyMagnitude")
@@ -182,8 +213,9 @@ class FrequencyMagnitude(TimeSeriesRepresentation):
@register_representation([ModalityType.TIMESERIES])
+@register_context_representation_operator(ModalityType.TIMESERIES)
class SpectralCentroid(TimeSeriesRepresentation):
- def __init__(self, fs=1.0):
+ def __init__(self, fs=1.0, params=None):
super().__init__("SpectralCentroid", parameters={"fs": [0.5, 1.0, 2.0]})
self.fs = fs
@@ -196,8 +228,9 @@ class SpectralCentroid(TimeSeriesRepresentation):
@register_representation([ModalityType.TIMESERIES])
+@register_context_representation_operator(ModalityType.TIMESERIES)
class BandpowerFFT(TimeSeriesRepresentation):
- def __init__(self, fs=1.0, f1=0.0, f2=0.5):
+ def __init__(self, fs=1.0, f1=0.0, f2=0.5, params=None):
super().__init__(
"BandpowerFFT",
parameters={"fs": [0.5, 1.0], "f1": [0.0, 1.0], "f2": [0.5, 1.0]},
diff --git a/src/main/python/systemds/scuro/representations/unimodal.py b/src/main/python/systemds/scuro/representations/unimodal.py
index a1a1632c26..3ad787c677 100644
--- a/src/main/python/systemds/scuro/representations/unimodal.py
+++ b/src/main/python/systemds/scuro/representations/unimodal.py
@@ -42,7 +42,7 @@ class UnimodalRepresentation(Representation):
self.initial_context_length = None
@abc.abstractmethod
- def transform(self, data):
+ def transform(self, data, aggregation=None):
# TODO: check if there is a way to time the transformation in here (needed for chunked execution)
raise f"Not implemented for {self.name}"
diff --git a/src/main/python/systemds/scuro/representations/vgg.py b/src/main/python/systemds/scuro/representations/vgg.py
index 8bc4a15b95..fa8f121438 100644
--- a/src/main/python/systemds/scuro/representations/vgg.py
+++ b/src/main/python/systemds/scuro/representations/vgg.py
@@ -30,19 +30,31 @@ import re
import torchvision.models as models
import numpy as np
from systemds.scuro.modality.type import ModalityType
-from systemds.scuro.utils.static_variables import get_device
+from systemds.scuro.utils.static_variables import (
+ get_device,
+)
+from systemds.scuro.dataloader.image_loader import ImageStats
+from systemds.scuro.representations.representation import RepresentationStats
@register_representation([ModalityType.IMAGE, ModalityType.VIDEO])
class VGG19(UnimodalRepresentation):
- def __init__(self, layer="classifier.0", output_file=None):
+ def __init__(
+ self, layer="classifier.0", output_file=None, params=None, batch_size=32
+ ):
self.data_type = torch.bfloat16
- self.model = models.vgg19(weights=models.VGG19_Weights.DEFAULT).to(get_device())
+ self.model = None
+ self.gpu_id = None
+ self.device = get_device()
+ 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)
self.output_file = output_file
self.layer_name = layer
self.model.eval()
+ self.batch_size = batch_size
+
for param in self.model.parameters():
param.requires_grad = False
@@ -52,6 +64,17 @@ class VGG19(UnimodalRepresentation):
self.model.fc = Identity()
+ @property
+ def gpu_id(self):
+ return self._gpu_id
+
+ @gpu_id.setter
+ def gpu_id(self, gpu_id):
+ self._gpu_id = gpu_id
+ self.device = get_device(gpu_id)
+ if self.model is not None:
+ self.model = self.model.to(self.device)
+
def _get_parameters(self):
parameters = {
"layer_name": [
@@ -64,20 +87,58 @@ class VGG19(UnimodalRepresentation):
return parameters
- def transform(self, modality):
+ def estimate_output_memory_bytes(self, input_stats: ImageStats) -> int:
+ return input_stats.num_instances * 4096 * np.dtype(np.float32).itemsize
+
+ def get_output_stats(self, input_stats) -> RepresentationStats:
+ return RepresentationStats(input_stats.num_instances, (4096,))
+
+ def estimate_peak_memory_bytes(self, input_stats: ImageStats) -> dict:
+ batch_size_bytes = 224 * 224 * 3 * self.data_type.itemsize * self.batch_size * 2
+ input_bytes = (
+ self.batch_size
+ * input_stats.max_width
+ * input_stats.max_height
+ * input_stats.max_channels
+ * self.data_type.itemsize
+ )
+ model = models.vgg19(weights=models.VGG19_Weights.DEFAULT)
+ param_bytes = sum(p.numel() for p in model.parameters())
+ buffer_bytes = sum(b.numel() for b in model.buffers())
+ model_size_bytes = param_bytes * 4 + buffer_bytes * 4
+
+ return {
+ "cpu_peak_bytes": (
+ self.estimate_output_memory_bytes(input_stats)
+ + self.estimate_output_memory_bytes(input_stats)
+ / input_stats.num_instances
+ * self.batch_size
+ + model_size_bytes
+ + input_bytes
+ )
+ * 2,
+ "gpu_peak_bytes": (
+ model_size_bytes
+ + batch_size_bytes
+ + self.estimate_output_memory_bytes(input_stats)
+ / input_stats.num_instances
+ * self.batch_size
+ )
+ * 5,
+ }
+
+ def transform(self, modality, aggregation=None):
self.data_type = torch.float32
if next(self.model.parameters()).dtype != self.data_type:
self.model = self.model.to(self.data_type)
- dataset = CustomDataset(modality.data, self.data_type, get_device())
- embeddings = {}
- activations = {}
+ self.activations = {}
def get_activation(name_):
def hook(
_module: torch.nn.Module, input_: Tuple[torch.Tensor], output: Any
):
- activations[name_] = output
+ self.activations[name_] = output
return hook
@@ -90,26 +151,51 @@ class VGG19(UnimodalRepresentation):
self.model.classifier[int(digit)].register_forward_hook(
get_activation(self.layer_name)
)
+ is_image = len(modality.data[0].shape) == 3
+ embeddings = (
+ self._transform_image_modality(modality)
+ if is_image
+ else self._transform_video_modality(modality)
+ )
+
+ transformed_modality = TransformedModality(
+ modality, self, self.output_modality_type
+ )
+ transformed_modality.data = embeddings
+
+ return transformed_modality
+
+ def _transform_image_modality(self, modality):
+ dataset = CustomDataset(modality.data, self.data_type, self.device)
+ embeddings = []
+ for instance in torch.utils.data.DataLoader(dataset, self.batch_size):
+ frames = instance["data"]
+
+ _ = self.model(frames)
+ output = self.activations[self.layer_name]
+
+ if len(output.shape) == 4:
+ output = torch.nn.functional.adaptive_avg_pool2d(output, (1, 1))
+
+ embeddings.extend(output.detach().cpu().float().numpy().astype(np.float32))
+
+ return np.array(embeddings)
+
+ def _transform_video_modality(self, modality):
+ dataset = CustomDataset(modality.data, self.data_type, self.device)
+ embeddings = {}
for instance in torch.utils.data.DataLoader(dataset):
video_id = instance["id"][0]
frames = instance["data"][0]
embeddings[video_id] = []
- if frames.dim() == 3:
- # Single image: (3, 224, 224) -> (1, 3, 224, 224)
- frames = frames.unsqueeze(0)
- batch_size = 1
- else:
- # Video: (T, 3, 224, 224) - process in batches
- batch_size = 32
-
- for start_index in range(0, frames.shape[0], batch_size):
- end_index = min(start_index + batch_size, frames.shape[0])
+ for start_index in range(0, frames.shape[0], self.batch_size):
+ end_index = min(start_index + self.batch_size, frames.shape[0])
frame_batch = frames[start_index:end_index]
_ = self.model(frame_batch)
- output = activations[self.layer_name]
+ output = self.activations[self.layer_name]
if len(output.shape) == 4:
output = torch.nn.functional.adaptive_avg_pool2d(output, (1, 1))
@@ -125,9 +211,4 @@ class VGG19(UnimodalRepresentation):
embeddings[video_id] = np.array(embeddings[video_id])
- transformed_modality = TransformedModality(
- modality, self, self.output_modality_type
- )
- transformed_modality.data = list(embeddings.values())
-
- return transformed_modality
+ return list(embeddings.values())
diff --git a/src/main/python/systemds/scuro/representations/video_context.py b/src/main/python/systemds/scuro/representations/video_context.py
new file mode 100644
index 0000000000..af3d3f4be8
--- /dev/null
+++ b/src/main/python/systemds/scuro/representations/video_context.py
@@ -0,0 +1,72 @@
+# -------------------------------------------------------------
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+# -------------------------------------------------------------
+import numpy as np
+
+from systemds.scuro.modality.type import ModalityType
+from systemds.scuro.modality.transformed import TransformedModality
+from systemds.scuro.representations.context import Context
+
+
+# TODO: this should not be a context but a sampling operator
+class UniformFrameSampling(Context):
+ """
+ Downsamples video by sampling every Nth frame uniformly.
+ """
+
+ def __init__(self, frame_interval=5):
+ parameters = {"frame_interval": [1, 2, 3, 5, 10, 15, 30]}
+ super().__init__("UniformFrameSampling", parameters)
+ self.frame_interval = max(1, int(frame_interval))
+
+ def execute(self, modality):
+ """
+ Sample frames uniformly from each video instance.
+
+ Returns:
+ Modality with downsampled video data.
+ """
+ transformed_modality = TransformedModality(
+ modality, self, modality.modality_type
+ )
+ sampled_data = []
+
+ for video_frames in modality.data:
+ if video_frames is None or video_frames.size == 0:
+ sampled_data.append(np.array([]))
+ continue
+
+ num_frames = video_frames.shape[0]
+
+ if num_frames == 0:
+ sampled_data.append(np.array([]))
+ continue
+
+ sampled_indices = list(range(0, num_frames, self.frame_interval))
+
+ if not sampled_indices:
+ sampled_indices = [0]
+
+ sampled_frames = video_frames[sampled_indices]
+ sampled_data.append(sampled_frames)
+
+ transformed_modality.data = sampled_data
+
+ return transformed_modality
diff --git a/src/main/python/systemds/scuro/representations/wav2vec.py b/src/main/python/systemds/scuro/representations/wav2vec.py
index 86145e3769..c599550be1 100644
--- a/src/main/python/systemds/scuro/representations/wav2vec.py
+++ b/src/main/python/systemds/scuro/representations/wav2vec.py
@@ -25,6 +25,7 @@ import torch
from systemds.scuro.modality.type import ModalityType
from systemds.scuro.modality.transformed import TransformedModality
+from systemds.scuro.representations.representation import RepresentationStats
from systemds.scuro.representations.unimodal import UnimodalRepresentation
from systemds.scuro.drsearch.operator_registry import register_representation
@@ -35,7 +36,7 @@ warnings.filterwarnings("ignore", message="Some weights of")
@register_representation(ModalityType.AUDIO)
class Wav2Vec(UnimodalRepresentation):
- def __init__(self):
+ def __init__(self, params=None):
super().__init__("Wav2Vec", ModalityType.TIMESERIES, {})
self.processor = Wav2Vec2Processor.from_pretrained(
"facebook/wav2vec2-base-960h"
@@ -44,7 +45,7 @@ class Wav2Vec(UnimodalRepresentation):
"facebook/wav2vec2-base-960h"
).float()
- def transform(self, modality):
+ def transform(self, modality, aggregation=None):
transformed_modality = TransformedModality(
modality, self, self.output_modality_type
)
@@ -68,3 +69,8 @@ class Wav2Vec(UnimodalRepresentation):
transformed_modality.data = result
return transformed_modality
+
+ def get_output_stats(self, input_stats) -> RepresentationStats:
+ num_instances = getattr(input_stats, "num_instances", 0)
+ embedding_dim = 768
+ return RepresentationStats(num_instances, (embedding_dim,))
diff --git a/src/main/python/systemds/scuro/representations/window_aggregation.py b/src/main/python/systemds/scuro/representations/window_aggregation.py
index 039387eb01..c9622f1a2a 100644
--- a/src/main/python/systemds/scuro/representations/window_aggregation.py
+++ b/src/main/python/systemds/scuro/representations/window_aggregation.py
@@ -18,7 +18,6 @@
# under the License.
#
# -------------------------------------------------------------
-import copy
import numpy as np
import math
@@ -28,24 +27,50 @@ from systemds.scuro.modality.type import DataLayout, ModalityType
from systemds.scuro.drsearch.operator_registry import register_context_operator
from systemds.scuro.representations.aggregate import Aggregation
from systemds.scuro.representations.context import Context
+from systemds.scuro.representations.representation import (
+ Representation,
+ RepresentationStats,
+)
class Window(Context):
def __init__(self, name, aggregation_function):
- is_ts_rep = False
- if isinstance(aggregation_function, str):
- parameters = {
- "aggregation_function": list(Aggregation().get_aggregation_functions()),
- }
+ self.aggregation_function = aggregation_function
+ parameters = {}
+ if isinstance(self.aggregation_function, Aggregation) or isinstance(
+ self.aggregation_function, Representation
+ ):
+ parameters["aggregation_function"] = self.aggregation_function.__class__
else:
- is_ts_rep = True
- parameters = {
- "aggregation_function": aggregation_function.name,
- "agg_params": aggregation_function.parameters,
- }
+ raise ValueError(
+ f"Invalid aggregation function type: {type(self.aggregation_function)}"
+ )
- super().__init__(name, parameters, is_ts_rep)
- self.aggregation_function = aggregation_function
+ super().__init__(name, parameters)
+ self.data_type = np.float32
+
+ def get_current_parameters(self):
+ current_params = {}
+ if not self.parameters:
+ return current_params
+
+ for parameter in list(self.parameters.keys()):
+ if parameter == "aggregation_function":
+ if "aggregation_function" in self.parameters:
+ current_params["aggregation_function"] = self.parameters[
+ "aggregation_function"
+ ]
+ if isinstance(self.aggregation_function, Aggregation) or isinstance(
+ self.aggregation_function, Representation
+ ):
+ for (
+ key,
+ value,
+ ) in self.aggregation_function.get_current_parameters().items():
+ current_params[f"aggregation_function_{key}"] = value
+ else:
+ current_params[parameter] = getattr(self, parameter)
+ return current_params
@property
def aggregation_function(self):
@@ -53,32 +78,117 @@ class Window(Context):
@aggregation_function.setter
def aggregation_function(self, value):
- if self.is_ts_rep:
+ if isinstance(value, Representation):
+ self._aggregation_function = value
+ elif isinstance(value, Aggregation):
self._aggregation_function = value
else:
self._aggregation_function = Aggregation(value)
+ def estimate_output_memory_bytes(self, input_stats: RepresentationStats) -> int:
+ return None
+
+ def estimate_peak_memory_bytes(self, input_stats: RepresentationStats) -> dict:
+ return None
+
+ def assert_output_stats(self, windowed_data):
+ if self.stats:
+ assert (
+ len(windowed_data) == self.stats.num_instances
+ ), f"Output shape: {windowed_data.shape}, Expected shape: {self.stats.output_shape}"
+ assert (
+ windowed_data.shape[1] == self.stats.output_shape[0]
+ ), f"Output shape: {windowed_data.shape}, Expected shape: {self.stats.output_shape}"
+ assert (
+ windowed_data.shape[2:] == self.stats.output_shape[1:]
+ ), f"Output shape: {windowed_data.shape}, Expected shape: {self.stats.output_shape}"
+
+ @staticmethod
+ def _shape_numel(shape):
+ return int(np.prod(shape)) if len(shape) > 0 else 1
+
+ @staticmethod
+ def _rest_numel(shape):
+ return int(np.prod(shape[1:])) if len(shape) > 1 else 1
+
@register_context_operator(
[ModalityType.TIMESERIES, ModalityType.AUDIO, ModalityType.EMBEDDING]
)
class WindowAggregation(Window):
- def __init__(self, aggregation_function="mean", window_size=10, pad=False):
+ def __init__(
+ self, aggregation_function="mean", window_size=10, pad=True, params=None
+ ):
+ if params is not None:
+ aggregation_function = params["aggregation_function"]
+ try:
+ aggregation_function = aggregation_function()
+ except:
+ pass
+ window_size = params["window_size"]
+ pad = True
super().__init__("WindowAggregation", aggregation_function)
self.parameters["window_size"] = [5, 10, 15, 25, 50, 100]
self.window_size = int(window_size)
self.pad = pad
+ def get_output_stats(self, input_stats: RepresentationStats) -> tuple:
+ in_shape = tuple(int(s) for s in input_stats.output_shape)
+ if len(in_shape) == 1:
+ self.stats = RepresentationStats(
+ input_stats.num_instances,
+ (math.ceil(in_shape[0] / self.window_size),),
+ )
+ if len(in_shape) >= 2:
+ self.stats = RepresentationStats(
+ input_stats.num_instances,
+ (math.ceil(in_shape[0] / self.window_size), *in_shape[1:]),
+ )
+ self.stats.output_shape_is_known = input_stats.output_shape_is_known
+ return self.stats
+
+ def estimate_output_memory_bytes(self, input_stats: RepresentationStats) -> int:
+ in_shape = tuple(int(s) for s in input_stats.output_shape)
+ if len(in_shape) == 0:
+ return 0
+
+ out_seq_len = math.ceil(in_shape[0] / self.window_size)
+ output_bytes = out_seq_len * self._rest_numel(in_shape)
+ return (
+ input_stats.num_instances * output_bytes * np.dtype(self.data_type).itemsize
+ )
+
+ def estimate_peak_memory_bytes(self, input_stats: RepresentationStats) -> dict:
+ in_shape = tuple(int(s) for s in input_stats.output_shape)
+ if len(in_shape) == 0:
+ return {"cpu_peak_bytes": 0, "gpu_peak_bytes": 0}
+
+ effective_seq_len = in_shape[0]
+ in_numel = effective_seq_len * self._rest_numel(in_shape)
+ output_bytes = self.estimate_output_memory_bytes(input_stats)
+ one_instance_bytes = in_numel * np.dtype(self.data_type).itemsize
+ cpu_peak = (
+ output_bytes * 2
+ + one_instance_bytes * input_stats.num_instances
+ + one_instance_bytes * self.window_size
+ + 8 * 1024 * 1024
+ )
+ return {"cpu_peak_bytes": cpu_peak, "gpu_peak_bytes": 0}
+
def execute(self, modality):
windowed_data = []
original_lengths = []
for instance in modality.data:
new_length = math.ceil(len(instance) / self.window_size)
if modality.get_data_layout() == DataLayout.SINGLE_LEVEL:
+ instance = np.array(instance)
+ instance.setflags(write=False)
windowed_instance = self.window_aggregate_single_level(
instance, new_length
)
else:
+ instance = np.array(instance)
+ instance.setflags(write=False)
windowed_instance = self.window_aggregate_nested_level(
instance, new_length
)
@@ -88,7 +198,6 @@ class WindowAggregation(Window):
if self.pad and not isinstance(windowed_data, np.ndarray):
target_length = max(original_lengths)
sample_shape = windowed_data[0].shape
- is_1d = len(sample_shape) == 1
padded_features = []
for i, features in enumerate(windowed_data):
@@ -97,13 +206,9 @@ class WindowAggregation(Window):
if current_len < target_length:
padding_needed = target_length - current_len
- if is_1d:
- padding = np.zeros(padding_needed)
- padded = np.concatenate([features, padding])
- else:
- feature_dim = features.shape[-1]
- padding = np.zeros((padding_needed, feature_dim))
- padded = np.concatenate([features, padding], axis=0)
+ pad_shape = (padding_needed,) + features.shape[1:]
+ padding = np.zeros(pad_shape)
+ padded = np.concatenate([features, padding], axis=0)
padded_features.append(padded)
else:
@@ -123,48 +228,31 @@ class WindowAggregation(Window):
if data_type != "str":
windowed_data = windowed_data.astype(data_type)
+ self.assert_output_stats(windowed_data)
return windowed_data
def window_aggregate_single_level(self, instance, new_length):
if isinstance(instance, str):
return instance
- instance = np.array(copy.deepcopy(instance))
result = []
for i in range(0, new_length):
- if self.is_ts_rep:
- result.append(
- self.aggregation_function.compute_feature(
- instance[
- i * self.window_size : i * self.window_size
- + self.window_size
- ]
- )
- )
- else:
- result.append(
- self.aggregation_function.aggregate_instance(
- instance[
- i * self.window_size : i * self.window_size
- + self.window_size
- ]
- )
+ result.append(
+ self.aggregation_function.compute_feature(
+ instance[
+ i * self.window_size : i * self.window_size + self.window_size
+ ]
)
+ )
return np.array(result)
def window_aggregate_nested_level(self, instance, new_length):
result = [[] for _ in range(0, new_length)]
- data = np.stack(copy.deepcopy(instance))
for i in range(0, new_length):
- if self.is_ts_rep:
- result[i] = self.aggregation_function.compute_feature(
- data[i * self.window_size : i * self.window_size + self.window_size]
- )
- else:
- result[i] = self.aggregation_function.aggregate_instance(
- data[i * self.window_size : i * self.window_size + self.window_size]
- )
+ result[i] = self.aggregation_function.compute_feature(
+ instance[i * self.window_size : i * self.window_size + self.window_size]
+ )
return np.array(result)
@@ -173,11 +261,52 @@ class WindowAggregation(Window):
[ModalityType.TIMESERIES, ModalityType.AUDIO, ModalityType.EMBEDDING]
)
class StaticWindow(Window):
- def __init__(self, aggregation_function="mean", num_windows=100):
+ # TODO
+ def __init__(self, aggregation_function="mean", num_windows=100, params=None):
super().__init__("StaticWindow", aggregation_function)
self.parameters["num_windows"] = [10, num_windows]
self.num_windows = int(num_windows)
+ def get_output_stats(self, input_stats: RepresentationStats) -> tuple:
+ in_shape = tuple(int(s) for s in input_stats.output_shape)
+ if len(in_shape) <= 1:
+ self.stats = RepresentationStats(
+ input_stats.num_instances, (self.num_windows,)
+ )
+ else:
+ self.stats = RepresentationStats(
+ input_stats.num_instances, (self.num_windows, *in_shape[1:])
+ )
+ self.stats.output_shape_is_known = input_stats.output_shape_is_known
+
+ return self.stats
+
+ def estimate_output_memory_bytes(self, input_stats: RepresentationStats) -> int:
+ in_shape = tuple(int(s) for s in input_stats.output_shape)
+ if len(in_shape) == 0:
+ return 0
+
+ out_seq_len = self.num_windows
+ output_bytes = out_seq_len * self._rest_numel(in_shape)
+ return (
+ input_stats.num_instances * output_bytes * np.dtype(self.data_type).itemsize
+ )
+
+ def estimate_peak_memory_bytes(self, input_stats: RepresentationStats) -> dict:
+ in_shape = tuple(int(s) for s in input_stats.output_shape)
+ if len(in_shape) == 0:
+ return {"cpu_peak_bytes": 0, "gpu_peak_bytes": 0}
+ effective_seq_len = in_shape[0]
+ in_numel = effective_seq_len * self._rest_numel(in_shape)
+ output_bytes = self.estimate_output_memory_bytes(input_stats)
+ one_instance_bytes = in_numel * np.dtype(self.data_type).itemsize
+ cpu_peak = (
+ output_bytes * 2
+ + one_instance_bytes * input_stats.num_instances
+ + 8 * 1024 * 1024
+ )
+ return {"cpu_peak_bytes": cpu_peak, "gpu_peak_bytes": 0}
+
def execute(self, modality):
windowed_data = []
@@ -189,9 +318,10 @@ class StaticWindow(Window):
for i in range(0, self.num_windows):
extra = 1 if i < remainder else 0
end = start + window_size + extra
- window = copy.deepcopy(instance[start:end])
+ window = instance[start:end]
+ window.setflags(write=False)
val = (
- self.aggregation_function.aggregate_instance(window)
+ self.aggregation_function.compute_feature(window)
if len(window) > 0
else np.zeros_like(output[i - 1])
)
@@ -199,18 +329,59 @@ class StaticWindow(Window):
start = end
windowed_data.append(output)
- return np.array(windowed_data)
+ windowed_data = np.array(windowed_data)
+ self.assert_output_stats(windowed_data)
+ return windowed_data
@register_context_operator(
[ModalityType.TIMESERIES, ModalityType.AUDIO, ModalityType.EMBEDDING]
)
class DynamicWindow(Window):
- def __init__(self, aggregation_function="mean", num_windows=100):
+ def __init__(self, aggregation_function="mean", num_windows=100, params=None):
super().__init__("DynamicWindow", aggregation_function)
self.parameters["num_windows"] = [10, num_windows]
self.num_windows = int(num_windows)
+ def get_output_stats(self, input_stats: RepresentationStats) -> tuple:
+ in_shape = tuple(int(s) for s in input_stats.output_shape)
+ if len(in_shape) <= 1:
+ self.stats = RepresentationStats(
+ input_stats.num_instances, (self.num_windows,)
+ )
+ else:
+ self.stats = RepresentationStats(
+ input_stats.num_instances, (self.num_windows, *in_shape[1:])
+ )
+ self.stats.output_shape_is_known = input_stats.output_shape_is_known
+ return self.stats
+
+ def estimate_output_memory_bytes(self, input_stats: RepresentationStats) -> int:
+ in_shape = tuple(int(s) for s in input_stats.output_shape)
+ if len(in_shape) == 0:
+ return 0
+
+ out_seq_len = self.num_windows
+ output_bytes = out_seq_len * self._rest_numel(in_shape)
+ return (
+ input_stats.num_instances * output_bytes * np.dtype(self.data_type).itemsize
+ )
+
+ def estimate_peak_memory_bytes(self, input_stats: RepresentationStats) -> dict:
+ in_shape = tuple(int(s) for s in input_stats.output_shape)
+ if len(in_shape) == 0:
+ return {"cpu_peak_bytes": 0, "gpu_peak_bytes": 0}
+ effective_seq_len = in_shape[0]
+ in_numel = effective_seq_len * self._rest_numel(in_shape)
+ output_bytes = self.estimate_output_memory_bytes(input_stats)
+ one_instance_bytes = in_numel * np.dtype(self.data_type).itemsize
+ cpu_peak = (
+ output_bytes * 2
+ + one_instance_bytes * input_stats.num_instances
+ + 8 * 1024 * 1024
+ )
+ return {"cpu_peak_bytes": cpu_peak, "gpu_peak_bytes": 0}
+
def execute(self, modality):
windowed_data = []
@@ -224,9 +395,10 @@ class DynamicWindow(Window):
output = []
start = 0
for end in indices:
- window = copy.deepcopy(instance[start:end])
+ window = instance[start:end]
+ window.setflags(write=False)
val = (
- self.aggregation_function.aggregate_instance(window)
+ self.aggregation_function.compute_feature(window)
if len(window) > 0
else np.zeros_like(instance[0])
)
@@ -234,5 +406,6 @@ class DynamicWindow(Window):
start = end
windowed_data.append(output)
-
- return np.array(windowed_data)
+ windowed_data = np.array(windowed_data)
+ self.assert_output_stats(windowed_data)
+ return windowed_data
diff --git a/src/main/python/systemds/scuro/representations/word2vec.py b/src/main/python/systemds/scuro/representations/word2vec.py
index 737d72b8b0..a744bc8db3 100644
--- a/src/main/python/systemds/scuro/representations/word2vec.py
+++ b/src/main/python/systemds/scuro/representations/word2vec.py
@@ -19,7 +19,9 @@
#
# -------------------------------------------------------------
import numpy as np
+from systemds.scuro.dataloader.text_loader import TextStats
from systemds.scuro.modality.transformed import TransformedModality
+from systemds.scuro.representations.representation import RepresentationStats
from systemds.scuro.representations.unimodal import UnimodalRepresentation
from systemds.scuro.representations.utils import save_embeddings
from gensim.models import Word2Vec
@@ -41,7 +43,7 @@ def get_embedding(sentence, model):
@register_representation(ModalityType.TEXT)
class W2V(UnimodalRepresentation):
- def __init__(self, vector_size=150, min_count=1, output_file=None):
+ def __init__(self, vector_size=150, min_count=1, output_file=None, params=None):
parameters = {
"vector_size": [50, 100, 150, 200],
"min_count": [1, 2, 4, 8],
@@ -50,8 +52,29 @@ class W2V(UnimodalRepresentation):
self.vector_size = vector_size
self.min_count = min_count
self.output_file = output_file
+ self.data_type = np.float32
- def transform(self, modality):
+ def get_output_stats(self, input_stats: TextStats) -> RepresentationStats:
+ return RepresentationStats(input_stats.num_instances, (self.vector_size,))
+
+ def estimate_output_memory_bytes(self, input_stats: TextStats) -> int:
+ return (
+ input_stats.num_instances
+ * self.vector_size
+ * np.dtype(self.data_type).itemsize
+ )
+
+ def estimate_peak_memory_bytes(self, input_stats: TextStats) -> dict:
+ output_bytes = (
+ self.estimate_output_memory_bytes(input_stats)
+ + input_stats.num_instances
+ * input_stats.max_length
+ * np.dtype(self.data_type).itemsize
+ ) * 2
+ model_estimate = 100 * 1024 * 1024
+ return {"cpu_peak_bytes": model_estimate + output_bytes, "gpu_peak_bytes": 0}
+
+ def transform(self, modality, params=None):
transformed_modality = TransformedModality(modality, self)
t = [list(tokenize(s.lower())) for s in modality.data]
model = Word2Vec(
@@ -66,6 +89,6 @@ class W2V(UnimodalRepresentation):
if self.output_file is not None:
save_embeddings(np.array(embeddings), self.output_file)
- transformed_modality.data_type = np.float32
+ transformed_modality.data_type = self.data_type
transformed_modality.data = np.array(embeddings)
return transformed_modality
diff --git a/src/main/python/systemds/scuro/representations/x3d.py b/src/main/python/systemds/scuro/representations/x3d.py
index 7701865f82..f7a7092153 100644
--- a/src/main/python/systemds/scuro/representations/x3d.py
+++ b/src/main/python/systemds/scuro/representations/x3d.py
@@ -18,10 +18,15 @@
# under the License.
#
# -------------------------------------------------------------
-from systemds.scuro.utils.static_variables import get_device
+from systemds.scuro.utils.static_variables import (
+ compute_batch_size,
+ get_device,
+ get_device_for_model,
+)
from systemds.scuro.utils.torch_dataset import CustomDataset
from systemds.scuro.modality.transformed import TransformedModality
from systemds.scuro.representations.unimodal import UnimodalRepresentation
+from systemds.scuro.representations.representation import RepresentationStats
from typing import Tuple, Any
import torch.utils.data
import torch
@@ -34,7 +39,9 @@ from systemds.scuro.drsearch.operator_registry import register_representation
@register_representation([ModalityType.VIDEO])
class X3D(UnimodalRepresentation):
- def __init__(self, layer="classifier.1", model_name="s3d", output_file=None):
+ def __init__(
+ self, layer="classifier.1", model_name="s3d", output_file=None, params=None
+ ):
self.data_type = torch.float32
self.model_name = model_name
parameters = self._get_parameters()
@@ -52,6 +59,9 @@ class X3D(UnimodalRepresentation):
self.model.fc = Identity()
+ def get_output_stats(self, input_stats) -> RepresentationStats:
+ return RepresentationStats(input_stats.num_instances, (512,))
+
@property
def model_name(self):
return self._model_name
@@ -60,9 +70,13 @@ class X3D(UnimodalRepresentation):
def model_name(self, model_name):
self._model_name = model_name
if model_name == "r3d":
- self.model = r3d_18(pretrained=True).to(get_device())
+ self.model = r3d_18(pretrained=True)
+ self.device = get_device_for_model(self.model, memory_factor=1.5)
+ self.model = self.model.to(self.device)
elif model_name == "s3d":
- self.model = s3d(weights=models.video.S3D_Weights.DEFAULT).to(get_device())
+ self.model = s3d(weights=models.video.S3D_Weights.DEFAULT)
+ self.device = get_device_for_model(self.model, memory_factor=1.5)
+ self.model = self.model.to(self.device)
else:
raise NotImplementedError
@@ -97,8 +111,17 @@ class X3D(UnimodalRepresentation):
parameters["layer_name"].append(name)
return parameters
- def transform(self, modality):
- dataset = CustomDataset(modality.data, self.data_type, get_device())
+ def transform(self, modality, aggregation=None):
+ sample = modality.data[0] if modality.data else ""
+ self.batch_size = compute_batch_size(
+ model=self.model,
+ device=self.device,
+ sample_data=sample,
+ tokenizer=None,
+ max_seq_length=None,
+ max_batch_size=128,
+ )
+ dataset = CustomDataset(modality.data, self.data_type, self.device)
embeddings = {}
@@ -121,7 +144,7 @@ class X3D(UnimodalRepresentation):
for instance in dataset:
video_id = instance["id"]
- frames = instance["data"].to(get_device())
+ frames = instance["data"].to(self.device)
embeddings[video_id] = []
frames = frames.unsqueeze(0).permute(0, 2, 1, 3, 4)
@@ -152,7 +175,9 @@ class I3D(UnimodalRepresentation):
self.model_name = model_name
self.model = torch.hub.load(
"facebookresearch/pytorchvideo", "i3d_r50", pretrained=True
- ).to(get_device())
+ )
+ self.device = get_device_for_model(self.model, memory_factor=1.5)
+ self.model = self.model.to(self.device)
parameters = self._get_parameters()
super().__init__("I3D", ModalityType.EMBEDDING, parameters)
@@ -181,7 +206,16 @@ class I3D(UnimodalRepresentation):
return parameters
def transform(self, modality):
- dataset = CustomDataset(modality.data, torch.float32, get_device())
+ sample = modality.data[0] if modality.data else ""
+ self.batch_size = compute_batch_size(
+ model=self.model,
+ device=self.device,
+ sample_data=sample,
+ tokenizer=None,
+ max_seq_length=None,
+ max_batch_size=128,
+ )
+ dataset = CustomDataset(modality.data, torch.float32, self.device)
embeddings = {}
features = None
@@ -204,7 +238,7 @@ class I3D(UnimodalRepresentation):
for instance in dataset:
video_id = instance["id"]
- frames = instance["data"].to(get_device())
+ frames = instance["data"].to(self.device)
embeddings[video_id] = []
batch = torch.transpose(frames, 1, 0)
diff --git a/src/main/python/systemds/scuro/utils/checkpointing.py b/src/main/python/systemds/scuro/utils/checkpointing.py
new file mode 100644
index 0000000000..e821798534
--- /dev/null
+++ b/src/main/python/systemds/scuro/utils/checkpointing.py
@@ -0,0 +1,164 @@
+# -------------------------------------------------------------
+#
+# 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 os
+import pickle
+import time
+from typing import Any, Callable, Dict, Optional, Tuple
+
+
+class CheckpointManager:
+ def __init__(
+ self,
+ checkpoint_dir: str,
+ prefix: str,
+ checkpoint_every: Optional[int] = None,
+ resume: bool = True,
+ ):
+ self.checkpoint_dir = checkpoint_dir or "."
+ self.prefix = prefix
+ self.resume_enabled = resume
+ self.checkpoint_every = checkpoint_every
+ self.eval_count = 0
+ self.counts_by_key: Dict[Any, int] = {}
+ self.skip_remaining_by_key: Dict[Any, int] = {}
+ self._last_checkpoint_eval_count = 0
+ self._resumed = False
+
+ def _checkpoint_path(self) -> str:
+ return os.path.join(self.checkpoint_dir, f"{self.prefix}latest.pkl")
+
+ def _meta_path(self, checkpoint_path: str) -> str:
+ return f"{checkpoint_path}.meta.pkl"
+
+ def find_latest_checkpoint(self) -> Optional[str]:
+ if not os.path.isdir(self.checkpoint_dir):
+ return None
+ checkpoint_path = self._checkpoint_path()
+ if os.path.isfile(checkpoint_path):
+ return checkpoint_path
+ return None
+
+ def load_latest(self) -> Optional[Tuple[Any, Dict[str, Any], str]]:
+ if not self.resume_enabled or self._resumed:
+ return None
+ checkpoint_path = self.find_latest_checkpoint()
+ if not checkpoint_path:
+ self._resumed = True
+ return None
+
+ with open(checkpoint_path, "rb") as f:
+ results = pickle.load(f)
+
+ meta_path = self._meta_path(checkpoint_path)
+ meta: Dict[str, Any] = {}
+ if os.path.isfile(meta_path):
+ with open(meta_path, "rb") as f:
+ meta = pickle.load(f)
+
+ self._resumed = True
+ return results, meta, checkpoint_path
+
+ def resume_from_checkpoint(
+ self,
+ count_key: str,
+ count_fallback_fn: Callable[[Any], Dict[Any, int]],
+ ) -> Optional[Tuple[Any, Dict[str, Any], str]]:
+ loaded = self.load_latest()
+ if not loaded:
+ self._resumed = True
+ return None
+
+ results, meta, checkpoint_path = loaded
+ counts_by_key = meta.get(count_key, {})
+ eval_count = meta.get("eval_count", 0)
+
+ counts_by_key = count_fallback_fn(results)
+ eval_count = sum(counts_by_key.values())
+ if not eval_count:
+ eval_count = sum(counts_by_key.values())
+
+ self.eval_count = eval_count
+ self.counts_by_key = counts_by_key
+ self.skip_remaining_by_key = dict(counts_by_key)
+ self._resumed = True
+ return results, meta, checkpoint_path
+
+ def increment(self, key: Any, count: int = 1, **kwargs):
+ if count <= 0:
+ return
+ self.eval_count += count
+ self.counts_by_key[key] = self.counts_by_key.get(key, 0) + count
+
+ def save(self, results: Any, meta: Dict[str, Any]) -> str:
+ try:
+ os.makedirs(self.checkpoint_dir, exist_ok=True)
+ checkpoint_path = self._checkpoint_path()
+ with open(checkpoint_path, "wb") as f:
+ pickle.dump(results, f)
+ with open(self._meta_path(checkpoint_path), "wb") as f:
+ pickle.dump(meta, f)
+ except Exception as e:
+ print(f"Error saving checkpoint: {e}")
+ return None
+ return checkpoint_path
+
+ def save_checkpoint(self, results: Any, extra_meta: Dict[str, Any] = {}):
+ meta = {"eval_count": self.eval_count}
+ # meta.update(extra_meta or {})
+ self.save(results, meta)
+
+ def checkpoint_if_due(self, results: Any, extra_meta: Dict[str, Any] = None):
+ if not self.checkpoint_every:
+ return
+ if self.eval_count <= 0:
+ return
+ if self.eval_count == self._last_checkpoint_eval_count:
+ return
+ if self.eval_count % self.checkpoint_every == 0:
+ self._last_checkpoint_eval_count = self.eval_count
+ self.save_checkpoint(results, extra_meta or {})
+
+ def cleanup(self):
+ """Remove the checkpoint and meta files managed by this instance."""
+ checkpoint_path = self._checkpoint_path()
+ meta_path = self._meta_path(checkpoint_path)
+ for path in [checkpoint_path, meta_path]:
+ if os.path.isfile(path):
+ try:
+ os.remove(path)
+ except OSError as e:
+ print(f"Warning: Could not remove checkpoint file {path}: {e}")
+
+ @staticmethod
+ def cleanup_by_prefix(checkpoint_dir, prefix):
+ """Remove all checkpoint files whose name starts with the given prefix."""
+ checkpoint_dir = checkpoint_dir or "."
+ if not os.path.isdir(checkpoint_dir):
+ return
+ for filename in os.listdir(checkpoint_dir):
+ if filename.startswith(prefix):
+ filepath = os.path.join(checkpoint_dir, filename)
+ if os.path.isfile(filepath):
+ try:
+ os.remove(filepath)
+ except OSError as e:
+ print(f"Warning: Could not remove {filepath}: {e}")
diff --git a/src/main/python/systemds/scuro/utils/memory_utility.py b/src/main/python/systemds/scuro/utils/memory_utility.py
new file mode 100644
index 0000000000..0d3cd9d786
--- /dev/null
+++ b/src/main/python/systemds/scuro/utils/memory_utility.py
@@ -0,0 +1,194 @@
+# -------------------------------------------------------------
+#
+# 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 resource
+import sys
+import numpy as np
+from sympy import Dict
+import torch
+from typing import List, Tuple
+import psutil
+
+
+def get_model_size_mb(model: torch.nn.Module) -> float:
+ param_bytes = sum(p.numel() * p.element_size() for p in model.parameters())
+ buffer_bytes = sum(b.numel() * b.element_size() for b in model.buffers())
+ return (param_bytes + buffer_bytes) / (1024**2)
+
+
+def get_gpu_free_memory_mb(device_index: int = 0) -> float:
+ free, _ = torch.cuda.mem_get_info(device_index)
+ return free / (1024**2)
+
+
+def get_best_gpu() -> Tuple[int, float]:
+ num_gpus = torch.cuda.device_count()
+ if num_gpus == 0:
+ return -1, 0.0
+
+ best_gpu = 0
+ best_free = 0.0
+ for i in range(num_gpus):
+ free_mb = get_gpu_free_memory_mb(i)
+ if free_mb > best_free:
+ best_free = free_mb
+ best_gpu = i
+
+ return best_gpu, best_free
+
+
+def get_device_for_model(
+ model: torch.nn.Module, memory_factor: float = 1.5
+) -> torch.device:
+ if not torch.cuda.is_available():
+ return torch.device("cpu")
+
+ model_mb = get_model_size_mb(model)
+ required_mb = model_mb * memory_factor
+
+ gpu_idx, free_mb = get_best_gpu()
+
+ if free_mb >= required_mb:
+ return torch.device(f"cuda:{gpu_idx}")
+
+ return torch.device("cpu")
+
+
+def estimate_batch_memory_mb(
+ sample_data, tokenizer=None, max_seq_length=512, dtype=torch.float32
+) -> float:
+ bytes_per_element = torch.tensor([], dtype=dtype).element_size()
+
+ if tokenizer is not None:
+ num_tokens = max_seq_length
+ num_input_tensors = 3
+ input_bytes = num_tokens * num_input_tensors * 8 # int64
+
+ activation_bytes = num_tokens * 768 * 4 * bytes_per_element
+
+ return (input_bytes + activation_bytes) / (1024**2)
+
+ elif isinstance(sample_data, torch.Tensor):
+ return sample_data.numel() * bytes_per_element / (1024**2)
+
+ elif isinstance(sample_data, np.ndarray):
+ return sample_data.nbytes / (1024**2)
+
+ else:
+ return 1.0
+
+
+def compute_batch_size(
+ model: torch.nn.Module,
+ device: torch.device,
+ sample_data,
+ tokenizer=None,
+ max_seq_length: int = 512,
+ dtype=torch.float32,
+ min_batch_size: int = 1,
+ max_batch_size: int = 128,
+ memory_fraction: float = 0.8,
+) -> int:
+
+ if device.type == "cpu":
+ return max_batch_size
+
+ gpu_idx = device.index if device.index is not None else 0
+ free_mb = get_gpu_free_memory_mb(gpu_idx)
+ usable_mb = free_mb * memory_fraction
+
+ per_sample_mb = estimate_batch_memory_mb(
+ sample_data, tokenizer=tokenizer, max_seq_length=max_seq_length, dtype=dtype
+ )
+
+ if per_sample_mb <= 0:
+ return max_batch_size
+
+ computed = int(usable_mb / per_sample_mb)
+
+ batch_size = max(min_batch_size, min(computed, max_batch_size))
+
+ return batch_size
+
+
+def get_device(gpu_id: int = None):
+ if gpu_id is not None:
+ return torch.device(f"cuda:{gpu_id}")
+ return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+def get_cpu_memory_mb():
+ rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
+ if __import__("sys").platform == "darwin":
+ rss /= 1024 # macOS reports bytes
+ return rss / 1024 # MB
+
+
+def get_gpu_memory_mb(device):
+ if not torch.cuda.is_available() or device.type == "cpu":
+ return 0.0
+ torch.cuda.synchronize(device)
+ return torch.cuda.memory_allocated(device) / (1024**2)
+
+
+def gpu_memory_info():
+ infos = []
+ num_gpus = torch.cuda.device_count()
+ for i in range(num_gpus):
+ torch.cuda.set_device(i)
+ free_b, total_b = torch.cuda.mem_get_info()
+ infos.append(dict(index=i, free_b=free_b, total_b=total_b))
+ return infos
+
+
+def log_memory(operation_name: str, device: torch.device):
+ cpu_mb = get_cpu_memory_mb()
+ gpu_mb = get_gpu_memory_mb(device)
+ current_mb = psutil.Process().memory_info().rss / (1024**2)
+ print(
+ f"[mem] {operation_name}: CPU={cpu_mb:.1f} MB, GPU={gpu_mb:.1f} MB, Current={current_mb:.1f} MB"
+ )
+
+
+def estimate_numpy_like_bytes(data) -> int:
+ if data is None:
+ return 0
+ if isinstance(data, np.ndarray):
+ return int(data.nbytes)
+ if isinstance(data, torch.Tensor):
+ return int(data.numel() * data.element_size())
+ if isinstance(data, memoryview):
+ return int(data.nbytes)
+ if isinstance(data, (bytes, bytearray)):
+ return int(len(data))
+ if isinstance(data, dict):
+ return int(sum(estimate_numpy_like_bytes(v) for v in data.values()))
+ if isinstance(data, (list, tuple)):
+ return int(sum(estimate_numpy_like_bytes(v) for v in data))
+ return int(sys.getsizeof(data))
+
+
+def estimate_modality_bytes(modality) -> int:
+ if modality is None:
+ return 0
+ data_bytes = estimate_numpy_like_bytes(getattr(modality, "data", None))
+ metadata = getattr(modality, "metadata", None)
+ metadata_bytes = estimate_numpy_like_bytes(metadata)
+ return int(data_bytes + metadata_bytes)
diff --git a/src/main/python/systemds/scuro/utils/static_variables.py b/src/main/python/systemds/scuro/utils/static_variables.py
index 807287cd95..0da186348e 100644
--- a/src/main/python/systemds/scuro/utils/static_variables.py
+++ b/src/main/python/systemds/scuro/utils/static_variables.py
@@ -20,6 +20,13 @@
# -------------------------------------------------------------
import numpy as np
import torch
+from typing import Tuple
+
+PY_LIST_HEADER_BYTES = 56
+PY_LIST_SLOT_BYTES = 8
+NP_ARRAY_HEADER_BYTES = 112
+
+DEBUG = False
global_rng = np.random.default_rng(42)
@@ -28,10 +35,108 @@ def get_seed():
return global_rng.integers(0, 1024)
-def get_device():
- return torch.device(
- "cuda:0"
- if torch.cuda.is_available()
- # else "mps" if torch.mps.is_available()
- else "cpu"
+def get_model_size_mb(model: torch.nn.Module) -> float:
+ param_bytes = sum(p.numel() * p.element_size() for p in model.parameters())
+ buffer_bytes = sum(b.numel() * b.element_size() for b in model.buffers())
+ return (param_bytes + buffer_bytes) / (1024**2)
+
+
+def get_gpu_free_memory_mb(device_index: int = 0) -> float:
+ free, _ = torch.cuda.mem_get_info(device_index)
+ return free / (1024**2)
+
+
+def get_best_gpu() -> Tuple[int, float]:
+ num_gpus = torch.cuda.device_count()
+ if num_gpus == 0:
+ return -1, 0.0
+
+ best_gpu = 0
+ best_free = 0.0
+ for i in range(num_gpus):
+ free_mb = get_gpu_free_memory_mb(i)
+ if free_mb > best_free:
+ best_free = free_mb
+ best_gpu = i
+
+ return best_gpu, best_free
+
+
+def get_device_for_model(
+ model: torch.nn.Module, memory_factor: float = 1.5
+) -> torch.device:
+ if not torch.cuda.is_available():
+ return torch.device("cpu")
+
+ model_mb = get_model_size_mb(model)
+ required_mb = model_mb * memory_factor
+
+ gpu_idx, free_mb = get_best_gpu()
+
+ if free_mb >= required_mb:
+ return torch.device(f"cuda:{gpu_idx}")
+
+ return torch.device("cpu")
+
+
+def estimate_batch_memory_mb(
+ sample_data, tokenizer=None, max_seq_length=512, dtype=torch.float32
+) -> float:
+ bytes_per_element = torch.tensor([], dtype=dtype).element_size()
+
+ if tokenizer is not None:
+ num_tokens = max_seq_length
+ num_input_tensors = 3
+ input_bytes = num_tokens * num_input_tensors * 8 # int64
+
+ activation_bytes = num_tokens * 768 * 4 * bytes_per_element
+
+ return (input_bytes + activation_bytes) / (1024**2)
+
+ elif isinstance(sample_data, torch.Tensor):
+ return sample_data.numel() * bytes_per_element / (1024**2)
+
+ elif isinstance(sample_data, np.ndarray):
+ return sample_data.nbytes / (1024**2)
+
+ else:
+ return 1.0
+
+
+def compute_batch_size(
+ model: torch.nn.Module,
+ device: torch.device,
+ sample_data,
+ tokenizer=None,
+ max_seq_length: int = 512,
+ dtype=torch.float32,
+ min_batch_size: int = 1,
+ max_batch_size: int = 128,
+ memory_fraction: float = 0.8,
+) -> int:
+
+ if device.type == "cpu":
+ return max_batch_size
+
+ gpu_idx = device.index if device.index is not None else 0
+ free_mb = get_gpu_free_memory_mb(gpu_idx)
+ usable_mb = free_mb * memory_fraction
+
+ per_sample_mb = estimate_batch_memory_mb(
+ sample_data, tokenizer=tokenizer, max_seq_length=max_seq_length, dtype=dtype
)
+
+ if per_sample_mb <= 0:
+ return max_batch_size
+
+ computed = int(usable_mb / per_sample_mb)
+
+ batch_size = max(min_batch_size, min(computed, max_batch_size))
+
+ return batch_size
+
+
+def get_device(gpu_id: int = None):
+ if gpu_id is not None:
+ return torch.device(f"cuda:{gpu_id}")
+ return torch.device("cuda" if torch.cuda.is_available() else "cpu")
diff --git a/src/main/python/systemds/scuro/utils/torch_dataset.py b/src/main/python/systemds/scuro/utils/torch_dataset.py
index ba3e24a317..fea59088d7 100644
--- a/src/main/python/systemds/scuro/utils/torch_dataset.py
+++ b/src/main/python/systemds/scuro/utils/torch_dataset.py
@@ -56,16 +56,14 @@ class CustomDataset(torch.utils.data.Dataset):
def __getitem__(self, index) -> Dict[str, object]:
data = self.data[index]
- output = torch.empty(
- (len(data), 3, self.size[1], self.size[0]),
- dtype=self.data_type,
- device=self.device,
- )
if isinstance(data, np.ndarray) and data.ndim == 3:
# image
- output = self.tf(data).to(self.device)
+ output = self.tf(data)
else:
+ output = torch.empty(
+ (len(data), 3, self.size[1], self.size[0]), dtype=self.data_type
+ )
for i, d in enumerate(data):
if data[0].ndim < 3:
d = torch.tensor(d)
@@ -76,7 +74,7 @@ class CustomDataset(torch.utils.data.Dataset):
tf = tf[:3, :, :]
output[i] = tf
- return {"id": index, "data": output}
+ return {"id": index, "data": output.to(self.device)}
def __len__(self) -> int:
return len(self.data)
diff --git a/src/main/python/tests/scuro/data_generator.py b/src/main/python/tests/scuro/data_generator.py
index cfa77b1dd6..b78ea31483 100644
--- a/src/main/python/tests/scuro/data_generator.py
+++ b/src/main/python/tests/scuro/data_generator.py
@@ -30,15 +30,21 @@ from sklearn import svm
from sklearn.metrics import classification_report
from sklearn.model_selection import train_test_split
+from systemds.scuro.dataloader.audio_loader import AudioStats
+from systemds.scuro.dataloader.image_loader import ImageStats
+from systemds.scuro.dataloader.text_loader import TextStats
+from systemds.scuro.dataloader.timeseries_loader import TimeseriesStats
+from systemds.scuro.dataloader.video_loader import VideoStats
from systemds.scuro.models.model import Model
from systemds.scuro.dataloader.base_loader import BaseLoader
-from systemds.scuro.dataloader.video_loader import VideoLoader
-from systemds.scuro.dataloader.audio_loader import AudioLoader
-from systemds.scuro.dataloader.text_loader import TextLoader
-from systemds.scuro.modality.unimodal_modality import UnimodalModality
+
from systemds.scuro.modality.transformed import TransformedModality
from systemds.scuro.modality.type import ModalityType
from systemds.scuro.drsearch.task import Task
+from systemds.scuro.representations.representation import RepresentationStats
+
+random_state = 42
+np.random.seed(random_state)
class TestDataLoader(BaseLoader):
@@ -47,6 +53,55 @@ class TestDataLoader(BaseLoader):
self.metadata = metadata
self.test_data = data
+ if modality_type == ModalityType.TEXT:
+ self.stats = TextStats(
+ len(data),
+ max(len(d) for d in data),
+ sum(len(d) for d in data) / len(data),
+ max(len(d.split(" ")) for d in data),
+ sum(len(d.split(" ")) for d in data) / len(data),
+ (max(len(d) for d in data),),
+ )
+ elif modality_type == ModalityType.AUDIO:
+ self.stats = AudioStats(
+ 16000,
+ max(len(d) for d in data),
+ sum(len(d) for d in data) / len(data),
+ len(data),
+ output_shape_is_known=True,
+ )
+ elif modality_type == ModalityType.VIDEO:
+ self.stats = VideoStats(
+ 30,
+ max(d.shape[0] for d in data),
+ max(d.shape[1] for d in data),
+ max(d.shape[2] for d in data),
+ max(d.shape[3] for d in data),
+ len(data),
+ )
+ elif modality_type == ModalityType.TIMESERIES:
+ self.stats = TimeseriesStats(
+ max(len(d) for d in data),
+ len(data),
+ sum(len(d) for d in data) / len(data),
+ (max(len(d) for d in data),),
+ True,
+ )
+ elif modality_type == ModalityType.IMAGE:
+ self.stats = ImageStats(
+ max(d.shape[0] for d in data),
+ max(d.shape[1] for d in data),
+ max(d.shape[2] for d in data),
+ len(data),
+ (
+ max(d.shape[0] for d in data),
+ max(d.shape[1] for d in data),
+ max(d.shape[2] for d in data),
+ ),
+ average_width=sum(d.shape[0] for d in data) / len(data),
+ average_height=sum(d.shape[1] for d in data) / len(data),
+ average_channels=sum(d.shape[2] for d in data) / len(data),
+ )
def reset(self):
self._next_chunk = 0
@@ -67,7 +122,58 @@ class ModalityRandomDataGenerator:
self.modality_type = None
self.metadata = {}
self.data_type = np.float32
- self.transform_time = None
+ self.transform_time = 0
+ self.stats = None
+
+ def create_stats(self, data):
+ if self.modality_type == ModalityType.TEXT:
+ self.stats = TextStats(
+ len(data),
+ max(len(d) for d in data),
+ sum(len(d) for d in data) / len(data),
+ max(len(d.split(" ")) for d in data),
+ sum(len(d.split(" ")) for d in data) / len(data),
+ (max(len(d) for d in data),),
+ )
+ elif self.modality_type == ModalityType.AUDIO:
+ self.stats = AudioStats(
+ 16000,
+ max(len(d) for d in data),
+ sum(len(d) for d in data) / len(data),
+ len(data),
+ )
+ elif self.modality_type == ModalityType.VIDEO:
+ self.stats = VideoStats(
+ 30,
+ max(d.shape[0] for d in data),
+ max(d.shape[1] for d in data),
+ max(d.shape[2] for d in data),
+ max(d.shape[3] for d in data),
+ len(data),
+ )
+ elif self.modality_type == ModalityType.TIMESERIES:
+ self.stats = TimeseriesStats(
+ len(data),
+ max(len(d) for d in data),
+ sum(len(d) for d in data) / len(data),
+ (max(len(d) for d in data),),
+ True,
+ )
+ elif self.modality_type == ModalityType.IMAGE:
+ self.stats = ImageStats(
+ max(d.shape[0] for d in data),
+ max(d.shape[1] for d in data),
+ max(d.shape[2] for d in data),
+ len(data),
+ (
+ max(d.shape[0] for d in data),
+ max(d.shape[1] for d in data),
+ max(d.shape[2] for d in data),
+ ),
+ )
+ else:
+ raise ValueError(f"Unsupported modality type: {self.modality_type}")
+ return self.stats
def create1DModality(
self,
@@ -104,7 +210,7 @@ class ModalityRandomDataGenerator:
return tf_modality
def create_audio_data(self, num_instances, max_audio_length):
- modality_type = ModalityType.AUDIO
+ self.modality_type = ModalityType.AUDIO
data = [
[
random.random()
@@ -116,29 +222,31 @@ class ModalityRandomDataGenerator:
for i in range(num_instances):
data[i] = np.array(data[i]).astype(self.data_type)
- metadata = {
- i: modality_type.create_metadata(16000, np.array(data[i]))
+ self.metadata = {
+ i: self.modality_type.create_metadata(16000, np.array(data[i]))
for i in range(num_instances)
}
- return data, metadata
+ return data, self.metadata
def create_timeseries_data(self, num_instances, sequence_length, num_features=1):
+ self.modality_type = ModalityType.TIMESERIES
data = [
np.random.rand(sequence_length, num_features).astype(self.data_type)
for _ in range(num_instances)
]
if num_features == 1:
data = [d.squeeze(-1) for d in data]
- metadata = {
- i: ModalityType.TIMESERIES.create_metadata(
+ self.metadata = {
+ i: self.modality_type.create_metadata(
[f"feature_{j}" for j in range(num_features)], data[i]
)
for i in range(num_instances)
}
- return data, metadata
+ return data, self.metadata
def create_text_data(self, num_instances, num_sentences_per_instance=1):
+ self.modality_type = ModalityType.TEXT
subjects = [
"The cat",
"A dog",
@@ -200,48 +308,77 @@ class ModalityRandomDataGenerator:
sentence += f" {verb} {obj}{punct}"
sentences.append(sentence)
- metadata = {
- i: ModalityType.TEXT.create_metadata(len(sentences[i]), sentences[i])
+ self.metadata = {
+ i: self.modality_type.create_metadata(len(sentences[i]), sentences[i])
for i in range(num_instances)
}
- return sentences, metadata
+ return sentences, self.metadata
+
+ def create_3d_modality(self, num_instances, dims=(100, 28, 28)):
+ self.modality_type = ModalityType.EMBEDDING
+ data = [
+ np.random.rand(dims[0], dims[1], dims[2]).astype(self.data_type)
+ for _ in range(num_instances)
+ ]
+ self.metadata = {
+ i: self.modality_type.create_metadata(data[i]) for i in range(num_instances)
+ }
+ return data, self.metadata
+
+ def create_2d_modality(self, num_instances, dims=(100, 28)):
+ self.modality_type = ModalityType.EMBEDDING
+ data = [
+ np.random.rand(dims[0], dims[1]).astype(self.data_type)
+ for _ in range(num_instances)
+ ]
+ self.metadata = {
+ i: self.modality_type.create_metadata(data[i]) for i in range(num_instances)
+ }
+ return data, self.metadata
def create_visual_modality(
- self, num_instances, max_num_frames=1, height=28, width=28
+ self, num_instances, max_num_frames=1, height=28, width=28, color_channels=3
):
if max_num_frames > 1:
+ self.modality_type = ModalityType.VIDEO
data = [
np.random.uniform(
0.0,
1.0,
- (np.random.randint(10, max_num_frames + 1), height, width, 3),
+ (
+ np.random.randint(10, max_num_frames + 1),
+ height,
+ width,
+ color_channels,
+ ),
)
for _ in range(num_instances)
]
- metadata = {
- i: ModalityType.VIDEO.create_metadata(
- 30, data[i].shape[0], width, height, 3
+ self.metadata = {
+ i: self.modality_type.create_metadata(
+ 30, data[i].shape[0], width, height, color_channels
)
for i in range(num_instances)
}
else:
+ self.modality_type = ModalityType.IMAGE
data = [
np.random.randint(
0,
256,
- (height, width, 3),
+ (height, width, color_channels),
dtype=np.uint8,
)
for _ in range(num_instances)
]
- metadata = {
- i: ModalityType.IMAGE.create_metadata(width, height, 3)
+ self.metadata = {
+ i: self.modality_type.create_metadata(width, height, color_channels)
for i in range(num_instances)
}
- return data, metadata
+ return data, self.metadata
def create_balanced_labels(self, num_instances, num_classes=2):
if num_instances % num_classes != 0:
@@ -256,7 +393,7 @@ class ModalityRandomDataGenerator:
return vector
-def setup_data(modalities, num_instances, path):
+def setup_data(modality_types, num_instances, path):
if os.path.isdir(path):
shutil.rmtree(path)
@@ -264,47 +401,31 @@ def setup_data(modalities, num_instances, path):
indizes = [str(i) for i in range(0, num_instances)]
- modalities_to_create = []
- for modality in modalities:
- mod_path = path + "/" + modality.name + "/"
-
- if modality == ModalityType.VIDEO:
- data_loader = VideoLoader(mod_path, indizes)
- elif modality == ModalityType.AUDIO:
- data_loader = AudioLoader(mod_path, indizes)
- elif modality == ModalityType.TEXT:
- data_loader = TextLoader(mod_path, indizes)
- else:
- raise "Modality not supported in DataGenerator"
-
- modalities_to_create.append(UnimodalModality(data_loader))
-
- data_generator = TestDataGenerator(modalities_to_create, path)
+ data_generator = TestDataGenerator(modality_types, path)
data_generator.create_multimodal_data(num_instances)
return data_generator
class TestDataGenerator:
- def __init__(self, modalities, path, balanced=True):
-
- self.modalities = modalities
- self.modalities_by_type = {}
- for modality in modalities:
- self.modalities_by_type[modality.modality_type] = modality
+ def __init__(self, modalities_types, path, balanced=True):
+ self.modalities_types = modalities_types
+ self.modalities_paths = {
+ modality_type.name: f"{path}/{modality_type.name}/"
+ for modality_type in modalities_types
+ }
self._indices = None
self.path = path
self.balanced = balanced
- for modality in modalities:
- mod_path = f"{self.path}/{modality.modality_type.name}/"
- os.mkdir(mod_path)
- modality.file_path = mod_path
+ for modality_path in self.modalities_paths.values():
+ os.mkdir(modality_path)
+
self.labels = []
self.label_path = f"{path}/labels.npy"
def get_modality_path(self, modality_type):
- return self.modalities_by_type[modality_type].data_loader.source_path
+ return self.modalities_paths[modality_type.name]
@property
def indices(self):
@@ -342,41 +463,64 @@ class TestDataGenerator:
else:
speed_slow += 1
- for modality in self.modalities:
- if modality.modality_type == ModalityType.VIDEO:
- self.__create_video_data(idx, duration, 30, speed_factor)
- if modality.modality_type == ModalityType.AUDIO:
+ for modality_type in self.modalities_types:
+ if modality_type == ModalityType.VIDEO:
+ self.__create_visual_data(
+ idx, duration, 30, speed_factor, modality_type
+ )
+ elif modality_type == ModalityType.IMAGE:
+ self.__create_visual_data(idx, 1, 1, speed_factor, modality_type)
+ elif modality_type == ModalityType.AUDIO:
self.__create_audio_data(idx, duration, speed_factor)
- if modality.modality_type == ModalityType.TEXT:
+ elif modality_type == ModalityType.TEXT:
self.__create_text_data(idx, speed_factor)
+ elif modality_type == ModalityType.TIMESERIES:
+ self.__create_timeseries_data(idx, duration, speed_factor)
+ else:
+ raise ValueError(f"Unsupported modality type: {modality_type}")
- np.save(f"{self.path}/labels.npy", np.array(self.labels))
-
- def __create_video_data(self, idx, duration, fps, speed_factor):
- path = f"{self.path}/VIDEO/{idx}.mp4"
+ np.save(self.label_path, np.array(self.labels))
- width, height = 160, 120
- fourcc = cv2.VideoWriter_fourcc(*"mp4v")
- out = cv2.VideoWriter(path, fourcc, fps, (width, height))
+ def __create_timeseries_data(self, idx, duration, speed_factor):
+ path = f"{self.path}/TIMESERIES/{idx}.npy"
+ data = np.random.rand(duration, 1)
+ np.save(path, data)
- num_frames = duration * fps
- ball_radius = 20
- center_x = width // 2
+ def __create_visual_data(self, idx, duration, fps, speed_factor, modality_type):
+ if modality_type == ModalityType.VIDEO:
+ ext = "mp4"
+ elif modality_type == ModalityType.IMAGE:
+ ext = "jpg"
+ else:
+ raise ValueError(f"Unsupported modality type: {modality_type}")
- amplitude = random.uniform(0.5, 1.5) * (height // 3)
+ path = f"{self.path}/{modality_type.name}/{idx}.{ext}"
- for i in range(num_frames):
- frame = np.ones((height, width, 3), dtype=np.uint8) * 255
- center_y = int(
- height // 2
- + amplitude * np.sin(speed_factor * 2 * np.pi * i / num_frames)
- )
- frame = cv2.circle(
- frame, (center_x, center_y), ball_radius, (0, 255, 0), -1
- )
- out.write(frame)
-
- out.release()
+ width, height = 160, 120
+ if modality_type == ModalityType.VIDEO:
+ fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+ out = cv2.VideoWriter(path, fourcc, fps, (width, height))
+ num_frames = duration * fps
+ ball_radius = 20
+ center_x = width // 2
+
+ amplitude = random.uniform(0.5, 1.5) * (height // 3)
+
+ for i in range(num_frames):
+ frame = np.ones((height, width, 3), dtype=np.uint8) * 255
+ center_y = int(
+ height // 2
+ + amplitude * np.sin(speed_factor * 2 * np.pi * i / num_frames)
+ )
+ frame = cv2.circle(
+ frame, (center_x, center_y), ball_radius, (0, 255, 0), -1
+ )
+ out.write(frame)
+ out.release()
+ elif modality_type == ModalityType.IMAGE:
+ out = cv2.imwrite(path, np.ones((height, width, 3), dtype=np.uint8) * 255)
+ else:
+ raise ValueError(f"Unsupported modality type: {modality_type}")
def __create_text_data(self, idx, speed_factor):
path = f"{self.path}/TEXT/{idx}.txt"
@@ -403,34 +547,16 @@ class TestSVM(Model):
super().__init__(name)
def fit(self, X, y, X_test, y_test):
- if X.ndim > 2:
- X = X.reshape(X.shape[0], -1)
- self.clf = svm.SVC(C=1, gamma="scale", kernel="rbf", verbose=False)
- self.clf = self.clf.fit(X, np.array(y))
- y_pred = self.clf.predict(X)
-
- return {
- "accuracy": classification_report(
- y, y_pred, output_dict=True, digits=3, zero_division=1
- )["accuracy"]
- }, 0
+ return {"accuracy": random.uniform(0.5, 1.0)}, 0
def test(self, test_X: np.ndarray, test_y: np.ndarray):
- if test_X.ndim > 2:
- test_X = test_X.reshape(test_X.shape[0], -1)
- y_pred = self.clf.predict(np.array(test_X)) # noqa]
-
- return {
- "accuracy": classification_report(
- np.array(test_y), y_pred, output_dict=True, digits=3, zero_division=1
- )["accuracy"]
- }, 0
+ return {"accuracy": random.uniform(0.5, 1.0)}, 0
class TestTask(Task):
def __init__(self, name, model_name, num_instances):
self.labels = ModalityRandomDataGenerator().create_balanced_labels(
- num_instances=10
+ num_instances=num_instances
)
split = train_test_split(
np.array(range(num_instances)),
... 1238 lines suppressed ...