Custom Datasets, Dataloaders, and Distributed Sampling

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
Process 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 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)))

Get Started

Core Concepts

Building Models

Advanced Topics

Reference

Custom Datasets, Dataloaders, and Distributed Sampling

The IDataset Protocol

Creating a Dataset with `data/dataset`

Defining Datasets with `data/defdataset`

Tensor Datasets

Dataloaders

Dataloader Options

Worker Backends

Distributed Sampling

Creating a DistributedSampler

Sampler Options

`data/set-epoch!`

`data/sample-indices`

Integrating with a Dataloader

Collation

Resource Cleanup

Build docs developers (and LLMs) love

Get Started

Core Concepts

Building Models

Advanced Topics

Reference

Documentation Index

​The IDataset Protocol

​Creating a Dataset with data/dataset

​Defining Datasets with data/defdataset

​Tensor Datasets

​Dataloaders

​Dataloader Options

​Worker Backends

​Distributed Sampling

​Creating a DistributedSampler

​Sampler Options

​data/set-epoch!

​data/sample-indices

​Integrating with a Dataloader

​Collation

​Resource Cleanup

Build docs developers (and LLMs) love

The IDataset Protocol

Creating a Dataset with `data/dataset`

Defining Datasets with `data/defdataset`

Tensor Datasets

Dataloaders

Dataloader Options

Worker Backends

Distributed Sampling

Creating a DistributedSampler

Sampler Options

`data/set-epoch!`

`data/sample-indices`

Integrating with a Dataloader

Collation

Resource Cleanup