data.DataLoader
是 PyTorch 中的一个工具,用于从给定的数据集中加载数据,实现批量化的训练。它是在训练神经网络时常用的一个组件,能够方便地迭代数据集并生成小批量数据。
下面是一个简单的例子,演示如何使用 DataLoader
:
import torch
from torch.utils.data import DataLoader, TensorDataset
# 假设有特征张量 features 和标签张量 labels
features = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
labels = torch.tensor([0, 1, 0])
# 使用 TensorDataset 创建数据集
dataset = TensorDataset(features, labels)
# 使用 DataLoader 加载数据集
batch_size = 2
data_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
# 迭代数据集中的小批量数据
for batch in data_loader:
batch_features, batch_labels = batch
print("Batch Features:", batch_features)
print("Batch Labels:", batch_labels)
在这个例子中,DataLoader
将 TensorDataset
中的数据按照指定的批量大小 (batch_size
) 进行加载。shuffle=True
表示在每个 epoch 开始时打乱数据,以增加模型的泛化能力。通过迭代 data_loader
,你可以逐批次地获取数据用于训练。
这种批量加载的方式对于大规模数据集和复杂模型的训练是非常有效的。