训练每隔几步就“卡半天”？PyTorch DataLoader 阻塞与吞吐骤降的三连坑：num_workers/预取与持久进程、pin_memory/非阻塞搬运

# dataloader_stall_repro.py
import os, time, argparse, random, numpy as np
import torch
from torch.utils.data import Dataset, DataLoader

SEED = 0
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)

class HeavyDataset(Dataset):
    """
    模拟重 CPU 变换与不规则 I/O：sleep + 大数组拷贝 + 小概率慢样本
    """
    def __init__(self, n=20000, slow_p=0.01, size=(3, 256, 256)):
        self.n, self.slow_p, self.size = n, slow_p, size
    def __len__(self): return self.n
    def __getitem__(self, idx):
        # 模拟解码/增强：20~40ms
        t0 = time.time()
        x = np.random.rand(*self.size).astype(np.float32)  # 生成/解码
        x = (x * 255).astype(np.uint8)                     # 模拟色彩/增强
        if random.random() < self.slow_p:
            time.sleep(0.15)                               # 极端慢样本
        # 转 Tensor
        x = torch.from_numpy(x).float() / 255.0
        y = torch.tensor(idx % 10, dtype=torch.long)
        return x, y, time.time() - t0

def run(args):
    if args.disable_opencv_threads:
        try:
            import cv2
            cv2.setNumThreads(0)   # 防止 OMP 过度并行
        except Exception:
            pass
    if args.set_torch_threads:
        torch.set_num_threads(args.set_torch_threads)

    ds = HeavyDataset(n=args.num_samples, slow_p=args.slow_p)
    loader = DataLoader(
        ds,
        batch_size=args.batch_size,
        shuffle=True,
        num_workers=args.num_workers,
        pin_memory=args.pin_memory,
        prefetch_factor=args.prefetch_factor if args.num_workers > 0 else None,
        persistent_workers=args.persistent_workers and args.num_workers > 0,
        drop_last=True,
    )

    device = torch.device("cuda" if torch.cuda.is_available() and not args.cpu_only else "cpu")
    model = torch.nn.Sequential(
        torch.nn.Conv2d(3, 16, 3, padding=1), torch.nn.ReLU(),
        torch.nn.Conv2d(16, 16, 3, padding=1), torch.nn.ReLU(),
        torch.nn.AdaptiveAvgPool2d(1), torch.nn.Flatten(), torch.nn.Linear(16, 10)
    ).to(device)

    def one_epoch(epoch):
        t0 = time.time()
        step_times, load_times = [], []
        for i, (x, y, load_t) in enumerate(loader):
            load_times.append(float(load_t.mean()))
            t1 = time.time()
            x = x.to(device, non_blocking=args.non_blocking and args.pin_memory)
            y = y.to(device, non_blocking=args.non_blocking and args.pin_memory)
            out = model(x); loss = torch.nn.functional.cross_entropy(out, y)
            loss.backward() if args.backward else None
            step_times.append(time.time() - t1)
            if (i + 1) % 50 == 0:
                print(f"[ep{epoch}] step {i+1:05d}  "
                      f"step_time={np.mean(step_times[-50:]):.4f}s  "
                      f"avg_load={np.mean(load_times[-50:]):.4f}s")
        dt = time.time() - t0
        print(f"[ep{epoch}] epoch_time={dt:.2f}s  avg_step_time={np.mean(step_times):.4f}s")
        return np.mean(step_times)

    # 热身一轮避免 JIT/CUDNN 影响
    one_epoch("warmup")
    times = [one_epoch(e+1) for e in range(2)]
    return times

def parse():
    ap = argparse.ArgumentParser()
    ap.add_argument("--num_workers", type=int, default=0)
    ap.add_argument("--prefetch_factor", type=int, default=2)
    ap.add_argument("--persistent_workers", action="store_true")
    ap.add_argument("--pin_memory", action="store_true")
    ap.add_argument("--non_blocking", action="store_true")
    ap.add_argument("--batch_size", type=int, default=64)
    ap.add_argument("--num_samples", type=int, default=5000)
    ap.add_argument("--slow_p", type=float, default=0.01)
    ap.add_argument("--backward", action="store_true")
    ap.add_argument("--cpu_only", action="store_true")
    ap.add_argument("--disable_opencv_threads", action="store_true")
    ap.add_argument("--set_torch_threads", type=int, default=0)
    return ap.parse_args()

if __name__ == "__main__":
    # Windows/macOS 必须守护 main，避免 spawn 死锁
    args = parse()
    run(args)

# A. 单进程（基线，通常卡 GPU）
python dataloader_stall_repro.py --num_workers 0 --batch_size 64

# B. 合理并行 + 预取 + 持久进程（训练多 epoch 时显著稳）
python dataloader_stall_repro.py --num_workers 4 --prefetch_factor 4 --persistent_workers \
    --pin_memory --non_blocking --batch_size 64

# C. 错误组合：persistent_workers=True 但 num_workers=0（无效且易误判）
python dataloader_stall_repro.py --num_workers 0 --persistent_workers --batch_size 64

# D. 线程过载：torch.set_num_threads 与 OpenCV 争抢（观察抖动/卡顿）
python dataloader_stall_repro.py --num_workers 8 --set_torch_threads 8 --disable_opencv_threads

# 1) DataLoader 基线（GPU 训练）
train_loader = DataLoader(
    dataset,
    batch_size=per_device_batch,
    shuffle=True,
    num_workers=min(8, os.cpu_count() // max(1, ngpu) // 2),  # 经验：每卡 2~4 个
    pin_memory=True,
    prefetch_factor=4,                 # 负载较重时适当增大
    persistent_workers=True,           # 训练多个 epoch 强烈建议开启
    drop_last=True
)

# 2) 训练循环中的非阻塞搬运
for x, y in train_loader:
    x = x.to(device, non_blocking=True)
    y = y.to(device, non_blocking=True)
    ...

# 3) 跨平台守护 + 线程抑制（放在入口处）
if __name__ == "__main__":
    import torch.multiprocessing as mp, cv2, os
    os.environ.setdefault("OMP_NUM_THREADS", "1")  # 先抑制 OMP
    try: cv2.setNumThreads(0)
    except Exception: pass
    if os.name != "posix":  # Windows/macOS
        mp.set_start_method("spawn", force=True)
    main()

def fast_collate(batch):
    # (示例) 批量堆叠而非逐个 append
    xs, ys = zip(*[(b[0], b[1]) for b in batch])
    return torch.stack(xs, 0), torch.tensor(ys, dtype=torch.long)

class LoaderHealth:
    def __init__(self): self.load_times, self.step_times = [], []
    def log(self, load_t, step_t):
        self.load_times.append(load_t); self.step_times.append(step_t)
        if len(self.step_times) % 100 == 0:
            lt, st = np.mean(self.load_times[-100:]), np.mean(self.step_times[-100:])
            ratio = lt / max(st, 1e-6)
            print(f"[health] load/step={ratio:.2f}  avg_load={lt:.3f}s  avg_step={st:.3f}s")
            # 若 load/step > 0.6，GPU 可能在等数据

def assert_persistent_valid(num_workers, persistent_workers):
    if persistent_workers and num_workers == 0:
        raise ValueError("persistent_workers=True 仅在 num_workers>0 时有效")

def set_data_threads(omp=1, torch_threads=None):
    os.environ.setdefault("OMP_NUM_THREADS", str(omp))
    if torch_threads: torch.set_num_threads(torch_threads)