pytorch-TensorDataset与DataLoader使用

pytorch-TensorDataset与DataLoader使用

1, TensorDataset：传入在第一维相同的张量。其实相当于zip.

from torch.utils.data import TensorDataset

a = torch.arange(4)
b = torch.arange(4)
c = TensorDataset(a, b)
for item in c:
    print(item)
'''
(tensor(0), tensor(0))
(tensor(1), tensor(1))
(tensor(2), tensor(2))
(tensor(3), tensor(3))
'''
for item in zip(a, b):
    print(item)
'''
(tensor(0), tensor(0))
(tensor(1), tensor(1))
(tensor(2), tensor(2))
(tensor(3), tensor(3))
'''

2,DataLoader： 按照第一维每次返回batch_size大小的数据。方便之处在于可以打乱数据，并且设置多线程。返回列表，并且数据被包装成[batch_size, 数据]张量形式。

from torch.utils.data import DataLoader

loader = DataLoader(c, batch_size=3, shuffle=True)
for item in loader:
    print(item)
 '''
[tensor([0, 2, 1]), tensor([0, 2, 1])]
[tensor([3]), tensor([3])]
 '''