Как эффективно загружать большие тренировочные данные, слишком большие для оперативной памяти, для обучения в тензорном потоке?

Question

У меня есть около 10 ГБ обучающих данных в формате массива numpy. Тем не менее, моя оперативная память недостаточно велика для одновременной загрузки данных и модели tenorflow 2.0. Я провел много исследований tf.data, tf.TFRecords и генераторов, но я застрял в фактической реализации тренинга l oop. В настоящее время у меня есть следующее стандартное обучение l oop, когда я использую загружаемое подмножество данных, чтобы проверить, все ли работает.

Варианты, которые я рассмотрел:

1) Разделение 10 ГБ файлов numpy на большое количество файлов tf.TFRecords (перемешанных), каждый из которых имеет размер BATCH_SIZE, а затем загрузка их в каждый train_step и epoch

2) Разделение 10 ГБ файлов numpy на более мелкие файлы по 1-2 ГБ каждый (в произвольном порядке), а затем последовательное обучение модели путем загрузки меньших файлов каждую эпоху последовательно

Есть ли лучший метод для этого или лучший вариант, который я не рассматриваю? Я полагаю, что это должно быть тривиальной проблемой, но я не могу найти каких-либо хороших решений в Интернете.

if __name__ == '__main__':
    # get the original_dataset
    #train_dataset, valid_dataset
    train_dataset = tf.data.Dataset.from_tensor_slices((mtr, labels))
    train_dataset = train_dataset.shuffle(buffer_size=mtr.shape[0]).batch(BATCH_SIZE)

    valid_dataset = tf.data.Dataset.from_tensor_slices((val_mtr, val_labels))
    valid_dataset = valid_dataset.batch(BATCH_SIZE)

    # create model
    model = get_model()

    # define loss and optimizer 
    loss_object = tf.keras.losses.SparseCategoricalCrossentropy()
    optimizer = tf.keras.optimizers.Adam(learning_rate=1e-04, beta_1=0.9, beta_2=0.999, epsilon=1e-08,)

    train_loss = tf.keras.metrics.Mean(name='train_loss')
    train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='train_accuracy')

    valid_loss = tf.keras.metrics.Mean(name='valid_loss')
    valid_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='valid_accuracy')

    @tf.function
    def train_step(images, labels):
        with tf.GradientTape() as tape:
            predictions = model(images, training=True)
            loss = loss_object(y_true=labels, y_pred=predictions)
        gradients = tape.gradient(loss, model.trainable_variables)
        optimizer.apply_gradients(grads_and_vars=zip(gradients, model.trainable_variables))
        train_loss(loss)
        train_accuracy(labels, predictions)

    @tf.function
    def valid_step(images, labels):
        predictions = model(images, training=False)
        v_loss = loss_object(labels, predictions)
        valid_loss(v_loss)
        valid_accuracy(labels, predictions)

    # start training
    for epoch in range(EPOCHS):
        train_loss.reset_states()
        train_accuracy.reset_states()
        valid_loss.reset_states()
        valid_accuracy.reset_states()
        step = 0
        for imgs, lbls in train_dataset:
            step += 1
            train_step(imgs, lbls)
            print("Epoch: {}/{}, step: {}/{}, loss: {:.5f}, accuracy: {:.5f}".format(epoch + 1, EPOCHS, step, math.ceil(mtr.shape[0] / BATCH_SIZE),
                                                                                     train_loss.result(), train_accuracy.result()))
        for valid_images, valid_labels in valid_dataset:
            valid_step(valid_images, valid_labels)

        print("Epoch: {}/{}, train loss: {:.5f}, train accuracy: {:.5f}, "
              "valid loss: {:.5f}, valid accuracy: {:.5f}".format(epoch + 1, EPOCHS, train_loss.result(),
                                                                  train_accuracy.result(), valid_loss.result(), valid_accuracy.result()))

    model.save_weights(filepath=save_model_dir, save_format='tf')

Основываясь на ответе ниже, это сработало:

def load_data(filename: tf.Tensor):
     coll = pickle.load(open("./TrainData2016/%s" %(str(filename.numpy().decode('utf-8'))), "rb")) #coll[0] = data, coll[1] = label
     return tf.convert_to_tensor(coll[0],dtype=tf.float32), tf.convert_to_tensor(coll[1],dtype=tf.int32)

filenames = os.listdir("./TrainData2016") # You can just shuffle this array and not use tf.data.Dataset.shuffle()
train_dataset = tf.data.Dataset.from_tensor_slices(filenames)
train_dataset = train_dataset.map(lambda x: tf.py_function(func=load_data, inp=[x], Tout=(tf.float32, tf.int32)))

Answer 1

Вы пытались использовать функцию map из tf.data.Dataset? Я имею в виду, что вместо загрузки изображений вы можете иметь массив с именами файлов, посмотрите только на приведенный ниже пример:

def load_img(filename):
   # Read image file
   img = tf.io.read_file(....)
   img = tf.io.decode_png(img) # have a look at tf.io.decode_*

   # Do some normalization or transformations if needed on img 

   # Add label
   label = associated_image_label() # associate the needed lable to your image

   return img, label

filenames = os.listdir("./images_dir") # You can just shuffle this array and not use tf.data.Dataset.shuffle()

train_dataset = tf.data.Dataset.from_tensor_slices(filenames)
train_dataset = train_dataset.map(load_img)

# Create batches and prefetch

Таким образом, вы будете загружать только серию изображений только для вашего train_step *