PyTorch学习笔记(二)数据加载器

发表于 分类于 ， Waline： 阅读次数： 本文字数： 718 阅读时长 ≈ 2 分钟

　　在定义好数据集后，PyTorch提供了一个DataLoader类作为数据集的迭代器，该类位于torch.utils.data模块中，我们仅需要在其构造函数中传入指定的数据集，即可直接使用for语句使用。除此之外，构造函数还可以接收一些实用的参数：

  • batch_size：指定每一轮迭代载入的数据数量(默认值：1)
  • shuffle：是否随机打乱顺序(默认值：否)
  • num_workers：指定用于载入数据的子线程数(默认值：0)

代码示例

　　在定义好data_loader后，直接使用for语句迭代，每次返回batch_size组数据，数据的类型即为dataset中定义的__getitem__(self, index)方法返回的类型。

1
2
3
4
5
6
7
8
9
10
11
 
from torch.utils.data import DataLoader

BATCH_SIZE = 64

dataset = BallDataset(os.path.join(".", "data", "ball.txt"))
data_loader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)

# 使用
for labels, inputs in data_loader:
    outputs = model(inputs)
    loss = criterion(outputs, labels)

参考文献

扩展阅读