用于训练 GoogLeNet 模型的快速数据生成器

Question

我尝试在 Keras 中从头开始训练 GoogLeNet。我构建了网络架构，并准备好进行训练。使用辅助输出训练 GoogLeNet，数据生成器应具有三个输出标签。我使用 tf.keras.utils.Sequence 编写自定义数据生成器。

我的自定义生成器是：

from skimage.transform import resize
from skimage.io import imread
import numpy as np
import math
from tensorflow.keras.utils import Sequence


class GoogLeNetDatasetGenerator(Sequence):

    def __init__(self, X_train_path, y_train, batch_size):
        """
        Initialize the GoogLeNet dataset generator.
        :param X_train_path: Path of train images
        :param y_train: Labels of train images
        :param batch_size:
        """

        self.X_train_path = X_train_path
        self.y_train = y_train
    self.batch_size = batch_size

        self.indexes = np.arange(len(self.X_train_path))
        np.random.shuffle(self.indexes)

    def __len__(self):
        """
        Denotes the number of batches per epoch
        :return:
        """

        return math.ceil(len(self.X_train_path) / self.batch_size)

    def __getitem__(self, index):
        """
        Get batch indexes from shuffled indexes
        :param index:
        :return:
        """

        indexes = self.indexes[index * self.batch_size:(index + 1) * self.batch_size]
        X_batch_names = [self.X_train_path[i] for i in indexes]
        y_batch_naive = self.y_train[indexes]

        X_batch = np.array([resize(imread(file_name), (224, 224) for file_name in X_batch_names],
                       dtype='float32')
        y_batch = [y_batch_naive, y_batch_naive, y_batch_naive]

        return X_batch, y_batch

    def on_epoch_end(self):
        """
        Updates indexes after each epoch
        :return:
        """

        self.indexes = np.arange(len(self.X_train_path))
        np.random.shuffle(self.indexes)

另外，我使用以下代码编译和训练模型：

# Compile model
model.compile(loss=[CategoricalCrossentropy(), CategoricalCrossentropy(), CategoricalCrossentropy()],
                  loss_weights=[1, 0.3, 0.3], optimizer='adam',
                  metrics=['accuracy'])

# Train model
history = model.fit(train_dataset, validation_data=test_dataset, epochs=100)

使用 GPU 版本的 TensorFlow 时，在数据生成器中加载图像非常耗时。它导致训练过程缓慢。有什么建议或其他解决方案可以加快加载数据的速度吗？

聚苯乙烯 我搜索 StackOverflow 问题，例如 this，但我没有找到任何想法。

Answer 1

我找到了另一个更快的解决方案。您可以使用tf.data.Dataset。在第一步中，我列出了所有训练图像目录。使用map方法帮助我读取图像并正确配置相应的标签。这是我的示例代码，用于加载带有三元标签的图像。

image_filenames = tf.constant(image_list)

slices_dataset = tf.data.Dataset.from_tensor_slices(image_filenames)
slices_labels = tf.data.Dataset.from_tensor_slices(label_list)

image_dataset = slices_dataset.map(map_func=process_image)
label_dataset = slices_labels.map(map_func=process_label)

x_dataset = image_dataset.shuffle(buffer_size=Cfg.BUFFER_SIZE, seed=0).\
    batch(batch_size=Cfg.BATCH_SIZE)
y_dataset = label_dataset.shuffle(buffer_size=Cfg.BUFFER_SIZE, seed=0).\
    batch(batch_size=Cfg.BATCH_SIZE)

dataset = tf.data.Dataset.zip((x_dataset, y_dataset))