浅谈torch.utils.data.TensorDataset和torch.utils.data.DataLoader

1.torch.utils.data.TensorDataset

功能定位

torch.utils.data.TensorDataset 是一个将多个张量(Tensor)数据进行简单包装整合的数据集类,它主要的作用是将相关联的数据(比如特征数据和对应的标签数据等)组合在一起,形成一个方便后续用于训练等操作的数据集对象。

例如,如果你有输入特征数据 x(形状为 [n_samples, feature_dim])和对应的标签数据 y(形状为 [n_samples]),且它们都是 torch.Tensor 类型,可以这样创建 TensorDataset

python 复制代码
import torch
from torch.utils.data import TensorDataset

x = torch.randn(100, 10)  # 模拟100个样本,每个样本特征维度为10
y = torch.randint(0, 2, (100,))  # 模拟二分类标签

dataset = TensorDataset(x, y)
特点
  • 简单包装:只是把传入的张量按照样本维度进行了对应组合,并没有对数据做复杂的预处理、采样等额外操作。

  • 索引支持 :支持像普通列表那样通过索引访问其中的数据元素,例如 dataset[0] 会返回由对应索引的特征和标签组成的元组(按照传入构造函数的张量顺序)。

  • 适用于小型数据集直接使用:当数据量不大且数据格式已经整理为张量形式时,可以直接基于它来进行简单的模型训练循环等操作,不过对于批量处理等更复杂的情况支持有限,需要进一步配合其他工具。

2.torch.utils.data.DataLoader

功能定位

torch.utils.data.DataLoader 是一个用于加载数据的工具类,它围绕着给定的数据集(比如 TensorDataset 或者自定义的继承自 Dataset 的类实例等),实现了诸如批量加载数据、打乱数据顺序、并行加载数据等功能,旨在让数据能够以合适的方式、合适的批量大小等被送入到模型中进行训练、验证或测试等操作。

示例:

python 复制代码
from torch.utils.data import DataLoader

batch_size = 10
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

for batch_x, batch_y in dataloader:
    # 这里的batch_x和batch_y就是每次迭代取出的一个批量的特征和标签数据
    pass
特点
  • 批量处理 :可以按照设定的 batch_size 参数,将数据集中的数据划分为一个个的小批量(mini-batch),方便模型以批量的方式进行梯度计算更新,有助于优化训练过程和提升效率,尤其在大数据集场景下优势明显。

  • 数据打乱 :通过设置 shuffle=True 可以在每个训练轮次(epoch)开始时对数据集里面的数据顺序进行随机打乱,使得数据的输入顺序具有随机性,这有助于提升模型训练的泛化能力,避免模型因数据顺序固定而产生过拟合等问题。

  • 并行加载 :支持多进程加载数据(通过设置 num_workers 参数大于 0),能够利用多核 CPU 的优势加快数据读取和预处理的速度,特别是在处理大规模数据集或者数据加载比较耗时的情况下,能显著提升整体训练效率。

  • 灵活性和通用性 :它可以适配各种不同类型的数据集,只要这些数据集继承自 torch.utils.data.Dataset 抽象类并实现了必要的 __len____getitem__ 等方法,因此无论是简单的 TensorDataset 还是复杂的自定义数据集都可以用它来加载数据。

总的来说,TensorDataset 侧重于对已有张量数据进行简单的整合包装形成数据集;而 DataLoader 侧重于围绕数据集实现数据的批量加载、打乱顺序、并行化等复杂的数据加载相关功能,它们通常配合使用,先使用 TensorDataset 组织好数据,再通过 DataLoader 按照训练需求来加载和处理这些数据并送入模型中。

python 复制代码
from torch.utils.data import TensorDataset
from torch.utils.data import DataLoader

train_ds = TensorDataset(x_train, y_train)
train_dl = DataLoader(train_ds, batch_size=bs, shuffle=True)

valid_ds = TensorDataset(x_valid, y_valid)
valid_dl = DataLoader(valid_ds, batch_size=bs)
相关推荐
阿坡RPA7 小时前
手搓MCP客户端&服务端:从零到实战极速了解MCP是什么?
人工智能·aigc
用户27784491049938 小时前
借助DeepSeek智能生成测试用例:从提示词到Excel表格的全流程实践
人工智能·python
机器之心8 小时前
刚刚,DeepSeek公布推理时Scaling新论文,R2要来了?
人工智能
算AI10 小时前
人工智能+牙科:临床应用中的几个问题
人工智能·算法
凯子坚持 c11 小时前
基于飞桨框架3.0本地DeepSeek-R1蒸馏版部署实战
人工智能·paddlepaddle
你觉得20511 小时前
哈尔滨工业大学DeepSeek公开课:探索大模型原理、技术与应用从GPT到DeepSeek|附视频与讲义下载方法
大数据·人工智能·python·gpt·学习·机器学习·aigc
8K超高清11 小时前
中国8K摄像机:科技赋能文化传承新图景
大数据·人工智能·科技·物联网·智能硬件
hyshhhh12 小时前
【算法岗面试题】深度学习中如何防止过拟合?
网络·人工智能·深度学习·神经网络·算法·计算机视觉
薛定谔的猫-菜鸟程序员12 小时前
零基础玩转深度神经网络大模型:从Hello World到AI炼金术-详解版(含:Conda 全面使用指南)
人工智能·神经网络·dnn
币之互联万物12 小时前
2025 AI智能数字农业研讨会在苏州启幕,科技助农与数据兴业成焦点
人工智能·科技