> ## Documentation Index
> Fetch the complete documentation index at: https://antlobach-clorch-182a83cb.mintlify.site/llms.txt
> Use this file to discover all available pages before exploring further.

# 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.

```clojure theme={null}
(require '[clorch.data :as data])
```

***

## The IDataset Protocol

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

```clojure theme={null}
(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:

```clojure theme={null}
(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`:

```clojure theme={null}
(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.

```clojure theme={null}
(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:

```clojure theme={null}
(def ds (data/tensor-dataset X-matrix y-vector))
```

***

## Dataloaders

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

```clojure theme={null}
(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 |

<Note>
  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.
</Note>

### Worker Backends

<Tabs>
  <Tab title="Thread workers">
    Thread workers share the same JVM heap and can access in-memory datasets directly. They are simpler but do not isolate failures.

    ```clojure theme={null}
    (data/dataloader dataset
      :num-workers    4
      :worker-backend :thread)
    ```
  </Tab>

  <Tab title="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.

    ```clojure theme={null}
    (data/dataloader dataset
      :num-workers    4
      :worker-backend :process)
    ```

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

    ```clojure theme={null}
    ;; In your dataset namespace
    (defn make-dataset [path]
      (data/dataset
       :size    (fn [] (count-files path))
       :get-item (fn [idx] (load-from-disk path idx))))
    ```
  </Tab>
</Tabs>

***

## 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

```clojure theme={null}
(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.

```clojure theme={null}
(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)))
```

<Warning>
  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.
</Warning>

### `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:

```clojure theme={null}
(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.

```clojure theme={null}
(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:

```clojure theme={null}
(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:

```clojure theme={null}
(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)))
```


This documentation is built and hosted on [Mintlify](https://mintlify.com), a developer documentation platform.