Dataset 与 DataLoader

主线位置：训练主线的入口，forward 之前：每一步的 batch 从哪来。主线 MLP 训练一步用的 4 个样本，在这里是数据集的第一个 batch。

版本：PyTorch 2.15.0（CPU）、burn 0.22.1（Flex 设备）

flowchart TB
  subgraph torch
    tg["Dataset getitem"] --> ts["sampler 给下标"]
    ts --> tc["default_collate"]
    tc --> tb["CPU 上的 batch"]
    tb --> tt["循环里 .to(device)"]
  end
  subgraph burn
    bg["Dataset get"] --> bl["DataLoader 给下标"]
    bl --> bb["Batcher batch"]
    bb --> bd["建在 set_device 的设备上"]
  end

两边的分工几乎一一对应：

PyTorch	burn 0.22.1
`Dataset.__getitem__(k)` / `__len__`	`Dataset::get(index) -> Result<I, E>` / `len()`
`default_collate`（不写就有）	`Batcher::batch(items, device)`（必须自己写）
`DataLoader(batch_size=, shuffle=, generator=, num_workers=)`	`DataLoaderBuilder.batch_size().shuffle(seed).num_workers().set_device()`
`drop_last=True`	没有
`for batch in loader` 得到 batch	`for batch in loader.iter()` 得到 `Result<Batch, DatasetError>`
worker 是进程	worker 是线程

实验代码：PyTorch 侧 exec_data.py；burn 侧 crate exec-data，feature 是本卷模板的 std、flex、autodiff、optim 再加 dataset（burn::data::dataloader 和 burn::data::dataset 都要这个 feature）。release 构建，平台 M2 Max、macOS 27、rustc 1.99.0。burn 内部的 DataLoader 结构见 burn 卷的 Learner 与数据「DataLoader」一节，本页只讲从 PyTorch 过来的差别。

同一份数据，两边各取一遍

10 个样本，第 k 个的输入是主线 MLP 那个公式的第 k 行，标签固定：

class MlpData(Dataset):
  """Sample k: the A1 input row k and its label."""

  def __len__(self) -> int:
    return len(LABELS)

  def __getitem__(self, k: int):
    x = torch.tensor([0.1 * (k * IN + c) - 0.5 for c in range(IN)])
    return k, x, LABELS[k]

loader = DataLoader(MlpData(), batch_size=4)

burn 里样本是普通的 Rust 值，不是张量；把一批样本拼成张量是 Batcher 的事，它拿到加载器的设备，张量直接建在那上面：

#[derive(Clone, Debug)]
struct Sample {
  id: usize,
  x: [f32; IN],
  y: i64,
}

fn sample(k: usize) -> Sample {
  // Same formula as the Python side: computed in f64, rounded to f32 once.
  let x = std::array::from_fn(|c| (0.1 * (k * IN + c) as f64 - 0.5) as f32);
  Sample { id: k, x, y: LABELS[k] }
}

impl Batcher<Sample, Batch> for MlpBatcher {
  fn batch(&self, items: Vec<Sample>, device: &Device) -> Batch {
    let n = items.len();
    let xs: Vec<f32> = items.iter().flat_map(|s| s.x).collect();
    let ys: Vec<i64> = items.iter().map(|s| s.y).collect();
    Batch {
      ids: items.iter().map(|s| s.id).collect(),
      x: Tensor::<2>::from_data(TensorData::new(xs, [n, IN]), device),
      y: Tensor::<1, Int>::from_data(TensorData::new(ys, [n]), device),
    }
  }
}

let loader = DataLoaderBuilder::new(MlpBatcher)
  .batch_size(4)
  .set_device(Device::flex())
  .build(InMemDataset::new(samples()));

x shape [4, 4], dtype torch.float32
x = [-0.5, -0.4000000059604645, -0.30000001192092896, -0.20000000298023224, -0.10000000149011612, 0.0, 0.10000000149011612, 0.20000000298023224, 0.30000001192092896, 0.4000000059604645, 0.5, 0.6000000238418579, 0.699999988079071, 0.800000011920929, 0.8999999761581421, 1.0]
y shape [4], dtype torch.int64, y = [0, 2, 1, 2]
ids [0, 1, 2, 3], then [([4, 5, 6, 7], [4, 4]), ([8, 9], [2, 4])]
len(loader) = 3, drop_last=True: [[0, 1, 2, 3], [4, 5, 6, 7]]

x shape [4, 4], dtype F32
x = [-0.5, -0.4, -0.3, -0.2, -0.1, 0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]
y shape [4], dtype I32, y = [0, 2, 1, 2]
ids [0, 1, 2, 3], then [([4, 5, 6, 7], [4, 4]), ([8, 9], [2, 4])]
num_items = 10

两边把一个 epoch 的所有 batch 写成 JSON，用 compare.py 逐个数比较，容差给 0：

ids                      n= 10  max|diff|=0.00e+00
x                        n= 40  max|diff|=0.00e+00
y                        n= 10  max|diff|=0.00e+00
worst=0.00e+00 tol=0e+00 OK

第一个 batch 就是主线 MLP 的 x 和 y。打印的样子不同（PyTorch 的 tolist() 打出 f32 对应的 double，burn 打出最短表示），值逐位相同。
标签在 PyTorch 是 int64，在 Flex 上是 I32：Batcher 里给的是 i64，from_data 按设备的默认 Int 类型存（见 dtype 与设备）。
len(loader) 是批数 3，burn 的 num_items() 是样本数 10（文档：「The number of items (not the number of batches nor the number of iterations)」）。
公式在哪种精度下算会影响逐位比较。Rust 里如果写成 0.1 * (k * IN + c) as f32 - 0.5，全程 f32，和 Python 先用 double 算再转 f32 的结果在 40 个数里差 1 到 2 个 ulp，compare.py 报 x max|diff|=2.38e-07、MISMATCH，退出码 1。这也是容差 0 的这道检查确实会变红的证据。

PyTorch 的 default_collate 在 CPU 上把样本叠起来，训练循环里再 x.to(device)；burn 没有自动 collate，Batcher 必须写，好处是 batch 直接建在加载器的设备上，循环里不用再搬。设备不设时是 Device::default()，交给 Learner 的加载器会被它改成训练设备（Learner 与数据）。

最后一批不满

两边默认都保留不满的最后一批（上面的 [8, 9]）。PyTorch 有 drop_last=True；burn 0.22.1 的 DataLoaderBuilder 没有这个选项，batch_size(n) 用的 FixBatchStrategy 在数据取完时强制出最后一批（FixBatchStrategy::batch，force 为真时不检查数量）：

    fn batch(&mut self, force: bool) -> Option<Vec<I>> {
        if self.items.len() < self.batch_size && !force {
            return None;
        }
        // ...

替代写法有两种，打乱时只有第二种和 drop_last 等价：

let whole = ds.len() - ds.len() % 4;
let ds = PartialDataset::new(ds, 0, whole);

PartialDataset 0..8: [[0, 1, 2, 3], [4, 5, 6, 7]]
PartialDataset 0..8, shuffle(7): [[5, 2, 7, 3], [1, 6, 4, 0]]

drop_last=True, shuffle, epoch 0: [[1, 3, 5, 7], [9, 4, 6, 2]]
drop_last=True, shuffle, epoch 1: [[5, 1, 6, 0], [9, 7, 8, 4]]
drop_last=True, shuffle, epoch 2: [[4, 3, 2, 8], [1, 0, 6, 5]]

PartialDataset 是在打乱之前截掉的，样本 8、9 永远不会被用到；PyTorch 的 drop_last 是打乱之后丢掉末尾，每个 epoch 丢的样本不同。要等价，就在训练循环里跳过不满的那一批（if batch.x.dims()[0] < 4 { continue; }）。

打乱：种子不通用

g = torch.Generator().manual_seed(7)
loader = DataLoader(MlpData(), batch_size=4, shuffle=True, generator=g)

DataLoaderBuilder::new(MlpBatcher).batch_size(4).shuffle(7)

shuffle, generator seed 7, epoch 0: [[1, 3, 5, 7], [9, 4, 6, 2], [8, 0]]
shuffle, generator seed 7, epoch 1: [[5, 1, 6, 0], [9, 7, 8, 4], [3, 2]]
shuffle, generator seed 7, epoch 2: [[4, 3, 2, 8], [1, 0, 6, 5], [9, 7]]

shuffle(7), epoch 0: [[5, 2, 7, 3], [1, 6, 8, 9], [4, 0]]
shuffle(7), epoch 1: [[1, 3, 7, 4], [5, 0, 2, 9], [8, 6]]
shuffle(7), epoch 2: [[5, 8, 6, 9], [1, 2, 3, 7], [4, 0]]

行为一样：每个 epoch 重新打乱，同一个种子整个过程可复现。顺序不一样：两边的随机数发生器和打乱算法都不同，没有哪个 burn 种子能得到 PyTorch 的顺序。连 PyTorch 自己的 torch.randperm(10, generator=torch.Generator().manual_seed(7)) 都是 [5, 0, 3, 4, 1, 7, 9, 6, 8, 2]，和加载器第一个 epoch 的顺序也不同，所以要对齐只能从 PyTorch 的加载器里把顺序取出来，交给 burn 的 SelectionDataset::from_indices_checked：

let idx: Vec<usize> = perm.split(',').map(|s| s.parse().unwrap()).collect();
let ds = SelectionDataset::from_indices_checked(InMemDataset::new(samples()), idx);
let l = DataLoaderBuilder::new(MlpBatcher).batch_size(4).set_device(Device::flex()).build(ds);

torch's permutation via SelectionDataset: [[1, 3, 5, 7], [9, 4, 6, 2], [8, 0]]

这个 epoch 的全部 batch 写成 JSON 再比，ids、x、y 三个字段的最大差都是 0；把 PyTorch 不打乱的结果拿来和它比，compare.py 报 worst=9.00e+00、MISMATCH。两边逐步对比训练（比如训练到收敛那样）时，用这个办法让两边看到同样的 batch 顺序。

burn 这一侧还有一个 PyTorch 没有的规则：对加载器做 to_device 或 slice 会派生一个新的随机数发生器，同一个种子的顺序就变了，Learner 正好会做这一步。细节见 burn 卷的同一个种子，顺序不一样。

多 worker：进程和线程

num_workers=2，同一个程序里新建 3 个加载器，每个跑 3 个 epoch，整个程序跑 3 遍。

PyTorch 的 3 个加载器、3 遍程序，每个 epoch 的批顺序都和单进程时一样：

2 workers, shuffle False, loader 0: [[[0, 1, 2, 3], [4, 5, 6, 7], [8, 9]], [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9]], [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9]]]
...
2 workers, shuffle True, loader 0: [[[1, 3, 5, 7], [9, 4, 6, 2], [8, 0]], [[5, 1, 6, 0], [9, 7, 8, 4], [3, 2]], [[4, 3, 2, 8], [1, 0, 6, 5], [9, 7]]]

DataLoader 的签名里 in_order: 'bool' = True：worker 先做完的批会被扣住，按下标顺序交出来。burn 的 worker 线程谁先做完谁先送（见 MultiThreadDataLoader 的 next），同一个程序第一遍的输出：

2 workers, shuffle false, loader 0: [[[8, 9], [0, 1, 2, 3], [4, 5, 6, 7]], [[0, 1, 2, 3], [8, 9], [4, 5, 6, 7]], [[0, 1, 2, 3], [8, 9], [4, 5, 6, 7]]]
2 workers, shuffle false, loader 1: [[[0, 1, 2, 3], [8, 9], [4, 5, 6, 7]], [[0, 1, 2, 3], [8, 9], [4, 5, 6, 7]], [[0, 1, 2, 3], [8, 9], [4, 5, 6, 7]]]
...
2 workers, shuffle true, loader 0: [[[6, 1, 0, 7], [9, 8, 2, 5], [3, 4]], [[0, 7, 2, 5], [1, 8, 9, 6], [3, 4]], [[9, 0, 2, 6], [7, 1, 5, 8], [4, 3]]]
2 workers, shuffle true, loader 1: [[[3, 4], [6, 1, 0, 7], [9, 8, 2, 5]], [[3, 4], [0, 7, 2, 5], [1, 8, 9, 6]], [[9, 0, 2, 6], [7, 1, 5, 8], [4, 3]]]
...

不打乱时，[8, 9] 有时第一个到、有时第二个到。打乱时还有一点更要紧：initialize 只在第一次 iter() 时用种子把整个数据集打乱一次，再按批切给各个 worker（PartialDataset::split_chunks：3 批分给 2 个 worker，第一个拿 2 批 8 个样本，第二个拿 1 批 2 个样本）；之后每个 epoch，每个 worker 只在自己那一段里重新打乱。

                let datasets = match self.strategy.batch_size() {
                    Some(batch_size) => {
                        PartialDataset::split_chunks(dataset, self.num_threads, batch_size)
                    }
                    None => PartialDataset::split(dataset, self.num_threads),
                };

❗IMPORTANT

动脑筋：2 个 worker、shuffle(7)，同一个加载器跑 3 个 epoch。PyTorch 的 3 个 epoch 里，只有 2 个样本的那一批分别是 [8, 0]、[3, 2]、[9, 7]。burn 的呢？数据、批大小同上，release 构建，平台同上。答案在本页末尾。

启动开销也不一样。从建加载器到拿到第一批，再到这个 epoch 结束（PyTorch 分成拿到最后一批和 StopIteration 两个时间点，负载约 3）：

0 workers: first batch after 0.496ms, last batch after 0.559ms, StopIteration after 0.570ms
2 workers: first batch after 670.598ms, last batch after 670.802ms, StopIteration after 10672.398ms

0 workers: first batch after 31.750µs, epoch done after 34.959µs
2 workers: first batch after 101.375µs, epoch done after 141.458µs

PyTorch 的 worker 是进程。这台机器上默认的启动方式是 spawn（multiprocessing.get_start_method() 的输出），每个 worker 重新导入 torch 和主模块，第一批要等 0.67 秒；数据集对象要能被 pickle，定义数据集的模块要有 if __name__ == "__main__": 保护。
这个 epoch 的最后一批 0.67 秒就到了，StopIteration 却在 10.7 秒才出来。换 1 个和 3 个 worker 实测是 5.6 秒和 15.7 秒，每个 worker 正好 5 秒，和关闭 worker 时 w.join(timeout=_utils.MP_STATUS_CHECK_INTERVAL)（这个值是 5.0）对得上：这台机器上 worker 进程没有在超时内退出。persistent_workers=True 时 epoch 之间不再关 worker，这 10 秒挪到了加载器被删除的时候（del loader: 10.006s）。worker 为什么不按时退出，本页没有追。
burn 的 worker 是线程，第一批 0.1 毫秒。Dataset 和 Batcher 要满足 Send + Sync（trait 定义里写着），编译器替你检查，不需要 pickle。线程只在第一次 iter() 时建一次，之后每个 epoch 复用（来历见 Learner 与数据的冷知识）。

取样本出错

数据集的第 5 个样本读不出来：

class Flaky(MlpData):
  def __getitem__(self, k: int):
    if k == 5:
      raise OSError(f"cannot read sample {k}")
    return super().__getitem__(k)

impl Dataset<Sample> for Flaky {
  fn get(&self, index: usize) -> Result<Sample, DatasetError> {
    if index == 5 {
      return Err(DatasetError::new(std::io::Error::other(format!("cannot read sample {index}"))));
    }
    Ok(sample(index))
  }
  fn len(&self) -> usize {
    LABELS.len()
  }
}

PyTorch 用 next(it) 一批一批取，接住异常后接着取；burn 打印每个 Result，只取前 6 个：

0 workers, batch 0: ok [0, 1, 2, 3]
0 workers, batch 1: OSError: cannot read sample 5
0 workers, batch 2: ok [8, 9]
0 workers, batch 3: StopIteration
2 workers, batch 0: ok [0, 1, 2, 3]
2 workers, batch 1: OSError: Caught OSError in DataLoader worker process 1.
    Original Traceback (most recent call last):
...
    OSError: cannot read sample 5
2 workers, batch 2: ok [8, 9]
2 workers, batch 3: StopIteration

0 workers, batch 0: Ok [0, 1, 2, 3]
0 workers, batch 1: Err(cannot read sample 5)
0 workers, batch 2: Err(cannot read sample 5)
0 workers, batch 3: Err(cannot read sample 5)
0 workers, batch 4: Err(cannot read sample 5)
0 workers, batch 5: Err(cannot read sample 5)
2 workers, batch 0: Ok [0, 1, 2, 3]
2 workers, batch 1: Err(dataloader worker 0 failed: cannot read sample 5)
2 workers, batch 2: Ok [8, 9]

PyTorch 的异常在取那一批时抛出，多进程时包一层「Caught OSError in DataLoader worker process 1」并带上 worker 里的原始 traceback。接住之后还能继续，坏的那一批被跳过。
burn 的迭代器不抛异常，交出 Err(DatasetError)，循环里要自己 match。
单线程加载器交出 Err 后，再调 next 还是同一个 Err，永远不结束：BatchDataloaderIterator::next 在 get_many 出错时直接返回，current_index 没有前进。去掉 take(6) 跑这个程序，两分钟左右写出了 5 GB 以上的输出，被强行终止。按 PyTorch 的习惯写成「记一条日志然后 continue」，就是死循环；在 burn 里遇到 Err 要 break 或用 ? 返回。origin/main 上这段代码没有变。
多线程时出错的 worker 报一条 Err 就退出，它那一段剩下的批不再交出；别的 worker 照常送完，迭代器随后结束。

Dataset 的文档还规定了另一半：下标越界（index >= len()）是 panic，不是 Err，「Err(_) is reserved for genuine retrieval failures on an in-bounds index (e.g. an I/O or deserialization failure)」。

📝NOTE

冷知识：get 以前返回 Option

Dataset::get 返回 Result、加载器的迭代器交出 Result，来自 PR #5198（「chore: change loop logic for data iterators and batchers」，首次进入 v0.22.0-pre.1，正式版 0.22.0），提交说明里的一条是「feat: change dataset contract to add Result<> and error propagation」。之前 get 是 fn get(&self, index: usize) -> Option<I>，单线程迭代器的循环是 while let Some(item) = self.dataset.get(self.current_index)：哪个样本取不出来（get 返回 None），这个 epoch 就在那里悄悄结束，后面的样本全部跳过。这个 PR 把循环改成按 len 走，取不出来就交出 Err；给 DataLoaderIterator 加的文档写着「Yields Err instead of panicking when a dataset retrieval fails, so callers can decide how to handle it (e.g. abort the training run with a clear message)」。上面那个重复交出同一个 Err 的行为，就是这次改写留下的：交出 Err 的那条路径在推进下标之前返回。

动脑筋答案

都是 [3, 4]（顺序有时是 [4, 3]）。前面第一遍程序里 loader 0 的三个 epoch：

2 workers, shuffle true, loader 0: [[[6, 1, 0, 7], [9, 8, 2, 5], [3, 4]], [[0, 7, 2, 5], [1, 8, 9, 6], [3, 4]], [[9, 0, 2, 6], [7, 1, 5, 8], [4, 3]]]

第一次 iter() 时整体打乱一次，切成 8 个样本和 2 个样本两段；第二段的 2 个样本是 3 和 4，分给第二个 worker。
之后每个 epoch，每个 worker 只在自己的段内打乱：第一个 worker 的 8 个样本 {0, 1, 2, 5, 6, 7, 8, 9} 每个 epoch 换着组成两批，样本 3 和 4 永远在同一批里，从来不和别的样本同批。
3 遍程序、每遍 3 个加载器，9 个加载器的结果都是这样（输出在实验目录里）。PyTorch 每个 epoch 对全部 10 个样本重新排列，样本之间的组合每次都变。数据少、worker 多时，这个差别会让 burn 的「打乱」明显弱于 PyTorch；要每个 epoch 全局打乱，用 num_workers(0)，或者自己每个 epoch 用新的排列建 SelectionDataset。

小结

Dataset.__getitem__ 对 Dataset::get（返回 Result），default_collate 对必须自己写的 Batcher，DataLoader(...) 对 DataLoaderBuilder；同一份数据两边的 batch 逐位相同（compare.py，容差 0）。
两边默认保留不满的最后一批；burn 没有 drop_last，在循环里跳过才等价，PartialDataset 截断会让末尾样本永远用不上。
同一个种子两边顺序不同；要对齐，把 PyTorch 的顺序交给 SelectionDataset::from_indices_checked。
PyTorch 的 worker 是进程、按顺序交批；burn 的 worker 是线程、谁先做完谁先交，打乱只在第一次 iter() 时全局做一次，之后每个 worker 只在自己的段内打乱。
burn 的迭代器交出 Err 而不是抛异常；单线程加载器出错后会一直交出同一个 Err，循环里遇到 Err 要退出。