pytorch api文档:DataLoader() 作者:马育民 • 2026-01-09 20:42 • 阅读:10002 # 介绍 PyTorch中 **加载批量训练/测试数据的核心工具**,能高效处理数据迭代、批量划分、打乱、多线程加载等问题,是深度学习训练流程中不可或缺的组件(不管是CV、NLP还是大模型训练,都会用到)。 # 功能 `DataLoader` 封装了PyTorch的`Dataset`(数据集),提供以下核心能力: 1. 按指定批次大小(`batch_size`)切分数据; 2. 训练时打乱数据(`shuffle=True`),避免过拟合; 3. 多线程/多进程加载数据(`num_workers`),提升训练效率; 4. 自动拼接张量、处理不同长度的数据(如NLP中的padding); 5. 支持自定义数据处理逻辑(如采样器、整理函数)。 # 语法 ```python 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`封装: ```python 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): ```python 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迭代逻辑(训练时的典型用法) ```python # 训练循环示例 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}%") ``` # 总结 1. `DataLoader`是PyTorch加载批量数据的核心,需配合`Dataset`使用,核心能力是批次划分、打乱、多线程加载; 2. 关键参数:`batch_size`(批次大小)、`shuffle`(训练打乱)、`num_workers`(多线程)、`collate_fn`(自定义数据整理); 3. 避坑要点:Windows下`num_workers=0`,训练集`shuffle=True`,GPU训练`pin_memory=True`; 4. 核心场景:CV/NLP/大模型训练的批量数据迭代,是连接数据集和模型训练的桥梁。 掌握`DataLoader`的用法,是实现高效、稳定的深度学习训练的基础,尤其是处理大规模数据或变长序列时,`collate_fn`和`num_workers`的配置直接影响训练效率。 原文出处:http://malaoshi.top/show_1GW2YvI5QpvU.html