add updated memory estimators

Author: christinadionysio

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

@@ -36,6 +36,9 @@ class ImageStats:
     max_channels: int
     num_instances: int
     output_shape: tuple
+    average_width: int
+    average_height: int
+    average_channels: int
 
 
 class ImageLoader(BaseLoader):
@@ -79,7 +82,9 @@ def get_stats(self, source_path: str):
         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:
@@ -98,10 +103,19 @@ def get_stats(self, source_path: str):
             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,
-            (max_width, max_height, max_channels),
+            (average_width, average_height, average_channels),
+            average_width,
+            average_height,
+            average_channels,
         )

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

@@ -33,7 +33,9 @@ class JSONStats:
     num_instances: int
     max_length: int
     avg_length: float
-    output_shape: Tuple[int, int]
+    max_words: int
+    avg_words: float
+    output_shape: Tuple[int]
 
 
 class JSONLoader(BaseLoader):
@@ -74,6 +76,8 @@ def get_stats(self, source_path: str):
         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)
@@ -96,10 +100,20 @@ def get_stats(self, source_path: str):
                     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
-        return JSONStats(num_instances, max_length, avg_length, (max_length,))
+            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

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

@@ -31,6 +31,8 @@ class TextStats:
     num_instances: int
     max_length: int
     avg_length: float
+    max_words: int
+    avg_words: float
     output_shape: tuple
 
 
@@ -63,6 +65,8 @@ 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:
@@ -74,5 +78,10 @@ def get_stats(self, source_path: str):
                     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
-        return TextStats(num_instances, max_length, avg_length, (max_length,))
+            avg_words /= num_instances
+        return TextStats(
+            num_instances, max_length, avg_length, max_words, avg_words, (max_length,)
+        )