PyTorch DataLoader详解

小蓝 • 2024-12-15 12:41 • 编程

一、简介

torch.utils.data.DataLoader是PyTorch自带的一个数据加载器，常用于加载大规模数据集，尤其是超越了内存大小的数据集。其主要作用是把数据按照batch_size大小分成若干个小batch，然后在每个batch内部进行并行读取数据，最后把一个batch的数据在返回给用户。

DataLoader主要有以下几个特点：

1. 多线程：DataLoader有一个num_workers参数，可以设置多个线程同时读取数据，可以加快数据读取速度。

2. 数据打乱：DataLoader有一个shuffle参数，可以打乱数据集，让模型学习更加robust。

3. 预处理：DataLoader可以传入自己的预处理函数，对数据集进行必要的变换，如数据增强、标准化等。

4. 可迭代：DataLoader继承了Python的迭代器协议，可以方便地使用Python的for循环进行迭代。

二、用法

创建DataLoader实例的方式非常简单，只需要传入数据集、batch_size即可。下面的代码演示了如何使用DataLoader进行数据集的读取：

    import torch
    import torch.utils.data as data
    
    train_data = data.TensorDataset(torch.Tensor([1, 2, 3, 4]), torch.Tensor([2, 4, 6, 8]))
    train_loader = data.DataLoader(train_data, batch_size=2, shuffle=True)

    for x, y in train_loader:
        print(x, y)

上述代码创建了一个包含四个样本的TensorDataset，并利用DataLoader将其划分为batch_size为2的小batch，并进行了数据打乱，然后通过for循环来遍历整个数据集。

三、常用参数

1. dataset

dataset参数是DataLoader的第一个参数，一般是一个Dataset对象，可以是PyTorch自带的一些数据集，也可以是用户自定义的数据集。Dataset对象本身也是一个抽象基类，需要实现__len__和__getitem__两个方法。

2. batch_size

batch_size是指每个小batch的大小，默认是1。一般会根据内存大小和数据集的大小来进行设置，过小会造成CPU、GPU的空闲时间增加，过大会导致内存不足。一般情况下，batch_size的值是2的指数。

3. shuffle

shuffle参数是用于打乱数据集，让模型更加robust。它可以在每个epoch开始时打乱数据集，也可以在DataLoader初始化时就进行打乱。一般情况下，打乱数据集的方式是随机打乱数据集的样本顺序，从而避免网络过拟合，提高模型泛化性能。

4. num_workers

num_workers参数是用于设置使用的线程数，默认值是0，即不使用线程。如果需要用到多线程读取数据，可以设置num_workers参数，一般设置成CPU核数的一半即可。常用取值范围是0-8。

5. pin_memory

pin_memory参数是用于设置是否将数据保存在CUDA支持的固定内存中，这样可以避免重复的显存和内存之间的数据传输，提高数据读取和使用的速度。但是，这个参数只在使用CUDA方式时生效。

6. drop_last

drop_last参数是用于当batch_size不能整除数据集长度时，是否丢弃最后一个缺少数据的batch。一般情况下，不建议丢弃缺少数据的batch，因为这会导致一些数据得不到使用，从而影响模型性能。

四、总结

通过本文对PyTorch DataLoader进行详细的介绍，我们可以发现DataLoader是PyTorch中一个很重要的模块，可以实现数据加载、数据打乱、数据预处理、多线程等功能，避免了手动完成这些繁琐的工作。因此，我们可以在实际的深度学习任务中积极地使用DataLoader模块，从而提高整个模型训练过程的效率。

原创文章，作者：小蓝，如若转载，请注明出处：https://www.506064.com/n/256539.html