Rust in Everything
基于 burn 0.22.1

同一个种子，顺序不一样

主线位置：训练主线的数据一侧。每个 step 的 batch 从 DataLoader 来，Learner 在第一个 epoch 之前对它做了一次 to_device。加载器的结构见 Learner 与数据。

flowchart TB
  seed["shuffle 7"] --> rng["StdRng"]
  rng -->|每个 epoch fork 一次| ep["ShuffledDataset"]
  rng -->|to_device 再 fork 一次| rng2["新加载器的 StdRng"]
  rng2 --> ep2["另一种顺序"]
  w["num_workers 2"] --> ch["同步通道 谁先做完谁先送"]

实验 crate pc-shuffle：一个 Batcher 直接返回样本编号，12 个样本、每批 4 个，shuffle(7)，不需要张量就能看到顺序。debug 构建，M2 Max、macOS 27、rustc 1.99.0。

现象

#[derive(Clone)]
struct Ids;

impl Batcher<u32, Vec<u32>> for Ids {
  fn batch(&self, items: Vec<u32>, _device: &Device) -> Vec<u32> {
    items
  }
}

fn loader(seed: u64, workers: usize) -> Arc<dyn DataLoader<Vec<u32>>> {
  DataLoaderBuilder::new(Ids)
    .batch_size(4)
    .shuffle(seed)
    .num_workers(workers)
    .set_device(Device::flex())
    .build(InMemDataset::new((0..12).collect::<Vec<u32>>()))
}

两个同样构造的加载器 a、b，再加一个构造完马上 to_device 的 c（单设备训练策略对训练加载器做的就是这一步）：

let a = loader(7, 0);
let b = loader(7, 0);
println!("loader a, epoch 1: {:?}", epoch(&a));
println!("loader b, epoch 1: {:?}", epoch(&b));
let c = loader(7, 0).to_device(&device);
println!("to_device, epoch 1: {:?}", epoch(&c));
println!("loader a, epoch 2: {:?}", epoch(&a));

loader a, epoch 1: [[11, 10, 7, 3], [1, 6, 8, 9], [4, 0, 2, 5]]
loader b, epoch 1: [[11, 10, 7, 3], [1, 6, 8, 9], [4, 0, 2, 5]]
to_device, epoch 1: [[8, 5, 11, 7], [1, 6, 0, 2], [10, 3, 4, 9]]
loader a, epoch 2: [[1, 3, 7, 4], [5, 0, 2, 10], [11, 6, 9, 8]]

同一个种子，to_device 之后的顺序和原来不一样。手写训练循环和 Learner 用同一个种子、同一份数据，loss 从第一步就对不上，原因就在这里（Learner 与数据用 loss 比对过）。

多 worker 的情况：同一个程序里新建 5 次 loader(7, 2)，各跑一个 epoch；整个程序再跑 3 遍，统计出现过的顺序：

2 workers, run 0: [[6, 1, 0, 7], [3, 9, 4, 10], [11, 8, 2, 5]]
2 workers, run 1: [[3, 9, 4, 10], [6, 1, 0, 7], [11, 8, 2, 5]]
2 workers, run 2: [[6, 1, 0, 7], [3, 9, 4, 10], [11, 8, 2, 5]]
2 workers, run 3: [[3, 9, 4, 10], [6, 1, 0, 7], [11, 8, 2, 5]]
2 workers, run 4: [[3, 9, 4, 10], [6, 1, 0, 7], [11, 8, 2, 5]]

   1 2 workers, run 0: [[3, 9, 4, 10], [6, 1, 0, 7], [11, 8, 2, 5]]
   2 2 workers, run 0: [[6, 1, 0, 7], [3, 9, 4, 10], [11, 8, 2, 5]]
...
   1 2 workers, run 3: [[6, 1, 0, 7], [11, 8, 2, 5], [3, 9, 4, 10]]
...
   2 2 workers, run 4: [[6, 1, 0, 7], [11, 8, 2, 5], [3, 9, 4, 10]]

每一批里有哪些样本是固定的，三批的先后有三种排法。

出处

单线程的 BatchDataLoader 每次 iter() 用自己的 StdRng 建一个 ShuffledDataset。传进去的是 &mut StdRng，From<&mut StdRng> for RngSource 做的是 rng.fork()：从父发生器取一段随机数当种子，派生一个新的。所以每个 epoch 从父发生器上「切」一次。

to_device 也是 fork：

    fn to_device(&self, device: &Device) -> Arc<dyn DataLoader<O>> {
        let rng = self.rng.as_ref().map(|rng| {
            let mut rng = rng.lock();
            rng.fork()
        });
        Arc::new(Self::new(

新加载器拿到的是「父发生器切下的第一段」，它的第一个 epoch 再从这一段上切，自然和原来的第一个 epoch 不同。fork 要 &mut，父发生器也往前走了一段：

a after a.to_device(), epoch 1: [[1, 3, 7, 4], [5, 0, 2, 10], [11, 6, 9, 8]]

a 被 to_device 过一次之后，它的第一个 epoch 正好是没被碰过时的第二个 epoch。slice 同理。Learner 的单设备策略在开始前对训练加载器 to_device(&self.device)，验证加载器也一样。

多线程的 MultiThreadDataLoader 是另一回事。它第一次 iter() 时用种子把数据集整体打乱一次，按批切给各个 worker（initialize），每个 worker 线程做完一批就往同一个通道里送；next 收到什么就返回什么，不按 worker 编号排序。每批的内容由种子决定，先后由线程调度决定。它的 to_device 直接复用种子，不 fork。

📝NOTE
冷知识：to_device 以前是复制状态

to_device 和 slice 里的 fork() 来自 PR #4500（「Upgrade to rand 0.10」，首次进入 v0.21.0-pre.2，正式版 0.21.0）。之前这两处是 rng.clone()：新加载器拿到和原来一模一样的发生器状态，打乱顺序和原加载器相同，原加载器也不往前走。提交说明只列了「Upgrade to rand 0.10」「Clippy」「Fix examples」这些条目，没解释这一处；能看到的原因在 rand 本身：rand 0.9.5 的 StdRng 是 #[derive(Clone, Debug, PartialEq, Eq)]，0.10.3 是 #[derive(Debug, PartialEq, Eq)]，不能再 clone。

修法

要和 Learner 逐位对得上，自己的循环也对加载器做一次 to_device。两个都 to_device 过的加载器顺序相同：

let x = loader(7, 0).to_device(&device);
let y = loader(7, 0).to_device(&device);

to_device x, epoch 1: [[8, 5, 11, 7], [1, 6, 0, 2], [10, 3, 4, 9]]
to_device y, epoch 1: [[8, 5, 11, 7], [1, 6, 0, 2], [10, 3, 4, 9]]

要求批的先后可复现时，不要开多个 worker。num_workers(1) 只有一个线程，顺序固定（和 num_workers(0) 的顺序不同，因为多线程加载器先用派生的种子整体打乱了一次）：

1 worker, run 0: [[4, 1, 0, 3], [9, 8, 2, 5], [6, 7, 10, 11]]
1 worker, run 1: [[4, 1, 0, 3], [9, 8, 2, 5], [6, 7, 10, 11]]
1 worker, run 2: [[4, 1, 0, 3], [9, 8, 2, 5], [6, 7, 10, 11]]

多 worker 时只能做到「每批内容可复现」，如果训练对批的先后敏感（比如要和另一次运行逐步比 loss），就把 worker 数降到 0 或 1 来比。

怎么确认没中招

  • 像本页这样写一个返回编号的 Batcher，在同一个进程里打印两个加载器的第一个 epoch，比较的两边要经过同样的 to_device / slice 调用。
  • 不要拿「同一个加载器的第二个 epoch」和「另一个加载器的第一个 epoch」比：每个 epoch 本来就重新打乱。
  • num_workers > 0 时，把整个程序跑几遍，看批的顺序是否变化。

小结

  • shuffle(seed) 保证的是：同样构造、同样经过 to_device / slice、同样迭代次数的加载器顺序相同。
  • to_device 和 slice 会 fork 随机数发生器，新旧两个加载器的顺序都会变；Learner 对训练加载器做了这一步。
  • 多 worker 时每批内容固定，批的先后看线程调度。

💬 关联讨论

发起讨论
加载中...