最近做实验跑baseline就很想有一个自己的统一的深度学习框架，于是乎，从读取数据集开始吧。

TensorDataset和Dataloader

使用torch.utils.data下的TensorDataset和Dataloader，可以很方便地读取数据集，一个简单的例子如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
import torch
import torch.utils.data as Data  # 直接用torch.utils.data.xxx是不好使的

features = torch.randn(200,3)
labels = torch.randn(200,1)
mask = torch.randn(20,1)
batch_size = 10

# 将训练数据的特征和标签组合
dataset = Data.TensorDataset(features, labels)
# 随机读取小批量
data_iter = Data.DataLoader(dataset, batch_size, shuffle=True)
for x,y in data_iter:
    # 打印一个batch的数据，可以看到size为10x3
    print(x.shape,y.shape) # torch.Size([10, 3]) torch.Size([10, 1])
    break

但是上述的TensorDataset只能读入数据的特征和对应的标签，当我们需要复杂一些的输入数据时，比如BERT输入还要带上mask什么的，就需要自己改写这个Dataset，然后利用DataLoader方便地读入。

继承Dataset类自定义

PyTorch允许自由地对Dataset类执行任何操作，只要重写改类中的两个函数即可：

  • __len__ 函数：返回数据集大小（忘记在哪看的，说这个不定义也可，无石锤）
  • __getitem__ 函数：返回对应索引的数据集中的样本

数据集的大小有时难以确定，但它等于整个数据集中的样本数量。因此，如果数据集中有10,000个样本（数据点，图像，句子等），则__len__函数应返回10,000。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
import torch
import torch.utils.data as Data  # 直接用torch.utils.data.xxx是不好使的

class MyDataset(Data.Dataset):
    def __init__(self): # 参数可以写数据集的路径，然后读取
        super(MyDataset, self).__init__()
        self.data1 = list(range(1,50))
        self.data2 = list(range(51,100))
        self.data3 = list(range(101,150))
    
    def __getitem__(self, index):
        return self.data1[index], self.data2[index],self.data3[index] # 可以自己定义需要返回的数据
    
    def __len__(self):
        return len(self.data1)

data_set = MyDataset()
print(data_set.data1) # 1-50个数的列表
print(data_set[1])    # (2, 52, 102)
print(data_set[1:3])  # ([2, 3], [52, 53], [102, 103])

data_loader = Data.DataLoader(data_set, batch_size=4,shuffle=False)
print(len(data_set))  # 49
for batch_data in data_loader: 
	# 打印一个batch的数据
    print(batch_data) # [tensor([1, 2, 3, 4]), tensor([51, 52, 53, 54]), tensor([101, 102, 103, 104])]
    break