From 9b6176d2d13dea4a9a308ec362595e67877c087c Mon Sep 17 00:00:00 2001 From: Patrick Robert Miles Date: Wed, 8 Apr 2026 15:09:54 -0700 Subject: [PATCH 1/2] Sync benchmark testing config --- ScaFFold/configs/benchmark_testing.yml | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/ScaFFold/configs/benchmark_testing.yml b/ScaFFold/configs/benchmark_testing.yml index 6b8c30a..956a9a0 100644 --- a/ScaFFold/configs/benchmark_testing.yml +++ b/ScaFFold/configs/benchmark_testing.yml @@ -1,23 +1,23 @@ # 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. 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. From 13b87e915c4ed9892fa0fa4359c423ab918abe15 Mon Sep 17 00:00:00 2001 From: Patrick Robert Miles Date: Wed, 8 Apr 2026 15:13:56 -0700 Subject: [PATCH 2/2] Improve dataloader throughput path --- ScaFFold/cli.py | 10 ++++ ScaFFold/configs/benchmark_default.yml | 2 + ScaFFold/configs/benchmark_testing.yml | 2 + ScaFFold/utils/config_utils.py | 6 ++ ScaFFold/utils/data_loading.py | 77 +++++++++++++++++++------- ScaFFold/utils/trainer.py | 12 +++- docs/perf/dataloader-throughput.md | 63 +++++++++++++++++++++ 7 files changed, 151 insertions(+), 21 deletions(-) create mode 100644 docs/perf/dataloader-throughput.md diff --git a/ScaFFold/cli.py b/ScaFFold/cli.py index 5b76c87..96c9128 100644 --- a/ScaFFold/cli.py +++ b/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, diff --git a/ScaFFold/configs/benchmark_default.yml b/ScaFFold/configs/benchmark_default.yml index 180a0dc..456eb0c 100644 --- a/ScaFFold/configs/benchmark_default.yml +++ b/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 diff --git a/ScaFFold/configs/benchmark_testing.yml b/ScaFFold/configs/benchmark_testing.yml index 956a9a0..66d89eb 100644 --- a/ScaFFold/configs/benchmark_testing.yml +++ b/ScaFFold/configs/benchmark_testing.yml @@ -8,6 +8,8 @@ problem_scale: 7 # 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, 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 diff --git a/ScaFFold/utils/config_utils.py b/ScaFFold/utils/config_utils.py index 640ad19..be6cd59 100644 --- a/ScaFFold/utils/config_utils.py +++ b/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"]) diff --git a/ScaFFold/utils/data_loading.py b/ScaFFold/utils/data_loading.py index 725854c..7f6ea78 100644 --- a/ScaFFold/utils/data_loading.py +++ b/ScaFFold/utils/data_loading.py @@ -32,15 +32,30 @@ 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}" @@ -48,14 +63,44 @@ def __init__( 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(), diff --git a/ScaFFold/utils/trainer.py b/ScaFFold/utils/trainer.py index 80a4d4e..3d8b628 100644 --- a/ScaFFold/utils/trainer.py +++ b/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 diff --git a/docs/perf/dataloader-throughput.md b/docs/perf/dataloader-throughput.md new file mode 100644 index 0000000..4b13c06 --- /dev/null +++ b/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.