介绍
PyTorch中 加载批量训练/测试数据的核心工具，能高效处理数据迭代、批量划分、打乱、多线程加载等问题，是深度学习训练流程中不可或缺的组件（不管是CV、NLP还是大模型训练，都会用到）。
功能
DataLoader 封装了PyTorch的Dataset（数据集），提供以下核心能力：
- 按指定批次大小（
batch_size）切分数据； - 训练时打乱数据（
shuffle=True），避免过拟合； - 多线程/多进程加载数据（
num_workers），提升训练效率； - 自动拼接张量、处理不同长度的数据（如NLP中的padding）；
- 支持自定义数据处理逻辑（如采样器、整理函数）。
语法
from torch.utils.data import DataLoader
dataloader = DataLoader(
dataset, # 必须：封装好的数据集（torch.utils.data.Dataset实例）
batch_size=1, # 批次大小（默认1，常用32/64/128）
shuffle=False, # 是否打乱数据（训练集设True，测试集设False）
sampler=None, # 采样器（自定义数据采样方式，与shuffle互斥）
batch_sampler=None, # 批次采样器（自定义批次采样）
num_workers=0, # 加载数据的线程数（Windows建议设0，Linux/Mac可设CPU核心数）
collate_fn=None, # 整理函数（处理一个批次的数据，如NLP的padding）
pin_memory=False, # 是否将数据加载到锁页内存（GPU训练时设True，加速）
drop_last=False, # 是否丢弃最后一个不完整的批次（如总数据数不是batch_size整数倍）
timeout=0, # 数据加载超时时间（秒）
worker_init_fn=None, # 每个工作线程的初始化函数
multiprocessing_context=None # 多进程上下文（如"spawn"）
)
关键参数详解
| 参数 | 核心作用 & 注意事项 |
|---|---|
batch_size |
批次越大，训练越快，但显存占用越高（需根据GPU显存调整，如32/64/128）；测试集可设更大。 |
shuffle |
训练集必须设True（打乱数据，避免模型记住顺序），测试集必须设False（保证结果可复现）。 |
num_workers |
Windows系统建议设0（多线程易报错），Linux/Mac可设CPU核心数-1（如8核设7）；过大可能导致内存溢出。 |
collate_fn |
处理非标量/变长数据（如NLP padding、CV多模态数据），是自定义数据处理的核心。 |
drop_last |
训练时建议设True（避免最后一个小批次导致BatchNorm等层不稳定），测试时设False（不丢弃数据）。 |
pin_memory |
GPU训练时设True（将数据加载到锁页内存，减少CPU→GPU的数据拷贝时间），CPU训练设False。 |
例子
1. 基础流程：Dataset + DataLoader
先定义数据集（Dataset），再用DataLoader封装：
import torch
from torch.utils.data import Dataset, DataLoader
# 1. 定义自定义数据集（继承Dataset）
class MyDataset(Dataset):
def __init__(self, data):
self.data = data # 示例数据：[(输入1, 标签1), (输入2, 标签2), ...]
# 必须实现：返回数据集长度
def __len__(self):
return len(self.data)
# 必须实现：返回索引为idx的样本（输入+标签）
def __getitem__(self, idx):
x, y = self.data[idx]
return torch.tensor(x), torch.tensor(y)
# 2. 准备数据
train_data = [(1, 0), (2, 1), (3, 0), (4, 1), (5, 0), (6, 1)] # 输入+标签
test_data = [(7, 0), (8, 1)]
# 3. 创建Dataset实例
train_dataset = MyDataset(train_data)
test_dataset = MyDataset(test_data)
# 4. 创建DataLoader（核心）
train_loader = DataLoader(
dataset=train_dataset,
batch_size=2, # 每个批次2个样本
shuffle=True, # 训练集打乱
num_workers=0, # Windows设0，避免多线程报错
drop_last=False # 保留最后一个不完整批次（如果有）
)
test_loader = DataLoader(
dataset=test_dataset,
batch_size=2,
shuffle=False, # 测试集不打乱
num_workers=0
)
# 5. 迭代DataLoader（训练/测试时的核心用法）
print("训练集迭代：")
for batch_idx, (x_batch, y_batch) in enumerate(train_loader):
print(f"批次{batch_idx+1}：")
print(f" 输入：{x_batch}, 形状：{x_batch.shape}")
print(f" 标签：{y_batch}, 形状：{y_batch.shape}")
print("\n测试集迭代：")
for x_batch, y_batch in test_loader:
print(f" 输入：{x_batch}, 标签：{y_batch}")
输出示例（训练集shuffle=True，顺序随机）：
训练集迭代：
批次1：
输入：tensor([3, 1]), 形状：torch.Size([2])
标签：tensor([0, 0]), 形状：torch.Size([2])
批次2：
输入：tensor([4, 2]), 形状：torch.Size([2])
标签：tensor([1, 1]), 形状：torch.Size([2])
批次3：
输入：tensor([6, 5]), 形状：torch.Size([2])
标签：tensor([1, 0]), 形状：torch.Size([2])
测试集迭代：
输入：tensor([7, 8]), 标签：tensor([0, 1])
2. 高频场景：NLP中的collate_fn（处理变长序列）
NLP任务中，不同样本的序列长度不同，需要用collate_fn统一长度（padding）：
import torch
from torch.utils.data import DataLoader
from torch.nn.utils.rnn import pad_sequence
# 模拟NLP数据：不同长度的句子（token id序列）
nlp_data = [
([1, 2, 3], 0), # 句子1：长度3，标签0
([4, 5], 1), # 句子2：长度2，标签1
([6, 7, 8, 9], 0) # 句子3：长度4，标签0
]
class NLPDataset(Dataset):
def __init__(self, data):
self.data = data
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
tokens, label = self.data[idx]
return torch.tensor(tokens), torch.tensor(label)
# 自定义collate_fn：统一批次内的序列长度
def collate_fn(batch):
# batch是列表，每个元素是 (tokens_tensor, label_tensor)
tokens_list = [item[0] for item in batch]
labels_list = [item[1] for item in batch]
# 对tokens做padding，补0到当前批次的最大长度
tokens_padded = pad_sequence(tokens_list, batch_first=True, padding_value=0)
labels = torch.tensor(labels_list)
return tokens_padded, labels
# 创建DataLoader
nlp_loader = DataLoader(
dataset=NLPDataset(nlp_data),
batch_size=2,
shuffle=True,
collate_fn=collate_fn # 应用自定义整理函数
)
# 迭代查看效果
for tokens, labels in nlp_loader:
print("padding后的tokens：", tokens)
print("标签：", labels)
print("tokens形状：", tokens.shape)
print("---")
输出示例：
padding后的tokens： tensor([[1, 2, 3, 0],
[4, 5, 0, 0]])
标签： tensor([0, 1])
tokens形状： torch.Size([2, 4])
---
padding后的tokens： tensor([[6, 7, 8, 9]])
标签： tensor([0])
tokens形状： torch.Size([1, 4])
DataLoader迭代逻辑（训练时的典型用法）
# 训练循环示例
epochs = 3
for epoch in range(epochs):
print(f"===== Epoch {epoch+1} =====")
# 训练模式
model.train()
total_loss = 0.0
for batch_idx, (x, y) in enumerate(train_loader):
# 1. 数据移到GPU
x = x.to(device)
y = y.to(device)
# 2. 清零梯度
optimizer.zero_grad()
# 3. 前向传播
outputs = model(x)
loss = criterion(outputs, y)
# 4. 反向传播+更新参数
loss.backward()
optimizer.step()
total_loss += loss.item()
# 打印批次信息
if (batch_idx + 1) % 100 == 0:
print(f"Batch {batch_idx+1}, Loss: {total_loss/(batch_idx+1):.4f}")
# 验证/测试
model.eval()
with torch.no_grad():
correct = 0
total = 0
for x, y in test_loader:
x = x.to(device)
y = y.to(device)
outputs = model(x)
_, predicted = torch.max(outputs.data, 1)
total += y.size(0)
correct += (predicted == y).sum().item()
print(f"Test Accuracy: {100 * correct / total:.2f}%")
总结
DataLoader是PyTorch加载批量数据的核心，需配合Dataset使用，核心能力是批次划分、打乱、多线程加载；- 关键参数：
batch_size（批次大小）、shuffle（训练打乱）、num_workers（多线程）、collate_fn（自定义数据整理）； - 避坑要点：Windows下
num_workers=0，训练集shuffle=True，GPU训练pin_memory=True； - 核心场景：CV/NLP/大模型训练的批量数据迭代，是连接数据集和模型训练的桥梁。
掌握DataLoader的用法，是实现高效、稳定的深度学习训练的基础，尤其是处理大规模数据或变长序列时，collate_fn和num_workers的配置直接影响训练效率。