Pytorch批量数据读取方法详解: DataLoader

在训练模型的过程中，我们需要不断的读取小批量的数据样本。Pytorch提供了data包来读取数据。接下来我将人工生成一些数据，然后使用data包来处理数据。

import torch
import numpy as np
'''
The features number is 3, and the number of examples is 1000.
''' 
true_w = [2, 3, 5.3]
true_b = 9.7
features = torch.tensor(np.random.normal(0, 1, (1000, 3)), dtype=torch.float)
labels = true_w[0] * features[:, 0] + true_w[1] * features[:, 1] \
						+ true_w[2] * features[:, 2] + true_b
noise = torch.tensor(np.random.normal(0, 0.01, size = labels.size()),
						dtype = torch.float32)
labels += noise ## add the noise item
print(features[0],'\n',labels[0])

outputs is:

tensor([0.9225, 0.4987, 0.5109]) 
 tensor(15.7612)

接下来我们使用data.TensorDataset()来讲特征和标签结合起来,然后使用DataLoader来生成一个数据生成器。

import torch.utils.data as Data
dataset = Data.TensorDataset(features, labels)
data_iter = Data.DataLoader(dataset, batch_size=10, shuffle = True)

for item in data_iter:
    print(type(item[0]))
    print(item[0].size())
    print(item)
    break

output is :

torch.Size([10, 3])
[tensor([[-0.1667,  0.0887, -0.7024],
        [ 1.2368, -1.2652,  0.1885],
        [-1.0656,  0.3612,  0.7720],
        [-0.4624, -0.5067, -0.5172],
        [ 1.4083,  0.0131, -0.0313],
        [ 1.5364, -0.6841,  0.3231],
        [ 0.2451,  1.0406, -1.6667],
        [-1.6337, -0.2157,  0.9186],
        [-0.7980,  0.0130, -0.1350],
        [ 0.7879,  1.4658,  0.9814]]), tensor([ 5.9068,  9.3762, 12.7405,  4.5073, 12.3995, 12.4275,  4.4938, 10.6610,
         7.4216, 20.8734])]

在实践中，数据读取的速度往往对模型的训练性能有着很大的影响，为了提高读取数据的性能，我们可能需要使用多个进程来读取数据，DataLoader可以很方便的使用多进程读取数据。
只需在上述代码中加入一个num_workers参数即可：

import torch.utils.data as Data
from multiprocessing import cpu_count

print('CPU核的数量：',cpu_count())##你可以先确定你的计算机上有几个核。

dataset = Data.TensorDataset(features, labels)
data_iter = Data.DataLoader(dataset, batch_size=10, shuffle = True, 
							num_workers = 4)

感谢趣的朋友可以自己试试在不同num_workers的情景下，数据读取速度的变化。