Pytorch加载数据集的方式总结及补充

import torch
import numpy as np

# 定义GetLoader类，继承Dataset方法，并重写__getitem__()和__len__()方法
class GetLoader(torch.utils.data.Dataset):
	# 初始化函数，得到数据
    def __init__(self, data_root, data_label):
        self.data = data_root
        self.label = data_label
    # index是根据batchsize划分数据后得到的索引，最后将data和对应的labels进行一起返回
    def __getitem__(self, index):
        data = self.data[index]
        labels = self.label[index]
        return data, labels
    # 该函数返回数据大小长度，目的是DataLoader方便划分，如果不知道大小，DataLoader会一脸懵逼
    def __len__(self):
        return len(self.data)

# 随机生成数据，大小为10 * 20列
source_data = np.random.rand(10, 20)
# 随机生成标签，大小为10 * 1列
source_label = np.random.randint(0,2,(10, 1))
# 通过GetLoader将数据进行加载，返回Dataset对象，包含data和labels
torch_data = GetLoader(source_data, source_label)

torch.utils.data.DataLoader(dataset,batch_size,shuffle,drop_last，num_workers)

...
torch_data = GetLoader(source_data, source_label)

from torch.utils.data import DataLoader
datas = DataLoader(torch_data, batch_size = 4, shuffle = True, drop_last = False, num_workers = 2)
for i, (data, label) in enumerate(datas):
	# i表示第几个batch， data表示batch_size个原始的数据，label代表batch_size个数据的标签
    print("第 {} 个Batch \n{}".format(i, data))

dataset=torchvision.datasets.ImageFolder(
                       root, transform=None, 
                       target_transform=None, 
                       loader=<function default_loader>, 
                       is_valid_file=None)

print(dataset.classes)  #根据分的文件夹的名字来确定的类别
print(dataset.class_to_idx) #按顺序为这些类别定义索引为0,1...
print(dataset.imgs) #返回从所有文件夹中得到的图片的路径以及其类别
'''
输出：
['cat', 'dog']
{'cat': 0, 'dog': 1}
[('./data/train\\cat\\1.jpg', 0), 
 ('./data/train\\cat\\2.jpg', 0), 
 ('./data/train\\dog\\1.jpg', 1), 
 ('./data/train\\dog\\2.jpg', 1)]
'''

# 5. 将文件夹数据导入
train_loader = torch.utils.data.DataLoader(dataset,
                                           batch_size = batch_size, shuffle=True,
                                           num_workers = 2)
# 6. 传入网络进行训练
for epoch in range(epochs):
    train_bar = tqdm(train_loader, file = sys.stdout)
    for step, data in enumerate(train_bar):
    ...

# 5. 将文件夹数据导入
train_loader = torch.utils.data.DataLoader(dataset,
                                           batch_size = batch_size, shuffle=True,
                                           num_workers = 2)
# 6. 传入网络进行训练
for epoch in range(epochs):
    train_bar = tqdm(train_loader, file = sys.stdout)
    for step, data in enumerate(train_bar):
    ...

from torch.utils.data import DataLoader
from torchvision import datasets, transforms

def dataloader(dataset, input_size, batch_size, split='train'):
    transform = transforms.Compose([
        					transforms.Resize((input_size, input_size)), 
       					    transforms.ToTensor(), 
        					transforms.Normalize(mean=[0.5], std=[0.5])
    ])
    if dataset == 'mnist':
        data_loader = DataLoader(
            datasets.MNIST('data/mnist', train=True, download=True, transform=transform),
            batch_size=batch_size, shuffle=True)
    elif dataset == 'fashion-mnist':
        data_loader = DataLoader(
            datasets.FashionMNIST('data/fashion-mnist', train=True, download=True, transform=transform),
            batch_size=batch_size, shuffle=True)
    elif dataset == 'cifar10':
        data_loader = DataLoader(
            datasets.CIFAR10('data/cifar10', train=True, download=True, transform=transform),
            batch_size=batch_size, shuffle=True)
    elif dataset == 'svhn':
        data_loader = DataLoader(
            datasets.SVHN('data/svhn', split=split, download=True, transform=transform),
            batch_size=batch_size, shuffle=True)
    elif dataset == 'stl10':
        data_loader = DataLoader(
            datasets.STL10('data/stl10', split=split, download=True, transform=transform),
            batch_size=batch_size, shuffle=True)
    elif dataset == 'lsun-bed':
        data_loader = DataLoader(
            datasets.LSUN('data/lsun', classes=['bedroom_train'], transform=transform),
            batch_size=batch_size, shuffle=True)

    return data_loader

transforms.Compose([transforms.RandomResizedCrop(224),
 		    	   transforms.RandomHorizontalFlip(),
                    transforms.ToTensor(),
                    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])

"""
    批训练，把数据变成一小批一小批数据进行训练。
    DataLoader就是用来包装所使用的数据，每次抛出一批数据
"""
import torch
import torch.utils.data as Data

BATCH_SIZE = 5

x = torch.linspace(1, 10, 10)
y = torch.linspace(10, 1, 10)
# 把数据放在数据库中
torch_dataset = Data.TensorDataset(x, y)
loader = Data.DataLoader(
    # 从数据库中每次抽出batch size个样本
    dataset=torch_dataset,
    batch_size=BATCH_SIZE,
    shuffle=True,
    num_workers=2,
)

def show_batch():
    for epoch in range(3):
        for step, (batch_x, batch_y) in enumerate(loader):
            # training
            print("steop:{}, batch_x:{}, batch_y:{}".format(step, batch_x, batch_y))

if __name__ == '__main__':
    show_batch()

steop:0, batch_x:tensor([ 3., 10.,  6.,  2.,  8.]), batch_y:tensor([8., 1., 5., 9., 3.])
steop:1, batch_x:tensor([5., 9., 7., 4., 1.]), batch_y:tensor([ 6.,  2.,  4.,  7., 10.])
steop:0, batch_x:tensor([8., 3., 1., 2., 9.]), batch_y:tensor([ 3.,  8., 10.,  9.,  2.])
steop:1, batch_x:tensor([10., 5.,  4.,  7.,  6.]), batch_y:tensor([1., 6., 7., 4., 5.])
steop:0, batch_x:tensor([5., 8., 4., 3., 7.]), batch_y:tensor([6., 3., 7., 8., 4.])
steop:1, batch_x:tensor([ 2., 10.,  6.,  9.,  1.]), batch_y:tensor([ 9.,  1.,  5.,  2., 10.])