import torch
import numpy as np

# 定義GetLoader類，繼承Dataset方法，并重寫__getitem__()和__len__()方法
class GetLoader(torch.utils.data.Dataset):
	# 初始化函數(shù)，得到數(shù)據(jù)
    def __init__(self, data_root, data_label):
        self.data = data_root
        self.label = data_label
    # index是根據(jù)batchsize劃分?jǐn)?shù)據(jù)后得到的索引，最后將data和對應(yīng)的labels進(jìn)行一起返回
    def __getitem__(self, index):
        data = self.data[index]
        labels = self.label[index]
        return data, labels
    # 該函數(shù)返回?cái)?shù)據(jù)大小長度，目的是DataLoader方便劃分，如果不知道大小，DataLoader會(huì)一臉懵逼
    def __len__(self):
        return len(self.data)

# 隨機(jī)生成數(shù)據(jù)，大小為10 * 20列
source_data = np.random.rand(10, 20)
# 隨機(jī)生成標(biāo)簽，大小為10 * 1列
source_label = np.random.randint(0,2,(10, 1))
# 通過GetLoader將數(shù)據(jù)進(jìn)行加載，返回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表示第幾個(gè)batch， data表示batch_size個(gè)原始的數(shù)據(jù)，label代表batch_size個(gè)數(shù)據(jù)的標(biāo)簽
    print("第 {} 個(gè)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)  #根據(jù)分的文件夾的名字來確定的類別
print(dataset.class_to_idx) #按順序?yàn)檫@些類別定義索引為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. 將文件夾數(shù)據(jù)導(dǎo)入
train_loader = torch.utils.data.DataLoader(dataset,
                                           batch_size = batch_size, shuffle=True,
                                           num_workers = 2)
# 6. 傳入網(wǎng)絡(luò)進(jìn)行訓(xùn)練
for epoch in range(epochs):
    train_bar = tqdm(train_loader, file = sys.stdout)
    for step, data in enumerate(train_bar):
    ...

# 5. 將文件夾數(shù)據(jù)導(dǎo)入
train_loader = torch.utils.data.DataLoader(dataset,
                                           batch_size = batch_size, shuffle=True,
                                           num_workers = 2)
# 6. 傳入網(wǎng)絡(luò)進(jìn)行訓(xùn)練
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])])

"""
    批訓(xùn)練，把數(shù)據(jù)變成一小批一小批數(shù)據(jù)進(jìn)行訓(xùn)練。
    DataLoader就是用來包裝所使用的數(shù)據(jù)，每次拋出一批數(shù)據(jù)
"""
import torch
import torch.utils.data as Data

BATCH_SIZE = 5

x = torch.linspace(1, 10, 10)
y = torch.linspace(10, 1, 10)
# 把數(shù)據(jù)放在數(shù)據(jù)庫中
torch_dataset = Data.TensorDataset(x, y)
loader = Data.DataLoader(
    # 從數(shù)據(jù)庫中每次抽出batch size個(gè)樣本
    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.])