# Adapt a customized data loading and data sampling pipeline to hf datasets

**URL:** <https://discuss.huggingface.co/t/adapt-a-customized-data-loading-and-data-sampling-pipeline-to-hf-datasets/172895>\
**Category:** Beginners\
**Created:** [January 28, 2026, 9:31am UTC](https://discuss.huggingface.co/t/adapt-a-customized-data-loading-and-data-sampling-pipeline-to-hf-datasets/172895 "2026-01-28T09:31:59Z")\
**Posts on this page:** 1\
**Showing post:** 2

<div class="post-metadata">

**Author:** ![John6666](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/john6666/32/27664_2.png) [@John6666](https://discuss.huggingface.co/u/John6666)\
**Post date:** [January 29, 2026, 10:27am UTC](https://discuss.huggingface.co/t/adapt-a-customized-data-loading-and-data-sampling-pipeline-to-hf-datasets/172895/2 "2026-01-29T10:27:47Z")

</div>

It [seems like there might be a lot of tricky spots to achieve this](https://huggingface.co/datasets/John6666/forum3/blob/main/adapt_pipeline_to_hf_ds_1.md)…

* * *

## 1) Converting the corpus into 🤗 Datasets (recommended approach)

### What to store in Arrow vs keep external

A `datasets.Dataset` is an Arrow table, which gives fast reads and convenient column operations, but it doesn’t require you to store the heavy bytes inside Arrow. In fact, the “Image” feature typically expects **paths or embedded bytes** , which would either force you to (a) extract images to files or (b) duplicate LMDB bytes into Arrow. ([Hugging Face](https://huggingface.co/docs/datasets/en/about_dataset_features "Dataset features"))

**Recommended** : store **metadata only** in Arrow and keep **LMDB as the image store**. At training time, decode by `file_name` using a transform.

### Shape of the dataset(s)

You can represent your multi-source corpus in either of these two equivalent ways:

- **Option A (common)**: `DatasetDict` with one `Dataset` per source
- **Option B** : one merged `Dataset` with a `source` column

Option A is usually nicer when each source has different policies (sampling freq/caps/epoch splits), because you can keep per-source stats and debug easily.

### Flatten the JSON grouped by `"W_H"`

Your idea to de-group and create explicit columns is correct. Suggested columns (superset of what you proposed):

- `file_name` (LMDB key)
- `label`, `type`
- `width`, `height` (parsed from `"W_H"`)
- `resolution` (string, optional)
- `bucket_id` (int; exact `(W,H)` or binned)
- `is_doc` (per-source constant, but store it per-row so transforms can branch)
- `source` (string)

### Why not store `image` as a Datasets “Image” feature here?

HF’s `Image` feature decodes from **paths** or **stored bytes** , and `decode=False` can return `{path, bytes}` instead of `PIL.Image`. This is useful when your data is already file-based or already stored as bytes in Arrow, but it doesn’t directly match “bytes live in LMDB keyed by file\_name” without copying. ([Hugging Face](https://huggingface.co/docs/datasets/en/about_dataset_features "Dataset features"))

### Training-time decode/augment via `set_transform` / `with_transform`

HF explicitly supports on-the-fly transforms for images (`set_transform`) and recommends using `map()` for one-time preprocessing and `set_transform` for per-epoch augmentations. ([Hugging Face](https://huggingface.co/docs/datasets/v2.2.1/en/image_process "Process image data"))

So the “HF datasets” part becomes:

1. offline: build metadata datasets and `save_to_disk()`

2. training: `load_from_disk()`, then attach a transform that:

* * *

## 2) Can HF Datasets implement your sampling strategy?

### What HF can do directly (partial)

HF provides `interleave_datasets()` with:

- **probabilities** (stochastic mixing)
- **stopping strategies** like `first_exhausted` and oversampling modes like `all_exhausted` ([Hugging Face](https://huggingface.co/docs/datasets/en/process "Process"))

This is helpful for _rough_ mixing, but it does **not** naturally express:

- “exactly 10 epochs with explicit epoch semantics”
- “source A is exactly 1× per epoch and source B is exactly 5× per epoch”
- “bucket by resolution and form batches within buckets”
- “DDP-safe, deterministic, disjoint work per rank”

### What you should do for strict semantics (recommended)

Use HF Datasets for **indexable metadata storage** and PyTorch for **policy** :

- HF `Dataset`: metadata table, deterministic indexing

- Training transform: LMDB decode + augmentation (`set_transform` / `with_transform`) ([Hugging Face](https://huggingface.co/docs/datasets/v2.2.1/en/image_process "Process image data"))

- **Custom PyTorch BatchSampler** : the control plane for:

This division matches what HF is best at (dataset storage + column ops) and what PyTorch is best at (sampling/batching control).

* * *

## 3) DDP sharding: memory-mapped Arrow is not sharding

Arrow memory mapping improves I/O and avoids loading everything into RAM, but **it doesn’t prevent two ranks from reading the same indices**. You still need an explicit sharding mechanism (sampler or dataset sharding).

HF provides `Dataset.shard(num_shards, index)` as a deterministic way to split a dataset into N pieces. ([Hugging Face](https://huggingface.co/docs/datasets/v3.1.0/en/process "Process"))  
HF’s Trainer/Accelerate will do distributed sharding for you, but if you’re building a custom loop you must ensure your own per-rank split. ([Hugging Face Forums](https://discuss.huggingface.co/t/should-i-shard-dataset-in-distributed-training/12494 "Should I shard dataset in distributed training? - Datasets - Hugging Face Forums"))

Separately, with PyTorch’s `DistributedSampler`, you must call `sampler.set_epoch(epoch)` each epoch to get different shuffles across epochs; otherwise, you can repeat the same order. ([PyTorch Forums](https://discuss.pytorch.org/t/why-is-sampler-set-epoch-epoch-needed-for-distributedsampler/149672 "Why is 'sampler.set\_epoch(epoch)' needed for DistributedSampler? - distributed - PyTorch Forums"))

**Practical recommendation for your case** : partition at the **batch** level (global deterministic batch schedule, then rank takes every `world_size`-th batch). This makes “disjoint work + identical step counts” easy to reason about.

* * *

## 4) Resolution grouping (“bucketing”) with HF + PyTorch

HF itself doesn’t have a built-in “bucketed batch sampler” for your exact needs. The clean approach is:

1. store `bucket_id` as a metadata column
2. in your sampler, build `bucket_id -> indices` per source
3. shuffle within bucket per epoch
4. emit batches that are bucket-pure (or near-pure if you allow spillover)

HF docs explicitly note that if you want efficient batching and transforms, using a `BatchSampler` is the right tool on the PyTorch side. ([Hugging Face](https://huggingface.co/docs/datasets/en/use_with_pytorch "Use with PyTorch"))

* * *

## 5) Are predefined per-epoch split files necessary/common?

They’re not “standard” unless you’re doing something like:

- curriculum learning (difficulty schedule)
- strict replay / auditability requirements
- staged inclusion (e.g., progressively adding noisy sources)
- hard constraints from upstream labeling processes

If you don’t have one of those reasons, they mainly add complexity.

**Simpler alternative** : add a deterministic `split_id` per sample (e.g., hash of `file_name`) and define epoch subsets as rules like “epoch 0 uses split\_id in [0..k)”, epoch 1 uses [k..2k), etc. This avoids managing 10 external JSON lists and keeps the logic local and reproducible.

* * *

## Summary recommendation

- Convert each source to a **metadata-only** HF `Dataset` (`file_name,label,type,is_doc,width,height,bucket_id,source,...`) and `save_to_disk`.
- Use `set_transform` / `with_transform` to **decode LMDB + augment + image\_processor** on the fly. ([Hugging Face](https://huggingface.co/docs/datasets/v2.2.1/en/image_process "Process image data"))
- Implement strict mixing (1×/5×), epoch semantics, bucketing, and DDP disjointness in a **custom PyTorch BatchSampler** , not in `interleave_datasets`. ([Hugging Face](https://huggingface.co/docs/datasets/en/process "Process"))
- Do not rely on Arrow memory mapping for sharding; use `Dataset.shard`/`DistributedSampler`/batch-level partitioning. ([Hugging Face](https://huggingface.co/docs/datasets/v3.1.0/en/process "Process"))

---

_[View the full topic](https://discuss.huggingface.co/t/adapt-a-customized-data-loading-and-data-sampling-pipeline-to-hf-datasets/172895)._