torch.utils.data简介

import torch
import torch.utils.data as data
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import matplotlib.pyplot as plt
import numpy as np
from torch.utils.data.dataset import T_co


# 定义获取数据集的类
class TestDataset(data.Dataset):

    def __init__(self):
        self.Data = np.array([[1, 2], [3, 4], [2, 1], [3, 4], [4, 5]])
        self.label = np.array([0, 1, 0, 1, 2])

    def __getitem__(self, index):
        # 将numpy转换为Tensor
        txt = torch.from_numpy(self.Data[index])
        label = torch.tensor(self.label[index])
        return txt, label

    def __len__(self):
        return len(self.Data)


# 获取数据集中的数据
test = TestDataset()
# print(test[2])  # 相当于调用__getitem__(2)
# print(test.__len__())

'''
结果：
(tensor([2, 1]), tensor(0))
5

以上数据以tuple返回，每次只返回一个样本。实际上，Dateset只负责数据的抽取，调 
用一次__getitem__只返回一个样本。如果希望批量处理（batch），还要同时进行shuffle和 
并行加速等操作，可选择DataLoader。
'''
test_loader = data.DataLoader(dataset=test, batch_size=2, shuffle=True)
for step, train_data in enumerate(test_loader):
    print('step:{}'.format(step))
    x, y = train_data
    print(f'x:{x}')
    print(f'y:{y}')

'''
结果:
step:0
x:tensor([[4, 5],
        [2, 1]])
y:tensor([2, 0])
step:1
x:tensor([[1, 2],
        [3, 4]])
y:tensor([0, 1])
step:2
x:tensor([[3, 4]])
y:tensor([1])
'''