PyTorch中能否通过双DataLoader衔接模型预测实现训练加速？

PyTorch没有直接提供双DataLoader异步调度的专属模块，但可以通过多进程预加载+手动调度或异步DataLoader实现你要的“数据加载与模型计算重叠”效果，以下是具体实现方案：

方法1：基于标准DataLoader的手动预取 #

利用DataLoader的多进程加载能力（num_workers参数），手动提前预取目标数据，让模型计算当前批次的同时，后台进程加载下一批目标数据：

import torch
from torch.utils.data import DataLoader

# 初始化两个对齐的DataLoader（确保批次顺序一致）
input_loader = DataLoader(input_dataset, batch_size=32, num_workers=4, pin_memory=True)
target_loader = DataLoader(target_dataset, batch_size=32, num_workers=4, pin_memory=True)

input_iter = iter(input_loader)
target_iter = iter(target_loader)

# 预取第一批目标数据，提前启动加载
next_target = next(target_iter)

for input_batch in input_loader:
    # 1. 模型计算（此时后台已在加载下一批目标数据）
    pred = model(input_batch.cuda(non_blocking=True))
    
    # 2. 用预取的目标数据计算损失
    loss = criterion(pred, next_target.cuda(non_blocking=True))
    loss.backward()
    optimizer.step()
    
    # 3. 预取下一批目标数据，和下一次模型计算并行
    try:
        next_target = next(target_iter)
    except StopIteration:
        # 数据集遍历完后重置迭代器
        target_iter = iter(target_loader)
        next_target = next(target_iter)

方法2：使用PyTorch 2.0+的AsyncDataLoader #

PyTorch 2.0推出的AsyncDataLoader是原生异步预取的DataLoader，结合torch.jit.fork/wait可以更精准地控制并行逻辑：

from torch.utils.data import AsyncDataLoader
import torch

# 初始化异步DataLoader
input_loader = AsyncDataLoader(input_dataset, batch_size=32, num_workers=4, pin_memory=True)
target_loader = AsyncDataLoader(target_dataset, batch_size=32, num_workers=4, pin_memory=True)

input_iter = iter(input_loader)
target_iter = iter(target_loader)

# 预取第一批输入和目标数据
next_input = next(input_iter)
next_target = next(target_iter)

for _ in range(len(input_loader)):
    # 异步启动模型计算
    pred_future = torch.jit.fork(model, next_input.cuda(non_blocking=True))
    
    # 同时加载下一批输入和目标数据
    try:
        next_input = next(input_iter)
        next_target = next(target_iter)
    except StopIteration:
        input_iter = iter(input_loader)
        target_iter = iter(target_loader)
        next_input = next(input_iter)
        next_target = next(target_iter)
    
    # 等待模型计算完成
    pred = torch.jit.wait(pred_future)
    
    # 计算损失与反向传播
    loss = criterion(pred, next_target.cuda(non_blocking=True))
    loss.backward()
    optimizer.step()

关键注意事项 #

• 批次对齐：如果开启shuffle=True，必须给两个DataLoader设置相同的随机种子和worker_init_fn，保证输入与目标数据的配对正确：

  def seed_worker(worker_id):
      worker_seed = torch.initial_seed() % 2**32
      numpy.random.seed(worker_seed)
      random.seed(worker_seed)
  
  g = torch.Generator()
  g.manual_seed(42)
  
  input_loader = DataLoader(..., shuffle=True, worker_init_fn=seed_worker, generator=g)
  target_loader = DataLoader(..., shuffle=True, worker_init_fn=seed_worker, generator=g)
  

• 资源控制：num_workers不宜设置过大，一般为CPU核心数的1-2倍，避免CPU资源耗尽拖慢整体速度。

内容的提问来源于stack exchange，提问作者XYZCODE123