Pytorch中TensorDataset和DataLoader用法

原创已于 2023-12-19 16:13:14 修改 · 1.2k 阅读

7 ·

本内容遵循CC 4.0 BY-SA版权协议

标签

#pytorch #人工智能 #python #深度学习

于 2023-12-18 19:18:52 首次发布

本文详细介绍了PyTorch中的TensorDataset和DataLoader功能，包括如何使用它们打包和组织数据，以及DataLoader的批处理、打乱数据和并行加载特性。重点展示了如何创建和操作这两个工具，以及在训练循环中迭代数据的方法。

我们在Pytorch中通常使用torch.utils模块进行数据操作，其中TensorDataset是 PyTorch 中的一个工具类，用于将一组张量包装成一个数据集。DataLoader则是对数据进行采样、分割等操作，将数据划分为许多组特征张量+对应标签的形式。

具有划分小批量的功能DataLoader提供了一个简单的方式来迭代地加载数据集，它可以接受任意形式的数组、张量作为输入，并把他们一次性转换为神经网络可以接入的tensor类型。

TensorDataset

主要特性

TensorDataset可以将多个张量打包成一个数据集，其中每个张量包含相同数量的元素。即将最外面的维度一致的tensor进行打包，也就是将第一个维度一致的tensor进行打包。

TensorDataset允许你使用索引和切片来访问数据，这使得它的行为类似于一个标准的 Python 列表或数组。

示例

引用自：三、数据集打包—TensorDataset-CSDN博客

from torch.utils.data import TensorDataset
import torch
from torch.utils.data import DataLoader

data_a = torch.tensor(
[[11, 22, 33], [44, 55, 66], [77, 88, 99], [11, 22, 33], [44, 55, 66], [77, 88, 99], [11, 22, 33], [44, 55, 66],[77, 88, 99], [11, 22, 33], [44, 55, 66], [77, 88, 99]])
data_b = torch.tensor([0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2])
train_ids = TensorDataset(data_a, data_b)
# 切片输出
print(train_ids[0:2])
print('#' * 30)
# 循环取数据
for x_train, y_label in train_ids:
    print(x_train, y_label)
# DataLoader进行数据封装
print('#' * 30)
train_loader = DataLoader(dataset=train_ids, batch_size=4, shuffle=True)
for i, data in enumerate(train_loader, 1):  # 注意enumerate返回值有两个,一个是序号,一个是数据(包含训练数据和标签)
    x_data, label = data
    print('batch:{0} x_data:{1} label: {2}'.format(i, x_data, label))   # y data (torch tensor)

索引和切片操作

import torch
from torch.utils.data import TensorDataset, DataLoader

# 示例特征和标签张量
features = torch.tensor([[1, 2], [3, 4], [5, 6]])
labels = torch.tensor([0, 1, 0])

# 创建 TensorDataset
dataset = TensorDataset(features, labels)

# 获取第一个样本
first_sample = dataset[0]

# 获取前两个样本
first_two_samples = dataset[:2]

DataLoader

主要特性

批处理（Batching）：DataLoader允许您以批次的形式加载数据，这对于内存效率和训练速度都非常关键。
打乱数据（Shuffling）：在每个 epoch 开始时，DataLoader可以选择性地打乱数据，这有助于减少模型的过拟合。
并行加载（Parallelism）：DataLoader可以使用多个工作进程来并行加载数据，从而提高数据加载的速度。
自定义数据加载: 它允许使用自定义的数据集类（只要其继承自torch.utils.data.Dataset），几乎可以加载任何类型的数据。

主要参数

Dataset：加载的数据集
batch_size：每个批次加载的样本数。较大的batch_size可以提高内存利用率和加速训练，但也可能导致内存不足。
shuffle：指定是否在每个 epoch 开始时随机打乱数据。对于训练数据，通常设置为True以减少模型过拟合。
num_workers：用于数据加载的子进程数。更多的num_workers可以加快数据加载速度，但也会增加内存消耗。
drop_last：如果数据集大小不能被batch_size整除，设置为True时将丢弃最后一个不完整的批次，反之则保留。
sampler：定义从数据集中抽取样本的策略。如果指定，则shuffle必须为False。

from torch.utils.data import DataLoader

# 假设我们有一个数据集对象
dataset = MyDataset(...)  # MyDataset 是一个继承自 torch.utils.data.Dataset 的类

# 创建 DataLoader
batchdata = DataLoader(dataset            # 加载的数据集
                        , batch_size=32     # 每个批次加载的样本数
                        , shuffle=True      # 是否随机打乱数据
                        , num_workers=4)    # 数据加载的子进程数

len(batchdata) #查看具体被分了多少个batch

#可以使用.datasets查看数据集相关的属性
len(batchdata.dataset) #总共有多少数据

batchdata.dataset[0] #查看其中一个样本

batchdata.dataset[0][0]#一个样本的特征张量

batchdata.dataset[0][1]#一个样本的标签

#属性batch_size，查看现在的batch_size是多少
batchdata.batch_size

迭代数据

使用DataLoader的一般模式是在训练循环中迭代它：(两种循环方式都可以）

for batch_idx, (x, y) in enumerate(data_loader):
    # 在这里处理数据和目标
    # 例如，将数据和目标移至 GPU，然后进行模型的前向和后向传播

for data, target in data_loader:

batch_idx是当前批次的索引号，它是一个整数，从 0 开始计数。每次循环迭代时batch_idx会自动增加。

（x, y）是一个元组，其中x和y分别代表当前批次中的数据和标签。

x通常是一个张量（Tensor），包含了一批样本数据。在深度学习中，这些样本数据可能是图像、文本、声音片段或其他类型的输入数据。x可以理解成三维或二维数据通过循环不断从data_loader中提出。

y也是一个张量，包含了与 x中每个样本相对应的标签或目标值。在监督学习中，这些标签用于训练模型，使其能够从输入数据 x预测出正确的输出。

这个循环的过程可以理解为一张张图片的数据‘x’提取后送入神经网络进行训练，后由‘y’进行损失计算。

结合使用

data_loader = DataLoader(dataset, batch_size=2, shuffle=True)

for batch_idx, (batch_features, batch_labels) in enumerate(data_loader):
    # 在这里处理每个批次的数据
    pass