Perf/dataloader throughput

ScaFFold/cli.py

@@ -140,6 +140,16 @@ def main():
     benchmark_parser.add_argument(
         "--batch-size", type=int, nargs="+", help="Batch sizes for each volume size."
     )
+    benchmark_parser.add_argument(
+        "--dataloader-num-workers",
+        type=int,
+        help="Number of DataLoader workers per rank.",
+    )
+    benchmark_parser.add_argument(
+        "--dataloader-prefetch-factor",
+        type=int,
+        help="Number of prefetched batches per worker when num_workers > 0.",
+    )
     benchmark_parser.add_argument(
         "--warmup-batches",
         type=int,

ScaFFold/configs/benchmark_default.yml

@@ -8,6 +8,8 @@ problem_scale: 8                   # Determines dataset resolution and number of
 unet_bottleneck_dim: 3             # Power of 2 of the unet bottleneck layer dimension. Default of 3 -> bottleneck layer of size 8.
 seed: 42                           # Random seed.
 batch_size: 1                      # Batch sizes for each vol size.
+dataloader_num_workers: 4          # Number of DataLoader workers per rank.
+dataloader_prefetch_factor: 2      # Prefetched batches per worker when num_workers > 0.
 optimizer: "ADAM"                  # "ADAM" is preferred option, otherwise training defautls to RMSProp.
 dc_num_shards: [1, 1, 2]           # DistConv param: number of shards to divide the tensor into. It's best to choose the fewest ranks needed to fit one sample in GPU memory, since that keeps communication at a minimum
 dc_shard_dims: [2, 3, 4]           # DistConv param: dimension on which to shard

ScaFFold/configs/benchmark_testing.yml

@@ -1,23 +1,25 @@
 # External/user-facing
 base_run_dir: "benchmark_runs"     # Subfolder of $(pwd) in which to run jobs.
-dataset_dir: "datasets"            # Directory in which to store and query for datasets.
+dataset_dir: "/p/lustre5/miles30/benchmark_datasets/"            # Directory in which to store and query for datasets.
 fract_base_dir: "fractals"         # Base directory for fractal IFS and instances.
 n_categories: 5                    # Number of fractal categories present in the dataset.
 n_instances_used_per_fractal: 145  # Number of unique instances to pull from each fractal class. There are 145 unique; exceeding this number will reuse some instances.
-problem_scale: 6                   # Determines dataset resolution and number of unet layers. Default is 6.
+problem_scale: 7                   # Determines dataset resolution and number of unet layers. Default is 6.
 unet_bottleneck_dim: 3             # Power of 2 of the unet bottleneck layer dimension. Default of 3 -> bottleneck layer of size 8.
 seed: 42                           # Random seed.
 batch_size: 1                      # Batch sizes for each vol size.
+dataloader_num_workers: 4          # Number of DataLoader workers per rank.
+dataloader_prefetch_factor: 2      # Prefetched batches per worker when num_workers > 0.
 optimizer: "ADAM"                  # "ADAM" is preferred option, otherwise training defautls to RMSProp.
-num_shards: [1, 1, 1]              # DistConv param: number of shards to divide the tensor into. It's best to choose the fewest ranks needed to fit one sample in GPU memory, since that keeps communication at a minimum
-shard_dim: [2, 3, 4]               # DistConv param: dimension on which to shard
+dc_num_shards: [1, 1, 1]              # DistConv param: number of shards to divide the tensor into. It's best to choose the fewest ranks needed to fit one sample in GPU memory, since that keeps communication at a minimum
+dc_shard_dims: [2, 3, 4]               # DistConv param: dimension on which to shard
 checkpoint_interval: 100           # Checkpoint every C epochs. More frequent checkpointing can be very expensive on slow filesystems.
 
 # Internal/dev use only
 variance_threshold: 0.15           # Variance threshold for valid fractals. Default is 0.15.
 n_fracts_per_vol: 3                # Number of fractals overlaid in each volume. Default is 3.
 val_split: 25                      # In percent.
-epochs: 10                         # Number of training epochs.
+epochs: 3                         # Number of training epochs.
 learning_rate: .0001               # Learning rate for training.
 disable_scheduler: 1               # If 1, disable scheduler during training to use constant LR.
 more_determinism: 0                # If 1, improve model training determinism.

ScaFFold/utils/config_utils.py

@@ -50,6 +50,12 @@ def __init__(self, config_dict):
         self.n_instances_used_per_fractal = config_dict["n_instances_used_per_fractal"]
         self.scale = 1
         self.batch_size = config_dict["batch_size"]
+        self.dataloader_num_workers = int(
+            config_dict.get("dataloader_num_workers", 4)
+        )
+        self.dataloader_prefetch_factor = int(
+            config_dict.get("dataloader_prefetch_factor", 2)
+        )
         self.epochs = config_dict["epochs"]
         self.optimizer = config_dict["optimizer"]
         self.disable_scheduler = bool(config_dict["disable_scheduler"])

ScaFFold/utils/data_loading.py

@@ -32,30 +32,75 @@ def __init__(
         self.mask_dir = Path(mask_dir)
         self.mask_suffix = mask_suffix
 
-        self.ids = [
-            splitext(file)[0]
+        image_files = {
+            splitext(file)[0]: self.images_dir / file
             for file in listdir(images_dir)
             if isfile(join(images_dir, file)) and not file.startswith(".")
-        ]
+        }
+        mask_files = {}
+        for file in listdir(mask_dir):
+            if not isfile(join(mask_dir, file)) or file.startswith("."):
+                continue
+            mask_stem = splitext(file)[0]
+            if not mask_stem.endswith(mask_suffix):
+                continue
+            sample_id = mask_stem[: -len(mask_suffix)]
+            mask_files[sample_id] = self.mask_dir / file
+
+        self.ids = sorted(set(image_files) & set(mask_files))
         if not self.ids:
             raise RuntimeError(
                 f"No input file found in {images_dir}, make sure you put your images there"
             )
+        self.sample_paths = [
+            (sample_id, image_files[sample_id], mask_files[sample_id])
+            for sample_id in self.ids
+        ]
 
         customlog(
             f"Creating dataset with {len(self.ids)} examples. Loading from {data_dir}"
         )
         with open(data_dir, "rb") as data_file:
             data = pickle.load(data_file)
         self.mask_values = data["mask_values"]
+        self.mask_value_lookup = self._build_mask_value_lookup(self.mask_values)
         customlog(f"Unique mask values: {self.mask_values}")
 
     def __len__(self):
         return len(self.ids)
 
     @staticmethod
-    def preprocess(mask_values, img, is_mask):
+    def _build_mask_value_lookup(mask_values):
+        normalized_mask_values = [int(v) for v in mask_values]
+        if normalized_mask_values == list(range(len(normalized_mask_values))):
+            return np.arange(len(normalized_mask_values), dtype=np.int16)
+
+        if min(normalized_mask_values, default=0) < 0:
+            return None
+
+        max_value = max(normalized_mask_values, default=0)
+        lookup = np.full(max_value + 1, -1, dtype=np.int16)
+        for idx, value in enumerate(normalized_mask_values):
+            lookup[value] = idx
+        return lookup
+
+    @staticmethod
+    def preprocess(mask_values, img, is_mask, mask_value_lookup=None):
         if is_mask:
+            if (
+                mask_value_lookup is not None
+                and img.ndim == 3
+                and np.issubdtype(img.dtype, np.integer)
+                and img.min() >= 0
+                and img.max() < len(mask_value_lookup)
+            ):
+                mapped_mask = mask_value_lookup[img]
+                if (mapped_mask < 0).any():
+                    raise ValueError(
+                        "Encountered a mask value that was not present in mask_values"
+                    )
+                return mapped_mask
+
             mask = np.zeros((img.shape[0], img.shape[1], img.shape[2]), dtype=np.short)
             for i, v in enumerate(mask_values):
                 if img.ndim == 3:
@@ -70,25 +115,19 @@ def preprocess(mask_values, img, is_mask):
             return img
 
     def __getitem__(self, idx):
-        name = self.ids[idx]
-        mask_file = list(self.mask_dir.glob(name + self.mask_suffix + ".*"))
-        img_file = list(self.images_dir.glob(name + ".*"))
-
-        assert len(img_file) == 1, (
-            f"Either no image or multiple images found for the ID {name}: {img_file}"
-        )
-        assert len(mask_file) == 1, (
-            f"Either no mask or multiple masks found for the ID {name}: {mask_file}"
-        )
-        with open(mask_file[0], "rb") as f:
+        _, img_path, mask_path = self.sample_paths[idx]
+        with open(mask_path, "rb") as f:
             mask = np.load(f)
-        f.close()
-        with open(img_file[0], "rb") as f:
+        with open(img_path, "rb") as f:
             img = np.load(f)
-        f.close()
 
         img = self.preprocess(self.mask_values, img, is_mask=False)
-        mask = self.preprocess(self.mask_values, mask, is_mask=True)
+        mask = self.preprocess(
+            self.mask_values,
+            mask,
+            is_mask=True,
+            mask_value_lookup=self.mask_value_lookup,
+        )
 
         return {
             "image": torch.as_tensor(img.copy()).float().contiguous(),

ScaFFold/utils/trainer.py

@@ -134,11 +134,19 @@ def create_dataloaders(self):
         self.create_dataset()
         self.create_sampler()
 
+        num_workers = max(0, int(self.config.dataloader_num_workers))
         loader_args = dict(
-            batch_size=self.config.batch_size, num_workers=1, pin_memory=True
+            batch_size=self.config.batch_size,
+            num_workers=num_workers,
+            pin_memory=True,
+            persistent_workers=num_workers > 0,
         )
+        if num_workers > 0:
+            loader_args["prefetch_factor"] = max(
+                1, int(self.config.dataloader_prefetch_factor)
+            )
         self.log.debug(
-            f"dataloader num_workers={loader_args['num_workers']}, os.cpu_count()={os.cpu_count()}, self.world_size={self.world_size} "
+            f"dataloader num_workers={loader_args['num_workers']}, prefetch_factor={loader_args.get('prefetch_factor', 'n/a')}, persistent_workers={loader_args['persistent_workers']}, os.cpu_count()={os.cpu_count()}, self.world_size={self.world_size} "
         )
         self.train_loader = DataLoader(
             self.train_set, sampler=self.train_sampler, **loader_args

docs/perf/dataloader-throughput.md

@@ -0,0 +1,63 @@
+# DataLoader Throughput
+
+This branch improves the input pipeline relative to `miles30/performance`.
+
+## What Changed
+
+- Added configurable DataLoader worker count and prefetch factor.
+- Enabled persistent workers automatically when worker count is greater than zero.
+- Cached image and mask paths during dataset construction instead of doing per-sample `glob` lookups.
+- Added a fast integer lookup-table path for mask remapping when mask values are compatible with direct indexing.
+
+## Public Interface
+
+- New config key: `dataloader_num_workers`
+- New config key: `dataloader_prefetch_factor`
+- New CLI override: `--dataloader-num-workers`
+- New CLI override: `--dataloader-prefetch-factor`
+
+Branch defaults:
+
+- `dataloader_num_workers: 4`
+- `dataloader_prefetch_factor: 2`
+
+## What Did Not Change
+
+- Dataset file format is unchanged.
+- Tensor shapes, dtypes, and channel ordering are unchanged.
+- Loss math, optimizer behavior, and DistConv behavior are unchanged.
+
+## Expected Effect
+
+- Lower filesystem metadata overhead in `__getitem__`.
+- Better overlap between file I/O, CPU preprocessing, and GPU work.
+- Reduced label-remap cost for the common integer-mask case.
+
+## Evaluation
+
+```bash
+torchrun-hpc -N 4 -n 4 --gpus-per-proc 1 $(which scaffold) benchmark \
+  -c $(pwd)/ScaFFold/configs/benchmark_testing.yml \
+  --problem-scale 7 --dc-num-shards 1 2 2
+
+torchrun-hpc -N 8 -n 4 --gpus-per-proc 1 $(which scaffold) benchmark \
+  -c $(pwd)/ScaFFold/configs/benchmark_testing.yml \
+  --problem-scale 8 --dc-num-shards 1 2 2
+```
+
+To sweep worker settings:
+
+```bash
+torchrun-hpc -N 4 -n 4 --gpus-per-proc 1 $(which scaffold) benchmark \
+  -c $(pwd)/ScaFFold/configs/benchmark_testing.yml \
+  --problem-scale 7 --dc-num-shards 1 2 2 \
+  --dataloader-num-workers 2 \
+  --dataloader-prefetch-factor 4
+```
+
+## Acceptance Criteria
+
+- DataLoader construction reflects the configured worker and prefetch values.
+- No per-sample `glob` remains in `__getitem__`.
+- Mask remapping uses the lookup-table fast path for the normal integer-mask dataset.
+- Epoch `2+` timing improves or stays flat for the required comparison matrix.