08 · DataLoader：worker 手里的不是你那个 Dataset

import os
import torch
from torch.utils.data import Dataset, DataLoader


class Counting(Dataset):
    def __init__(self):
        self.n_read = 0                      # ⚠️ 想统计「读了多少条」

    def __len__(self):
        return 8

    def __getitem__(self, i):
        self.n_read += 1
        return torch.tensor([float(i), float(self.n_read), float(os.getpid())])


if __name__ == "__main__":                   # ⭐ Windows 上必须有这一行，原因见 Python 板块
    for nw in (0, 2):
        ds = Counting()
        rows = [b for b in DataLoader(ds, batch_size=4, num_workers=nw)]
        pids = sorted({int(v) for b in rows for v in b[:, 2].tolist()})
        print(f"num_workers={nw}: 主进程里 ds.n_read = {ds.n_read}"
              f"  各条记录看到的 n_read = {[int(v) for b in rows for v in b[:, 1].tolist()]}"
              f"  进程数 = {len(pids)}")

import random
import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader


class RngDS(Dataset):
    def __len__(self):
        return 8

    def __getitem__(self, i):
        return {"torch": torch.randint(0, 10000, (1,)).item(),
                "numpy": int(np.random.randint(0, 10000)),
                "random": random.randint(0, 9999)}


if __name__ == "__main__":
    torch.manual_seed(0); np.random.seed(0); random.seed(0)
    for nw in (0, 2):
        out = {k: [] for k in ("torch", "numpy", "random")}
        for b in DataLoader(RngDS(), batch_size=8, num_workers=nw):
            for k in out:
                out[k] += [int(v) for v in b[k]]
        print(f"num_workers={nw}")
        for k, v in out.items():
            print(f"   {k:7}: 去重后 {len(set(v))}/8")

import os
import torch
from torch.utils.data import Dataset, DataLoader, get_worker_info


class WhoAmI(Dataset):
    def __len__(self):
        return 4

    def __getitem__(self, i):
        w = get_worker_info()
        if w is None:
            return torch.tensor([-1, -1, os.getpid(), 0])
        return torch.tensor([w.id, w.num_workers, os.getpid(), w.seed % 100000])


if __name__ == "__main__":
    for b in DataLoader(WhoAmI(), batch_size=2, num_workers=2):
        for r in b.tolist():
            print(f"worker id={r[0]} / 共 {r[1]}  pid={r[2]}  seed%100000={r[3]}")

import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader, get_worker_info


class OwnRng(Dataset):
    def __init__(self):
        self.rng = np.random.default_rng(0)      # ⚠️ 发生器存进了 self

    def __len__(self):
        return 8

    def __getitem__(self, i):
        return int(self.rng.integers(0, 10000))


def fix_seed(worker_id):                          # ⭐ 修法
    info = get_worker_info()
    info.dataset.rng = np.random.default_rng(torch.initial_seed() % 2**32)


if __name__ == "__main__":
    for nw in (0, 2, 4):
        v = [int(x) for b in DataLoader(OwnRng(), batch_size=4, num_workers=nw) for x in b]
        print(f"num_workers={nw}: {v}  去重后 {len(set(v))}/8")

    v = [int(x) for b in DataLoader(OwnRng(), batch_size=4, num_workers=2,
                                    worker_init_fn=fix_seed) for x in b]
    print(f"加了 worker_init_fn: {v}  去重后 {len(set(v))}/8")

import torch
from torch.utils.data import IterableDataset, DataLoader, get_worker_info


class Stream(IterableDataset):
    def __iter__(self):
        for i in range(6):
            yield i


class StreamFixed(IterableDataset):
    def __iter__(self):
        info = get_worker_info()
        wid = 0 if info is None else info.id
        n = 1 if info is None else info.num_workers
        for i in range(6):
            if i % n == wid:          # ⭐ 自己按 worker 分片
                yield i


if __name__ == "__main__":
    for nw in (0, 2, 3):
        v = [int(x) for b in DataLoader(Stream(), batch_size=6, num_workers=nw) for x in b]
        print(f"num_workers={nw}: {v}  共 {len(v)} 条（数据集只有 6 条）")
    v = [int(x) for b in DataLoader(StreamFixed(), batch_size=6, num_workers=3) for x in b]
    print(f"自己分片后 num_workers=3: {sorted(v)}  共 {len(v)} 条")

import torch
from torch.utils.data.dataloader import default_collate

print("① 三个 (2,) 张量 →", default_collate([torch.zeros(2), torch.ones(2), torch.ones(2) * 2]).shape)
print("② python int  →", default_collate([1, 2, 3]), default_collate([1, 2, 3]).dtype)
print("③ python float→", default_collate([1.0, 2.0]).dtype)
print("④ dict        →", default_collate([{"x": torch.zeros(2), "y": 1},
                                          {"x": torch.ones(2), "y": 2}]))
print("⑤ 字符串      →", default_collate(["a", "b", "c"]))

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader


class FloatDS(Dataset):
    def __len__(self):
        return 4

    def __getitem__(self, i):
        return torch.randn(3), float(i)          # ⚠️ 标签是 python float


if __name__ == "__main__":
    x, y = next(iter(DataLoader(FloatDS(), batch_size=4)))
    print("特征 dtype:", x.dtype, " 标签 dtype:", y.dtype)
    print("相减之后:", (nn.Linear(3, 1)(x).squeeze() - y).dtype)

RuntimeError: stack expects each tensor to be equal size, but got [3] at entry 0 and [5] at entry 1

import torch
from torch.nn.utils.rnn import pad_sequence


def pad_collate(batch):
    lens = torch.tensor([len(s) for s in batch])
    return pad_sequence(batch, batch_first=True), lens


x, lens = pad_collate([torch.ones(3), torch.ones(5) * 2])
print(x.shape, x.tolist(), lens.tolist())

TypeError: default_collate: batch must contain tensors, numpy arrays, numbers, dicts or lists; found <class '__main__.Foo'>