深度学习实战案例一：第一课：加载数据集

油条与肉包

已于 2023-10-30 11:29:50 修改

阅读量233

点赞数

文章标签： python 深度学习机器学习

于 2023-10-26 00:09:36 首次发布

本文链接：https://blog.youkuaiyun.com/qq_44738211/article/details/134046346

版权

文章介绍了如何设计一个名为Pokemon的自定义数据加载类，用于加载和转换皮卡丘等动漫角色的图片数据，以及如何使用DataLoader进行批量加载，并通过Visdom展示数据。还提到使用torchvision.datasets.ImageFolder简化数据加载的场景。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

1:如何设计Class 加载图片转换data set 并输出成Totensor：以pokemon dataset为例

Download: pokemon data

链接：https://pan.baidu.com/s/1o1iblvQyfYw47bk4UPFpbA?pwd=8888
提取码：8888

#数据集：皮卡丘：234，超梦：239，杰尼龟：223，小火龙：238，妙蛙种子：234
import torch
import  os,glob
import  random,csv
from torch.utils.data import Dataset ,DataLoader
from PIL import Image
from  torchvision import transforms

#自定义数据加载类
class Pokemon(Dataset):
    def __init__(self,root,resize,mode):            #root:文件所在目录，resize：图像分辨率调整一致,mode：当前类何功能
        super(Pokemon,self).__init__()

        self.root = root
        self.resize = resize

        self.name2label = {}                      #对每个加载的文件进行编码：'bulbasaur': 0, 'charmander': 1, 'mewtwo': 2, 'pikachu': 3, 'squirtle': 4
        for name in sorted(os.listdir((os.path.join(root)))):#对指定root中的文件进行排序
            if not os.path.isdir(os.path.join(root,name)):
                continue
            self.name2label[name] = len(self.name2label.keys())
        print(self.name2label)
        #images labels
        self.images,self.labels = self.load_csv('images.csv') #load_csv要么先创建images.csv，要么直接读取images.csv，
        if mode=='train':#train dataset 60% of ALL DATA
            self.images = self.images[:int(0.6*len(self.images))]
            self.labels = self.labels[:int(0.6*len(self.labels))]
        elif mode=='validation':#val dataset 60%-80% of ALL DATA
            self.images = self.images[int(0.6*len(self.images)):int(0.8*len(self.images))]
            self.labels = self.labels[int(0.6*len(self.labels)):int(0.8*len(self.labels))]
        else :#test dataset 80%-100% of ALL DATA
            self.images = self.images[int(0.8*len(self.images)):int(len(self.images))]
            self.labels = self.labels[int(0.8*len(self.labels)):int(len(self.labels))]
        # images[0]: D:\python pycharm learning\清华大佬课程\fisrt\pokemon\mewtwo\00000081.png
        # #labels[0]:2
        #images 还是图片的地址列表，需要__getitem__继续转换
    # image,label 不能把所有图片全部加载到内存，可能会爆内存
    def load_csv(self,filename):
        #filename 不存在：生成filename
        if not os.path.exists(os.path.join(self.root,filename)):
            images = []
            for name in self.name2label.keys():
                # .../pokemen/mewtwo/00001.png 加载进images列表
                # 实际上是加载每张图片的地址
                images += glob.glob(os.path.join(self.root, name, '*.png'))
                images += glob.glob(os.path.join(self.root, name, '*.jpg'))
                images += glob.glob(os.path.join(self.root, name, '*.jpeg'))

            print(len(images), images[0])
            random.shuffle(images)
            with open(os.path.join(self.root, filename), mode='w', newline='') as f:
                writer = csv.writer(f)
                for img in images:  # .....\bulbasaur\00000000.png
                    name = img.split(os.sep)[-2]  # 指：bulbasaur 图片真实类别
                    label = self.name2label[name]
                    # .....\bulbasaur\00000000.png ， 0
                    writer.writerow([img, label])
                print('writen into csv file:', filename)

        #filename 存在：直接读取filename
        images, labels = [], []
        with open(os.path.join(self.root, filename)) as f:
            reader = csv.reader(f)
            for row in reader:
                # '...pokemon\bulbasaur\00000000.png', 0
                img, label = row
                label = int(label)

                images.append(img)
                labels.append(label)

        assert len(images) == len(labels)
        return images,labels

    def __len__(self,x_hat):
        return len(self.images)

    def denormalize(self,x_hat):
        mean = [0.485, 0.456, 0.406]
        std = [0.229, 0.224, 0.225]

        #x_hat = (x-mean)
        #x = x_hat*std +mean
        #x:[c,h,w]
        #mean:[3]=>[3,1,1]
        mean = torch.tensor(mean).unsqueeze(1).unsqueeze(1)
        std = torch.tensor(std).unsqueeze(1).unsqueeze(1)

        x = x_hat * std + mean

        return x

    def __getitem__(self, idx):
        pass
        #idx~[0~len(images)]
        # self.iamges,self.labels
        #images[0]: D:\python pycharm learning\清华大佬课程\fisrt\pokemon\mewtwo\00000081.png
        # #labels[0]:2
        img, label = self.images[idx],self.labels[idx]
        tf = transforms.Compose([
            lambda x:Image.open(x).convert('RGB'),#string image => image data
            transforms.Resize((int(self.resize*1.25),int(self.resize*1.25))),#压缩到稍大
            transforms.RandomRotation(20),#图片旋转，增加图片的复杂度，但是又不会使网络太复杂
            transforms.CenterCrop(self.resize), #可能会有其他的底存在
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485,0.456,0.406],std=[0.229,0.224,0.225])
            #R mean:0.854,std:0.229
        ])
        img = tf(img)
        label = torch.tensor(label)

        return img,label

def main():
    import visdom      #  启动 python -m visdom.server，http://localhost:8097
    import time
    viz = visdom.Visdom()
    db = Pokemon('D:\python pycharm learning\清华大佬课程\\fisrt\pokemon',224,'train')
    x,y = next(iter(db))
    print('sample:',x.shape,y.shape,y)
    viz.images(db.denormalize(x), win='sample_x', opts=dict(title='sample_x'))
    #train,label= model.__getitem__(0)
    #print('train[0]:',train.shape)
    #print('label[0]:',label)

if __name__=='__main__':
        main()

D:\Anaconda\envs\study\python.exe "D:\python pycharm learning\清华大佬课程\fisrt\pokemon1.py"
Setting up a new session...
{'bulbasaur': 0, 'charmander': 1, 'mewtwo': 2, 'pikachu': 3, 'squirtle': 4}
1167 D:\python pycharm learning\清华大佬课程\fisrt\pokemon\bulbasaur\00000000.png
writen into csv file: images.csv
sample: torch.Size([3, 224, 224]) torch.Size([]) tensor(3)

Process finished with exit code 0

但是作业只能加载一张图片，我们要实现深度学习需要加载Batchsize张图片进行学习

2:使用DataLoader加载器实现批量加载，并使用visdom 显示

#数据集：皮卡丘：234，超梦：239，杰尼龟：223，小火龙：238，妙蛙种子：234
import torch
import  os,glob
import  random,csv
from torch.utils.data import Dataset ,DataLoader
from PIL import Image
from  torchvision import transforms

#自定义数据加载类
class Pokemon(Dataset):
    def __init__(self,root,resize,mode):            #root:文件所在目录，resize：图像分辨率调整一致,mode：当前类何功能
        super(Pokemon,self).__init__()

        self.root = root
        self.resize = resize

        self.name2label = {}                      #对每个加载的文件进行编码：'bulbasaur': 0, 'charmander': 1, 'mewtwo': 2, 'pikachu': 3, 'squirtle': 4
        for name in sorted(os.listdir((os.path.join(root)))):#对指定root中的文件进行排序
            if not os.path.isdir(os.path.join(root,name)):
                continue
            self.name2label[name] = len(self.name2label.keys())#keys返回列表当中的value,len计算列表长度
        print(self.name2label)#根据文件顺序，以idx:文件名，vlaue:0,1,2,3,4，生成列表
        #images labels
        self.images,self.labels = self.load_csv('images.csv') #load_csv要么先创建images.csv，要么直接读取images.csv，
        if mode=='train':#train dataset 60% of ALL DATA
            self.images = self.images[:int(0.6*len(self.images))]
            self.labels = self.labels[:int(0.6*len(self.labels))]
        elif mode=='validation':#val dataset 60%-80% of ALL DATA
            self.images = self.images[int(0.6*len(self.images)):int(0.8*len(self.images))]
            self.labels = self.labels[int(0.6*len(self.labels)):int(0.8*len(self.labels))]
        else :#test dataset 80%-100% of ALL DATA
            self.images = self.images[int(0.8*len(self.images)):int(len(self.images))]
            self.labels = self.labels[int(0.8*len(self.labels)):int(len(self.labels))]
        # images[0]: D:\python pycharm learning\清华大佬课程\fisrt\pokemon\mewtwo\00000081.png
        # #labels[0]:2
        #images 还是图片的地址列表，需要__getitem__继续转换
    # image,label 不能把所有图片全部加载到内存，可能会爆内存
    def load_csv(self,filename):#生成，读取filename文件
        #filename 不存在：生成filename
        if not os.path.exists(os.path.join(self.root,filename)):
            images = []
            for name in self.name2label.keys():
                # .../pokemen/mewtwo/00001.png 加载进images列表
                # 实际上是加载每张图片的地址
                images += glob.glob(os.path.join(self.root, name, '*.png'))
                images += glob.glob(os.path.join(self.root, name, '*.jpg'))
                images += glob.glob(os.path.join(self.root, name, '*.jpeg'))

            print(len(images), images[0])
            random.shuffle(images)
            with open(os.path.join(self.root, filename), mode='w', newline='') as f:
                writer = csv.writer(f)
                for img in images:  # .....\bulbasaur\00000000.png
                    name = img.split(os.sep)[-2]  # 指：bulbasaur 图片真实类别
                    label = self.name2label[name]#在name2label列表根据name找出对应的value:0,1...
                    # .....\bulbasaur\00000000.png ， 0
                    writer.writerow([img, label])
                print('writen into csv file:', filename)

        #filename 存在：直接读取filename
        images, labels = [], []
        with open(os.path.join(self.root, filename)) as f:
            reader = csv.reader(f)
            for row in reader:
                # '...pokemon\bulbasaur\00000000.png', 0
                img, label = row
                label = int(label)

                images.append(img)
                labels.append(label)

        assert len(images) == len(labels)
        return images,labels

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

    def denormalize(self,x_hat):#对已经进行规范化处理的totensor，去除规范化
        mean = [0.485, 0.456, 0.406]
        std = [0.229, 0.224, 0.225]

        #x_hat = (x-mean)/std
        #x = x_hat*std +mean
        #x:[c,h,w]
        #mean:[3]=>[3,1,1]
        mean = torch.tensor(mean).unsqueeze(1).unsqueeze(1)
        std = torch.tensor(std).unsqueeze(1).unsqueeze(1)

        x = x_hat * std + mean

        return x

    def __getitem__(self, idx):
        pass
        #idx~[0~len(images)]
        # self.iamges,self.labels
        #images[0]: D:\python pycharm learning\清华大佬课程\fisrt\pokemon\mewtwo\00000081.png
        # #labels[0]:2
        img, label = self.images[idx],self.labels[idx]
        tf = transforms.Compose([
            lambda x:Image.open(x).convert('RGB'),#string image => image data
            transforms.Resize((int(self.resize*1.25),int(self.resize*1.25))),#压缩到稍大
            transforms.RandomRotation(20),#图片旋转，增加图片的复杂度，但是又不会使网络太复杂
            transforms.CenterCrop(self.resize), #可能会有其他的底存在
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485,0.456,0.406],std=[0.229,0.224,0.225])
            #R mean:0.854,std:0.229
        ])
        img = tf(img)
        label = torch.tensor(label)
        #Pokemon类根据一个索引每次返回一个img（三位张量）,一个label（0维张量）
        return img,label #img,label打包成元组返回

def main():
    import visdom      #  启动 python -m visdom.server，http://localhost:8097
    import time
    viz = visdom.Visdom()
    db = Pokemon('D:\python pycharm learning\清华大佬课程\\fisrt\pokemon',64,'train')
    #x,y = next(iter(db))
    #print('sample:',x.shape,y.shape,y)
    #viz.images(db.denormalize(x), win='sample_x', opts=dict(title='sample_x'))
    #DataLoader加载器按batch_size打乱所有依次在内存当中按批次顺序加载每次批次，
    # 每个批次内含batch个Pokemon类返回的对象（元组，列表，字符串）
    loader = DataLoader(db,batch_size=64,shuffle=True)
    for x,y in loader:
        viz.images(db.denormalize(x),nrow=8,win='batch',opts=dict(title='batch'))
        viz.text(str(y.numpy()),win='label',opts=dict(title='batch-y'))

        time.sleep(10)

if __name__=='__main__':
        main()

3：如果数据集是安装二级目录分类，那可以使用

torchvision.datasets.ImageFolder（）

进行快速的加载，相当于用这一条函数替代了前面的 Pokemon class,非常简洁实用

import visdom
import time
import torchvision
from torch.utils.data import  DataLoader
from torchvision import transforms

def main():
    viz = visdom.Visdom()
    tf = transforms.Compose([
        transforms.Resize((64,64)),
        transforms.ToTensor(),
    ])
    db = torchvision.datasets.ImageFolder(root='D:\python pycharm learning\清华大佬课程\\fisrt\pokemon',transform=tf)
    loader = DataLoader(db,batch_size=32,shuffle=True)

    for x, y in loader:
        viz.images( x, nrow=8, win='batch', opts=dict(title='batch'))
        viz.text(str(y.numpy()), win='label', opts=dict(title='batch-y'))

        time.sleep(10)

if __name__ == '__main__':
    main()

在visdom显示结果：

如果要知道文件是怎么分类的可以在main函数中使用如下代码

 print(db.class_to_idx)

如下：

from torch.utils.data import DataLoader
import torchvision
from  torchvision import transforms
import visdom
import time
def main():
    viz = visdom.Visdom()
    tf = transforms.Compose([
        transforms.Resize((64, 64)),
        transforms.ToTensor(),
    ])
    db = torchvision.datasets.ImageFolder(root='D:\python pycharm learning\清华大佬课程\\fisrt\pokemon', transform=tf)
    loader = DataLoader(db, batch_size=32, shuffle=True, num_workers=8)  # 8个线程同时加速
    print(db.class_to_idx)
    for x, y in loader:
        viz.images(x, nrow=8, win='batch', opts=dict(title='batch'))
        viz.text(str(y.numpy()), win='label', opts=dict(title='batch-y'))

        time.sleep(10)

if __name__ == '__main__':
    main()

{'bulbasaur': 0, 'charmander': 1, 'mewtwo': 2, 'pikachu': 3, 'squirtle': 4}

使用torchvision.datasets.ImageFolder（）函数只有在pokemon数据集的二级分类各子文件下的picture准确对应文件名的情况下才可以使用，否则会出现图片的label与图片对应错误的情况发生