实战pytorch中的utils.data.TensorDataset和utils.data.DataLoader两个工具类

568 阅读1分钟

本文主要说明pytorch框架中utils.data.TensorDataset和utils.data.DataLoader两个工具类的用法。

一、定义

1、TensorDataset类:是torch中util.data包下面的一个工具类,用于把特征和标签整合一个数据集,便于将整个数据集拆分小批量数据集,用于后续小批量梯度下降训练。其输入参数是可变参数,可接受一个或多个张量数据。

2、DataLoader类:是torch中util.data包下面的一个工具类,用于将TensorDataset整合的数据集(比较大,比如1000个特征和1000个标签)划分为小批量数据级,比如批量大小为10,就会划分为100个。

DataLoader函数的参数后面如下:来源于pytorch中文文档, pytorch-cn.readthedocs.io/zh/latest/p…

image.png 在使用中,一般指定dataset, batch_size, shuffle三个参数即可。

二、实战

首先导入torch和torch.utils库,然后验证输入为一个张量的情况 image.png 再验证输入为两个张量(一个是特征,另一个是标签)的情况 image.png 利用TensorDataset()将上述定义的特征和标签合并为一个数据集,从循环打印的结果可以看到,数据集包括6组特征和对应的6组标签。 image.png 将数据集中的值打乱顺序,划分为3个批量大小为2的数据集 image.png

至此,utils.data.TensorDataset和utils.data.DataLoader两个工具类验证完毕。