三、数据集打包—TensorDataset

数据集打包—TensorDataset

TensorDataset 可以用来对 tensor 进行打包，就好像 python 中的 zip 功能。该类通过每一个 tensor 的第一个维度进行索引。因此，该类中的 tensor 第一维度必须相等。

from torch.utils.data import TensorDataset
import torch
from torch.utils.data import DataLoader

data_a = torch.tensor(
[[11, 22, 33], [44, 55, 66], [77, 88, 99], [11, 22, 33], [44, 55, 66], [77, 88, 99], [11, 22, 33], [44, 55, 66],[77, 88, 99], [11, 22, 33], [44, 55, 66], [77, 88, 99]])
data_b = torch.tensor([0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2])
train_ids = TensorDataset(data_a, data_b)
# 切片输出
print(train_ids[0:2])
print('#' * 30)
# 循环取数据
for x_train, y_label in train_ids:
    print(x_train, y_label)
# DataLoader进行数据封装
print('#' * 30)
train_loader = DataLoader(dataset=train_ids, batch_size=4, shuffle=True)
for i, data in enumerate(train_loader, 1):  # 注意enumerate返回值有两个,一个是序号,一个是数据(包含训练数据和标签)
    x_data, label = data
    print('batch:{0} x_data:{1} label: {2}'.format(i, x_data, label))   # y data (torch tensor)

结果：