pytorch api文档：DataLoader()-马育民老师

# 介绍

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