Custom Datasets, Dataloaders, and Distributed Sampling

Create IDataset implementations, configure batching and shuffle, prefetch with thread or process workers, and shard data across distributed training ranks.

Clorch’s clorch.data namespace provides a composable data loading pipeline built around the IDataset protocol. A dataset knows its size and how to produce one item; a dataloader sequences items into batches, handles shuffling, and optionally parallelizes item loading across worker threads or subprocess workers. For distributed training, the DistributedSampler ensures each rank receives a disjoint, deterministic partition of the data.

(require '[clorch.data :as data])

---

The IDataset Protocol

Any Clojure value that implements IDataset can be passed to a dataloader.

(defprotocol IDataset
  (get-size [this]   "Returns the total number of items.")
  (get-item [this idx] "Returns {:data tensor :target tensor} for index idx."))

Creating a Dataset with data/dataset

The data/dataset function builds a minimal dataset from two keyword arguments:

(def my-dataset
  (data/dataset
   :size    (fn [] (count my-paths))
   :get-item (fn [idx] {:data   (load-tensor (nth my-paths idx))
                        :target (load-label  (nth my-paths idx))})))

For process-worker support, also provide :process-spec:

(def my-dataset
  (data/dataset
   :size         (fn [] (count my-paths))
   :get-item     (fn [idx] (load-item idx))
   :process-spec {:factory 'my.ns/make-dataset
                  :args    ["/path/to/data"]}))

Defining Datasets with data/defdataset

defdataset generates a named constructor and a record that implements IDataset. Fields in the binding vector become record slots, so the dataset state is immutable.

(data/defdataset TinyDataset [txt tokenizer max-length]
  [token-ids (vec (.toArray (.encode tokenizer txt)))
   n         (count token-ids)
   inputs    (mapv #(torch/tensor (subvec token-ids % (+ % max-length))
                                  {:dtype :int64})
                   (range 0 (- n max-length)))
   ml        max-length]
  (get-size  []    (count inputs))
  (get-item  [idx] {:data   (nth inputs idx)
                    :target (torch/tensor
                             (subvec token-ids (inc idx) (+ idx ml 1))
                             {:dtype :int64})}))

;; Instantiate
(def dataset (TinyDataset text tokenizer 64))
(data/get-size dataset)   ;; → number of sequences

Tensor Datasets

For simple in-memory supervised learning, data/tensor-dataset creates a dataset from two tensors directly:

(def ds (data/tensor-dataset X-matrix y-vector))

---

Dataloaders

data/dataloader wraps a dataset and produces lazy sequences of batches.

(def loader
  (data/dataloader dataset
    :batch-size    32
    :shuffle?      true
    :drop-last?    false
    :num-workers   4
    :prefetch-factor 2
    :collate-fn    data/default-collate))

;; Iterate over batches
(doseq [{:keys [data target]} loader]
  (train-step model data target))

Dataloader Options

Option	Default	Description
:batch-size	32	Items per batch
:shuffle?	true	Randomly permute indices each iteration
:drop-last?	false	Drop the last partial batch
:num-workers	0	Number of parallel workers (0 = main thread)
:prefetch-factor	2	Batches to prefetch per worker
:collate-fn	data/default-collate	Function that merges a list of items into one batch
:worker-backend	:auto	:thread, :process, or :auto (picks :process when workers > 0 and :process-spec is present)
:sampler	nil	An ISampler instance; cannot be combined with :shuffle? true
:timeout-ms	nil	Worker response timeout in milliseconds

When :num-workers is zero, batches are built synchronously on the calling thread. This is fine for small datasets or when items are already tensor-backed in memory.

Worker Backends

Thread workers

Thread workers share the same JVM heap and can access in-memory datasets directly. They are simpler but do not isolate failures.

(data/dataloader dataset
  :num-workers    4
  :worker-backend :thread)

Process workers

Process workers launch separate JVM subprocesses for true isolation. The dataset must implement IProcessDataset and provide a :process-spec factory that can reconstruct it from EDN-serializable arguments.

(data/dataloader dataset
  :num-workers    4
  :worker-backend :process)

The worker spec factory is called once per subprocess at startup:

;; In your dataset namespace
(defn make-dataset [path]
  (data/dataset
   :size    (fn [] (count-files path))
   :get-item (fn [idx] (load-from-disk path idx))))

---

Distributed Sampling

In multi-rank training, every rank must receive a disjoint, deterministic subset of the dataset so the same item is never processed twice in the same epoch.

Creating a DistributedSampler

(def sampler
  (data/distributed-sampler
   dataset-size
   {:num-replicas (dist/world-size)  ;; defaults from WORLD_SIZE env var
    :rank         (dist/rank)        ;; defaults from RANK env var
    :seed         1337
    :shuffle?     true
    :drop-last?   false}))

distributed-sampler also accepts the dataset itself instead of a plain integer — it will call get-size to determine the count.

Sampler Options

Option	Default	Description
:num-replicas	From WORLD_SIZE	Total number of ranks
:rank	From RANK	This rank’s index
:seed	0	Base seed for the shuffled permutation
:shuffle?	true	Shuffle before partitioning
:drop-last?	false	Drop trailing items when dataset doesn’t divide evenly

data/set-epoch!

Call set-epoch! at the start of every training epoch to advance the shuffle seed. Each rank applies the same permutation independently, which guarantees the partition remains disjoint.

(dotimes [epoch epochs]
  ;; Advance the shuffle before computing indices
  (data/set-epoch! sampler epoch)

  (doseq [indices (partition-all batch-size (data/sample-indices sampler))]
    (train-batch! indices)))

If you omit set-epoch!, the sampler reuses the same shuffle every epoch. This causes every epoch to train on identical mini-batches, which can cause overfitting and poor convergence.

data/sample-indices

data/sample-indices returns a vector of integer indices for this rank and epoch. Pass it to partition-all to form batches:

(def indices (data/sample-indices sampler))
;; → [42 7 93 11 ...]  — this rank's share, shuffled

(doseq [batch-indices (partition-all batch-size indices)]
  (let [batch (map #(data/get-item dataset %) batch-indices)]
    (train! batch)))

Integrating with a Dataloader

Pass the sampler to data/dataloader directly. The loader reads indices from the sampler instead of shuffling internally.

(data/set-epoch! sampler epoch)

(def loader
  (data/dataloader dataset
    :sampler    sampler
    :batch-size 32
    :num-workers 4))

(doseq [batch loader]
  (train-step model batch))

---

Collation

The default collate function (data/default-collate) handles three cases:

• Tensors: stacks a list of tensors along a new batch dimension using torch/stack.
• Maps: recursively collates each key, producing a map of batched tensors.
• Other values: wraps them in a Clojure vector.

To override collation, provide a :collate-fn to the dataloader:

(defn my-collate [items]
  {:data   (torch/stack (mapv :data items) 0)
   :target (mapv :target items)})  ;; leave targets as a vector

(data/dataloader dataset :collate-fn my-collate)

---

Resource Cleanup

Tensors stacked by the default collate function are retained so they survive past the native-memory scope of the worker. If your training loop wraps each batch in t/with-torch, call data/cleanup-data! on the batch after you have extracted all JVM-scalar results:

(doseq [batch loader]
  (let [loss-value
        (t/with-torch
          (let [loss (train-step model batch)]
            (t/item-float loss)))]
    (data/cleanup-data! batch)
    (println "Loss:" loss-value)))