pytorch在torch.utils.data 对常用数据加载进行了封装，可用很容易实现数据读取和批量加载。

DataSet#

DataSet 是 pytorch 提供的用于包装数据的抽象类。可以通过继承并实现其中的抽象方法来自定义DataSet。

自定义DataSet要继承并实现两个成员方法：

__getitem__(self, idx) ：该方法需要实现通过索引获得一条数据。
__len__(self) ：该方法需要返回数据集的长度。

DataLoader#

DataLoader提供了对DataSet的读取操作，常用参数有：

batch_size ：每个批次的大小。
shuffle ：是否对数据进行洗牌操作。
num_work：加载数据时使用几个子进程。

示例#

在样例中，生成了 $[0, 400)$ 的整数的序列，并将其转换为 $[100, 2, 2]$ 的矩阵。

示例代码：

1
from torch.utils.data import Dataset
2
import torch
3

4
class MyDataset(Dataset):
5
    def __init__(self, data):
6
        self.data = data
7

8
    def __len__(self):
9
        return len(self.data)
10

11
    def __getitem__(self, index):
12
        return self.data[index]
13

14
if __name__ == '__main__':
15
    # 创建数据
16
    data = torch.range(0, 399)
17
    data = data.view(100, 2, 2)
18
    # 创建自己的DataSet
19
    ds = MyDataset(data)
20
    # 输出数据信息
21
    print('数据集长度：', len(ds))
22
    print('ds[0]: ', ds[0])
23

24
    # DataLoader提供了DataSet的读取操作
25
    dl = torch.utils.data.DataLoader(
26
        ds, batch_size=2, shuffle=True, num_workers=0
27
    )
28

29
    # 获取一批数据，即前2个数据
30
    itDl = iter(dl)
31
    print('next(itDl): ', next(itDl))

输出结果：

1
数据集长度： 100
2
ds[0]:  tensor([[0., 1.],
3
        [2., 3.]])
4
next(itDl):  tensor([[[276., 277.],
5
         [278., 279.]],
6

7
        [[284., 285.],
8
         [286., 287.]]])

从结果可以看出，DataLoader对数据进行了洗牌，并以每批次2个数据输出。

注意：

这里的MyDataSet 仅做了简单实现，并不一定只能传入数据，也可以传入文件路径等，然后对数据进行读取并保存。

如：

1
#定义一个数据集
2
class BulldozerDataset(Dataset):
3
    """ 数据集演示 """
4
    def __init__(self, csv_file):
5
        """实现初始化方法，在初始化的时候将数据读载入"""
6
        self.df=pd.read_csv(csv_file)
7
    def __len__(self):
8
        '''
9
        返回df的长度
10
        '''
11
        return len(self.df)
12
    def __getitem__(self, idx):
13
        '''
14
        根据 idx 返回一行数据
15
        '''
16
        return self.df.iloc[idx].SalePrice

原文链接：CSDN