【PyTorch深度学习快速入门|代码逐行解析|小白必备】1.加载数据

1.PyTorch加载数据集

from torch.utils.data import Dataset
from PIL import Image
import os
class MyData(Dataset):#继承Dataset类

    def __init__(self,root_dir,label_dir):#初始化，为class类提供全局变量，为后面函数提供变量
        self.root_dir = root_dir#self指定一个类中的全局变量
        self.label_dir = label_dir
        self.path = os.path.join(self.root_dir,self.label_dir)#拼接两个路径
        self.img_path = os.listdir(self.path)#获得图片所有list

    def __getitem__(self, idx):#index为编号(首先将图片文件放入list)
        img_name = self.img_path[idx]#读取某个对应位置图片
        img_item_path = os.path.join(self.root_dir,self.label_dir,img_name)#相对于程序的相对路径
        img = Image.open(img_item_path)#读取图片
        label = self.label_dir#读取label
        return img,label

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

root_dir = "dataset/train"#获取路径
ants_label_dir = "ants"#获取数据集名称
bees_label_dir = "bees"#获取数据集名称
ants_dataset = MyData(root_dir,ants_label_dir)#蚂蚁数据集
bees_dataset = MyData(root_dir,bees_label_dir)#蜜蜂数据集

train_dataset = ants_dataset + bees_dataset#数据集拼接

如有错误，欢迎批评指正。