def main():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print("using {} device.".format(device))
 
    train_images_path, train_images_label, val_images_path, val_images_label = read_split_data(root)
 
    data_transform = {
        "train": transforms.Compose([transforms.RandomResizedCrop(224),
                                     transforms.RandomHorizontalFlip(),
                                     transforms.ToTensor(),
                                     transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])]),
        "val": transforms.Compose([transforms.Resize(256),
                                   transforms.CenterCrop(224),
                                   transforms.ToTensor(),
                                   transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])}
 
    train_data_set = MyDataSet(images_path=train_images_path,
                               images_class=train_images_label,
                               transform=data_transform["train"])
 
    batch_size = 8
    nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, 8])  # number of workers
    print('Using {} dataloader workers'.format(nw))
    train_loader = torch.utils.data.DataLoader(train_data_set,
                                               batch_size=batch_size,
                                               shuffle=True,
                                               num_workers=nw,
                                               collate_fn=train_data_set.collate_fn)
 
    # plot_data_loader_image(train_loader)
 
    for step, data in enumerate(train_loader):
        images, labels = data

def read_split_data(root: str, val_rate: float = 0.2):
    random.seed(0)  # 保證隨機(jī)結(jié)果可復(fù)現(xiàn)
    assert os.path.exists(root), "dataset root: {} does not exist.".format(root)
 
    # 遍歷文件夾，一個(gè)文件夾對(duì)應(yīng)一個(gè)類別
    flower_class = [cla for cla in os.listdir(root) if os.path.isdir(os.path.join(root, cla))]
    # 排序，保證順序一致
    flower_class.sort()
    # 生成類別名稱以及對(duì)應(yīng)的數(shù)字索引
    class_indices = dict((k, v) for v, k in enumerate(flower_class))
    json_str = json.dumps(dict((val, key) for key, val in class_indices.items()), indent=4)
    with open('class_indices.json', 'w') as json_file:
        json_file.write(json_str)
 
    train_images_path = []  # 存儲(chǔ)訓(xùn)練集的所有圖片路徑
    train_images_label = []  # 存儲(chǔ)訓(xùn)練集圖片對(duì)應(yīng)索引信息
    val_images_path = []  # 存儲(chǔ)驗(yàn)證集的所有圖片路徑
    val_images_label = []  # 存儲(chǔ)驗(yàn)證集圖片對(duì)應(yīng)索引信息
    every_class_num = []  # 存儲(chǔ)每個(gè)類別的樣本總數(shù)
    supported = [".jpg", ".JPG", ".png", ".PNG"]  # 支持的文件后綴類型
    # 遍歷每個(gè)文件夾下的文件
    for cla in flower_class:
        cla_path = os.path.join(root, cla)
        # 遍歷獲取supported支持的所有文件路徑
        images = [os.path.join(root, cla, i) for i in os.listdir(cla_path)
                  if os.path.splitext(i)[-1] in supported]
        # 獲取該類別對(duì)應(yīng)的索引
        image_class = class_indices[cla]
        # 記錄該類別的樣本數(shù)量
        every_class_num.append(len(images))
        # 按比例隨機(jī)采樣驗(yàn)證樣本
        val_path = random.sample(images, k=int(len(images) * val_rate))
 
        for img_path in images:
            if img_path in val_path:  # 如果該路徑在采樣的驗(yàn)證集樣本中則存入驗(yàn)證集
                val_images_path.append(img_path)
                val_images_label.append(image_class)
            else:  # 否則存入訓(xùn)練集
                train_images_path.append(img_path)
                train_images_label.append(image_class)
 
    print("{} images were found in the dataset.".format(sum(every_class_num)))
    print("{} images for training.".format(len(train_images_path)))
    print("{} images for validation.".format(len(val_images_path)))
 
    plot_image = False
    if plot_image:
        # 繪制每種類別個(gè)數(shù)柱狀圖
        plt.bar(range(len(flower_class)), every_class_num, align='center')
        # 將橫坐標(biāo)0,1,2,3,4替換為相應(yīng)的類別名稱
        plt.xticks(range(len(flower_class)), flower_class)
        # 在柱狀圖上添加數(shù)值標(biāo)簽
        for i, v in enumerate(every_class_num):
            plt.text(x=i, y=v + 5, s=str(v), ha='center')
        # 設(shè)置x坐標(biāo)
        plt.xlabel('image class')
        # 設(shè)置y坐標(biāo)
        plt.ylabel('number of images')
        # 設(shè)置柱狀圖的標(biāo)題
        plt.title('flower class distribution')
        plt.show()
 
    return train_images_path, train_images_label, val_images_path, val_images_label

train_data_set = MyDataSet(images_path=train_images_path,
                           images_class=train_images_label,
                           transform=data_transform["train"])

from PIL import Image
import torch
from torch.utils.data import Dataset
 
 
class MyDataSet(Dataset):
    """自定義數(shù)據(jù)集"""
 
    def __init__(self, images_path: list, images_class: list, transform=None):
        self.images_path = images_path
        self.images_class = images_class
        self.transform = transform
 
    def __len__(self):
        return len(self.images_path)
 
    def __getitem__(self, item):
        img = Image.open(self.images_path[item])
        # RGB為彩色圖片，L為灰度圖片
        if img.mode != 'RGB':
            raise ValueError("image: {} isn't RGB mode.".format(self.images_path[item]))
        label = self.images_class[item]
 
        if self.transform is not None:
            img = self.transform(img)
 
        return img, label
 
    @staticmethod
    def collate_fn(batch):
        # 官方實(shí)現(xiàn)的default_collate可以參考
        # https://github.com/pytorch/pytorch/blob/67b7e751e6b5931a9f45274653f4f653a4e6cdf6/torch/utils/data/_utils/collate.py
        images, labels = tuple(zip(*batch))
 
        images = torch.stack(images, dim=0)
        labels = torch.as_tensor(labels)
        return images, labels

@staticmethod
def collate_fn(batch):
    # 官方實(shí)現(xiàn)的default_collate可以參考
    # https://github.com/pytorch/pytorch/blob/67b7e751e6b5931a9f45274653f4f653a4e6cdf6/torch/utils/data/_utils/collate.py
    images, labels = tuple(zip(*batch))
    images = torch.stack(images, dim=0)
    labels = torch.as_tensor(labels)
    return images, labels

train_loader = torch.utils.data.DataLoader(train_data_set,
                                           batch_size=batch_size,
                                           shuffle=True,
                                           num_workers=nw,
                                           collate_fn=train_data_set.collate_fn)

plot_data_loader_image(train_loader)