基于Pytorch、Keras、Tensorflow的图片数据生成器搭建

墨蓝 2022-09-09 12:44 277阅读 0赞

## 1. 前言 ##

作为一个对三种深度学习框架（**Tensorflow，Keras，Pytorch**）刚刚完成入门学习的菜鸟，在实战的过程中，遇到了一些菜鸟常见问题，即**图片数据加载与预处理**。在刚刚接触深度学习的时候，并不懂**数据生成器**的作用，后面随着数据量的增大，终于感受到了数据生成器的魅力了。我相信有部分新手和我一样，在训练的过程中，喜欢将所有图片数据读取到内存后，然后进行图片的预处理。那么会出现一个很大的问题 - **Out of memory**! 你可能会觉得是硬件的问题，实际上还是我们太菜了，所以这篇文章我想简单的说一下**三种深度学习框架**如何搭建数据生成器。

## 2． Pytorch下的图片数据生成器 ##

首先还是介绍下当下最火热的Pytorch深度学习框架，本来坚持Keras+Tensorflow的我，还是逃不过“**真香**”警告呀。

![efced2ee332399291f3ec44f68d95ef7.png][]

不得不说，Pytorch真好用，看Pytorch代码总有种**赏心悦目**的感觉，Pytorch中大量使用**类定义**，使得代码看起来更加清晰明了，不论从模型定义到数据生成，都是一气呵成。话不多说，来看一下Pytorch关于数据迭代器的生成，下面的例子都是以最常见的猫狗数据来进行定义的。这个数据集的结构如下图所示。

### 2.1 图片数据集结构 ###

Dogs\_cats\_data文件下包含三个子文件夹，分别放置**训练集，验证集和测试集**。

![f66b749430b0acc4081a8720033501ce.png][]

我们打开training\_set文件夹，有两个子文件夹dogs和cats，类似地，test和validation里面也有**这两个子文件夹**。

![364c5f6c9c4b4dd751c178e8dcea2b05.png][]

Cats文件夹下和dogs文件夹下的文件如下图所示。命名格式为“cat.xxx.jpg”和“dog.xxx.jpg”。

![3cf1f703f8f9ec84e0947fdc6499973d.png][] cats子文件夹的内容 ![188a7d51f79d0266132454807c1c850d.png][] dogs子文件夹的内容

### 2.2 Pytorch图片数据生成器解读 ###

介绍完了数据集的大致情况，我们就来写一个pytorch关于猫狗数据生成器。这里主要分别两步：

1.  写一个**torch.utils.data.Dataset的子类**，继承父类Dataset的一些属性和方法。一般来说，**修改**该类的三个**专有方法**（**init**、**getitem**、**len**）就可以完成一个新子类的生成。
2.  使用**torch.utils.data.DataLoader**结合上面定义的**Dataset子类**来创建一个数据生成器

为了**不占用过多内存**，我们需要将图片的所有**地址**（并不是所有数字化图片）加载到内存中，需要多少图片数据的时候就从**内存中解析多少图片地址**，这样有效且合理地使用内存，也不会耽误时间。

下面为数据生成器代码，定义了一个Dataset的子类Dogcat，**实例化类**的时候，需要**输入**

1.  图片的**根目录root**
2.  是否对图片数据进行转换（**标准化、resize、张量化**）
3.  Train还是val

from torchvision import transforms as T
    import os
    from torch.utils.data import Dataset
    from torch.utils.data import DataLoader
    from PIL import Image
    import os
    import numpy as np
    import pandas as pd
    import glob
    
    # 搭建一个猫狗数据集生成器
    class DogCat(Dataset):
        def __init__(self, root, transforms=None, train=True, val=False):
            """
            get images and execute transforms.
            """
            self.val = val
            imgs=glob.glob(os.path.join(root,'*/*/*.jpg'))
            imgs = sorted(imgs, key=lambda x: x.split('.')[-2]) #对图片进行排序，按照图片的索引
            self.imgs = imgs
            if transforms is None:
                # normalize
                normalize = T.Normalize(mean = [0.485, 0.456, 0.406],
                                         std = [0.229, 0.224, 0.225])
                # trainset and valset have different data transform
                # trainset need data augmentation but valset don't.
                # valset
    
                if self.val:
                    # T是组合的意思，将下面的操作组合起来
                    self.transforms = T.Compose([
                        T.Resize(224),
                        T.CenterCrop(224),
                        T.ToTensor(),
                        normalize
                    ])
                # trainset
                else:
                    self.transforms = T.Compose([
                        T.Resize(256),
                        T.RandomResizedCrop(224), #随机大小裁剪
                        T.RandomHorizontalFlip(),
                        T.ToTensor(),
                        normalize
                    ])
    
        def __getitem__(self, index):
            """
            return data and label
            """
            img_path = self.imgs[index]
            label = 1 if 'dog' in img_path.split('/')[-1] else 0 #根据图片名称给与label，这里是二分类问题
            data = Image.open(img_path)
            data = self.transforms(data)
            return data, label
    
        def __len__(self):
            """
            return images size.
            """
            return len(self.imgs)
    
    if __name__ == "__main__":
        train_dataset = DogCat('E:\Image_Dataset\dogs_cats_data', train=True)
        data_loader = DataLoader(dataset=train_dataset, batch_size=10, shuffle=True)
        for data,label in data_loader:
            print('label',label)
            break

这里，**getitem是类专有方法，给予了实例化类的索引方法**，即可以**通过索引访问实例化的Dataset子类（DogCat类）的元素**。比如对上述代码中，我可以使用**train\_dataset\[0\]** （这里的train\_dataset是实例化的DogCat类） 来获取**第一张图片**的信息。在\_\_getitem\_\_方法中，根据图片的索引，对相应位置的图片地址进行解析，读取图片，对其进行**transforms**，将其转换为图片数据并返回。这样一来，就完成了图片数据的处理。

再将**实例化的Dataset**作为**Dataloader输入**的时候，根据设置的batch\_size，不断地使用\_\_getitem\_\_方法获取对应**batch\_size数量**图片数据后，作为Dataloader返回值。

实践过程中，使用**for循环**不断地读取每批处理完后的图片数据，这样可以**有效地缓解内存不足问题**。

## 3 Keras下的图片数据生成器 ##

当初在学《Python深度学习》的时候，在猫狗分类的时候，使用了一个**数据迭代器**，代码是这样的，

from keras.preprocessing.image import ImageDataGenerator
    train_datagen = ImageDataGenerator(rescale=1./255)
    test_datagen = ImageDataGenerator(rescale=1./255)

上面这个是实例化一个图片数据生成器，这里包含了和pytorch类似的transforms，但是注意到我们**这里没有数据地址索引**，所以，我们需要使用实例化类的**flow\_from\_directory方法**，将图片的地址**train\_dir和target\_size**等作为输入，这样就**完成了一个数据生成器**的生成了，如下

train_generator = train_datagen.flow_from_directory(
                                 train_dir,
                                 target_size=(150, 150),
                                 batch_size=20,
                                 class_mode='binary')
    validation_generator = test_datagen.flow_from_directory(
                                 validation_dir,
                                 target_size=(150, 150),
                                 batch_size=20,
                                 class_mode='binary')

在使用的时候，使用循环就可以不断地生成数据了，每次循环生成batch\_size个数据了。下面代码，每次生成20个数据。

for data_batch, labels_batch in train_generator:
        print('data batch shape:', data_batch.shape)
        print('labels batch shape:', labels_batch.shape)
        break

训练过程中，使用fit\_generator就可以了

history = model.fit_generator(
                        train_generator,
                        steps_per_epoch=100,
                        epochs=30,
                        validation_data=validation_generator,
                        validation_steps=50)

但是按照上述方法生成生成器，最大的**弊端**就是**不灵活**，我**只要输入**一个地址端口train\_dir,他就生成了一个数据生成器，这其实是一个**高度封装**的函数，那么输入的**地址结构信息**一定要**固定**，想用自己的数据集做点事，就很不灵活了。

参考以下[这篇文章][Link 1]，我们也可以**自定义一个keras的数据生成器**。这里的定义就有点像Pytorch中数据生成器的定义，但是这里是继承**Sequence**生成一个**新的子类**，**不像Pytorch是继承Dataset的。**

Keras Sequence方法用于拟合一个数据序列，每一个Sequence必须提供\_\_getitem\_\_和\_\_len\_\_方法，这跟Torch的Dataset模块类似。Sequence是进行多进程处理的更安全的方法，这种结构保证网络在**每个时期每个样本只训练一次**，这与生成器不同。

这里我们仍然使用上述猫狗数据，来新建一个Keras数据生成器。代码如下：

from skimage.io import imread
    from skimage.transform import resize
    import numpy as np
    from keras.utils import Sequence
    import cv2
    import glob
    import os
    
    class Dogs_Cats_DataGenerator(Sequence):
        """
        基于Sequence的自定义Keras数据生成器
        """
        def __init__(self, filepath, batch_size=8, imgshape=(256, 472),
                     n_channels=3, n_classes=13, shuffle=True):
            """ 初始化方法
            :param filepath :数据文件地址
            :param batch_size: batch size
            :param imgshape: 图像大小
            :param n_channels: 图像通道
            :param n_classes: 标签类别
            :param shuffle: 每一个epoch后是否打乱数据
            """
            self.filepath=filepath
            self.pathlist=glob.glob(os.path.join(self.filepath,'*/*.jpg'))
            self.batch_size = batch_size
            self.imgshape = imgshape
            self.n_channels = n_channels
            self.n_classes = n_classes
            self.shuffle = shuffle
            # 每个epoch之后更新索引
            self.on_epoch_end()
            # 文件地址list
    
        def __getitem__(self, index):
            """生成每一批次训练数据
            :param index: 批次索引
            :return: 训练图像和标签
            """
            # 生成批次索引
            indexes = self.indexes[index * self.batch_size:(index + 1) * self.batch_size]
            # 索引列表
            batch_pathlist = [self.pathlist[k] for k in indexes]
            # 生成数据
            X = self._generate_X(batch_pathlist)
            y = self._generate_y(batch_pathlist)
            return X, y
    
        def __len__(self):
            """每个epoch下的批次数量,也就是每个epoch的iteration
            """
            return int(np.floor(len(self.pathlist) / self.batch_size))
    
        def _generate_X(self, batch_pathlist):
            """生成每一批次的图像
            :param list_IDs_temp: 批次数据索引列表
            :return: 一个批次的图像
            """
            # 初始化
            X = np.empty((self.batch_size, *self.imgshape, self.n_channels))
            # 生成数据
            for i, path in enumerate(batch_pathlist):
                # 存储一个批次
                X[i,] = self._load_image(path)
            return X
    
        def _generate_y(self, batch_pathlist):
            """生成每一批次的标签
            :param list_IDs_temp: 批次数据索引列表
            :return: 一个批次的标签
            """
            y = np.empty((self.batch_size, ), dtype=int)
            # Generate data
            for i, path in enumerate(batch_pathlist):
                # Store sample
                y[i,]= 1 if 'dog' in path.split('/')[-1] else 0
            return y
    
        def on_epoch_end(self):
            """每个epoch之后更新索引
            """
            self.indexes = np.arange(len(self.pathlist))
            if self.shuffle == True:
                np.random.shuffle(self.indexes)
    
        def _load_image(self, image_path):
            """cv2读取图像
            """
            # img = cv2.imread(image_path)
            img = cv2.imdecode(np.fromfile(image_path, dtype=np.uint8), cv2.IMREAD_COLOR)
            w, h, _ = img.shape
            if w>h:
                img = np.rot90(img)
            img = cv2.resize(img, self.imgshape)
            return img
    
    
    if __name__=='__main__':
        # Parameters
        params = {'batch_size': 10,
                'n_classes': 2,
                'n_channels': 3,
                'shuffle': True,
                'imgshape':(224,224)}
        train_filepath=r'E:\Image_Dataset\dogs_cats_data\training_set'
        val_filepath=r'E:\Image_Dataset\dogs_cats_data\validation'
        # Generators
        training_generator = Dogs_Cats_DataGenerator(train_filepath, **params)
        validation_generator =  Dogs_Cats_DataGenerator(val_filepath, **params)
    
        for x,y in training_generator:
            print(y)
            break

到这里，我们就完成了对**keras生成器**的搭建了，其实和Pytorch的生成器还是很像，只不过**继承的父类不一样**罢了。

## 4 Tensorflow下的图片数据生成器 ##

Tensorflow在我心中是老大哥，也是让我真正感受到**计算图美妙**的深度学习框架，虽然不如某torch简单易用（这确实是事实），如果不是大家都用pytorch，导致一些开源代码都是基于pytorch的，只用tensorflow和keras就能完全满足我的深度学习进程了。好了，话不多说，下面就用我的“**前任**”深度学习框架来创造一个图片数据生成器。

在Tensorflow中，使用**tf.data.Dataset**直接就可以**实例化一个图片数据生成器**。我将创建步骤分为以下几步：

1.  根据读取数据的类型实例化一个Dataset
2.  使用map函数对实例化的Dataset进行数据的预处理
3.  创建一个数据迭代器，根据自身需求，创建不同类型的迭代器
4.  get\_next()函数依次获取数据

讲了大致的思路后，接着我们开始定义**实例化Dataset**，当然还是以猫狗数据集为例，我们的思路仍然是将猫狗的数据的**地址**作为数据集的**输入**，**需要多少个图片数据的时候**，就**解析多少个图片地址**。代码如下

data_dir = config.dataset_dir
        data_root = pathlib.Path(data_dir)
        all_image_path = list(data_root.glob('*/*'))
        all_image_path = [str(path) for path in all_image_path]
        label_names = sorted(item.name for item in data_root.glob('*/'))
        # # dict: name->index ,如{'cats': 0, 'dogs': 1}
        label_to_index = dict((index, name) for name, index in enumerate(label_names))
        print('train dic index',label_to_index)
        # # get all images' labels
        all_image_label = [label_to_index[pathlib.Path(p).parent.name] for p in all_image_path]
        # # load dataset and preprocess images（the preprocess function in map  can preprocess all previous pictures)
        image_dataset = tf.data.Dataset.from_tensor_slices(all_image_path).map(load_and_preprocess_image_for_train)
        # print(image_dataset)
        label_dataset = tf.data.Dataset.from_tensor_slices(all_image_label)
        dataset = tf.data.Dataset.zip((image_dataset, label_dataset))

上面分别对image和label**各实例化**一个数据集，**再组合**成一个数据集。我们注意到，这里面使用了**map函数**，用来**解析batch\_size个图片地址**，并**加以预处理**。

定义如下：

def load_and_preprocess_image_for_train(img_path):
        # read pictures
        img_raw = tf.io.read_file(img_path)
        # decode pictures
        img_tensor = tf.image.decode_jpeg(img_raw, channels=channels)
        # resize
        img_tensor = tf.image.resize(img_tensor, [image_height, image_width])
        #tf.cast() function is a type conversion function that converts the data format of x into dtype
        img_tensor = tf.cast(img_tensor, tf.float32)
        # normalization
        img_tensor = img_tensor / 255.0
        # flip left or right
        img_tensor = tf.image.random_flip_left_right(img_tensor)
        # change color randomly
        img_tensor=change_color(img_tensor,NUM=np.random.randint(4))
        return img_tensor

定义完了相关的数据集后，还有一步，就是设定：

1.  batch\_size,
2.  shuffle
3.  epoch

只需一步就到位了

train_dataset =dataset.shuffle(buffer_size=10000).batch(batch_size=config.BATCH_SIZE).repeat(config.Epoches)

到此就完成了Tensorflow数据集的定义和预处理，接着我们需要**定义迭代器**，如下：

iteration_train = train_dataset.make_initializable_iterator()
     train_image_batch, train_label_batch = iteration_train.get_next()

然后使用**sess.run()来获取每组生成数据**，用来参与训练就好了

for i in range(config.training_step):
                images,labels=sess.run([train_image_batch,train_label_batch])
                _,train_loss,acc=sess.run([train_op,losses,accuracy],feed_dict={image:images,
                                                                   label:labels,
                                                                   is_training:True
                                                                   })
    
                train_loss_list.append(train_loss)
                train_acc_list.append(acc)

到这里，使用Tensorflow来创建数据迭代器就完成了，没有多少花哨的地方，简单又直接。

## 5.总结 ##

到这里，共计耗费我一天的时间，终于将**基于三个深度学习框架**的**图片数据迭代器**（又可以叫**数据生成器**）总结写完了，目的是为了我以后回头看的时候更加方便，也为了可以让广大网友们**学习、借鉴**，甚至**批评指正**。

**建议阅读：**

**[高考失利之后，属于我的大学本科四年][Link 2]**

[**【资源分享】对于时间序列，你所能做的一切.**][.]

[**【时空序列预测第一篇】什么是时空序列问题？这类问题主要应用了哪些模型？主要应用在哪些领域？**][Link 3]

[**【AI蜗牛车出品】手把手AI项目、时空序列、时间序列、白话机器学习、pytorch修炼**][AI_AI_pytorch]

**公众号：****AI蜗牛车**

保持谦逊、保持自律、保持进步

![633c05948f0e135579401d72d323d780.png][]

个人微信

备注：昵称+学校/公司+方向

如果没有备注不拉群！

拉你进AI蜗牛车交流群

![56b5940ca71625803de4fe07447de858.png][]

[efced2ee332399291f3ec44f68d95ef7.png]: /images/20220829/70327aca5c624b5bba6666db104fbc90.png
[f66b749430b0acc4081a8720033501ce.png]: /images/20220829/210fa0be99ef48b69d8f84c58ba72640.png
[364c5f6c9c4b4dd751c178e8dcea2b05.png]: /images/20220829/20f0aa93c5884ff8842845305be72088.png
[3cf1f703f8f9ec84e0947fdc6499973d.png]: https://img-blog.csdnimg.cn/img_convert/3cf1f703f8f9ec84e0947fdc6499973d.png
[188a7d51f79d0266132454807c1c850d.png]: /images/20220829/b96d386474174831934df6ac506cf103.png
[Link 1]: https://mp.weixin.qq.com/s?__biz=MzI4ODY2NjYzMQ%3D%3D&idx=1&mid=2247486630&scene=21&sn=739894a26e56a8d8fcb2f7824396586a#wechat_redirect
[Link 2]: http://mp.weixin.qq.com/s?__biz=MzA4ODUxNjUzMQ%3D%3D&chksm=9029a749a75e2e5fde587cf68ab729deaa3a6c2ec952e288af747c1591df1416c24c9a9bcd1a&idx=1&mid=2247488917&scene=21&sn=6b0a24a219f861bd146366b25f710beb#wechat_redirect
[.]: http://mp.weixin.qq.com/s?__biz=MzA4ODUxNjUzMQ%3D%3D&chksm=9029b4cda75e3ddb7cb31a55ed528f386b3adcba3a077d6344f4c868f19bac1185a58cb262f8&idx=1&mid=2247484945&scene=21&sn=d3b302099848ad3549a395750b9da2bf#wechat_redirect
[Link 3]: http://mp.weixin.qq.com/s?__biz=MzA4ODUxNjUzMQ%3D%3D&chksm=9029b7fba75e3eed7d9e647994cae086ee483c1e55d386a92d385e3648f0ce715097f2c33a62&idx=1&mid=2247484711&scene=21&sn=c4187946d58e277a45512d3146d4f062#wechat_redirect
[AI_AI_pytorch]: http://mp.weixin.qq.com/s?__biz=MzA4ODUxNjUzMQ%3D%3D&chksm=9029afcba75e26dd76bd9263b4cf22a9ef3e1538f90008bc64867f52a7b2705bdeaf526d4b37&idx=1&mid=2247490839&scene=21&sn=d9d0a9fad6efb225a92deafd80a2838d#wechat_redirect
[633c05948f0e135579401d72d323d780.png]: /images/20220829/d0c91f4ed4dc4ad183e90950d5821d71.png
[56b5940ca71625803de4fe07447de858.png]: /images/20220829/a4a3d2f44d68438a8cfb6691d48644df.png