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 :

<class 'torch.Tensor'>
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的情景下，數據讀取速度的變化。