Ожидаемая производительность обучения модели tf.keras.Sequential с использованием model.fit, model.fit_generator и model.train_on_batch

Я использую Keras с бекендом Tensorflow для обучения простейшей 1D КСН для обнаружения определённых событий на основе данных с датчиков. Хотя данные с десятками миллионов выборок легко помещаются в оперативную память в виде одномерного массива с плавающей запятой, очевидно, что для хранения данных в виде массива N x inputDim, который можно передать в model.fit для обучения, требуется огромное количество памяти. Хотя я могу использовать model.fit_generator или model.train_on_batch для генерации необходимых мини-батчей на лету, по непонятной причине я наблюдаю огромную разницу в производительности между model.fit и model.fit_generator & model.train_on_batch, даже несмотря на то, что всё хранится в памяти, а генерация мини-батчей быстро происходит, так как по сути состоит только в изменении формы данных. Поэтому я задаюсь вопросом, делаю ли я что-то ужасно неправильно, или такая разница в производительности вполне нормальна. Я использую версию Tensorflow 2.0 для процессора с процессором Intel Core i7 на 3.2 ГГц (4 ядра с поддержкой многопоточности) и Python 3.6.3 на Mac OS X Mojave.

Короче говоря, я создал простой скрипт на Python для воспроизведения проблемы, и он показывает, что при размере батча 64 требуется 407 секунд для выполнения 10 эпох с model.fit, 1852 секунды с model.fit_generator и 1985 секунд с model.train_on_batch. Нагрузки на процессор составляют ~220%, ~130% и ~120% соответственно, и особенно странно, что model.fit_generator и model.train_on_batch практически на одном уровне, в то время как model.fit_generator должен быть способен к параллельному созданию мини-батчей, а model.train_on_batch, безусловно, этого не делает. То есть, model.fit (с огромными требованиями к памяти) превосходит другие варианты решения с легко управляемыми требованиями к памяти в четыре раза. Очевидно, что нагрузки на процессор увеличиваются, а общее время обучения уменьшается при увеличении размера батча, но model.fit всегда быстрее с запасом минимум в два раза до размера батча 8096. В этом случае, model.fit занимает 99 секунд для выполнения 10 эпох с нагрузкой на процессор около ~860% (или практически всё, что у меня есть), model.fit_generator занимает 179 секунд с нагрузкой на процессор около ~700%, а model.train_on_batch занимает 198 секунд с нагрузкой на процессор около ~680%.

Является ли такое поведение нормальным (когда нет GPU) или что можно/нужно сделать, чтобы увеличить вычислительную производительность менее требовательных к памяти вариантов с разумными размерами батчей? В частности, model.fit_generator не обеспечивает достойной производительности. Кажется, что нет такой опции, которая позволяла бы разбить все данные на управляемые фрагменты, а затем запускать model.fit в итерационном режиме с постоянно изменяющимися тренировочными данными.

Обратите внимание, что предоставленный дартовый скрипт — это просто то, что подразумевает название, и количество данных было уменьшено, чтобы сделать все три варианта выполнимыми. Используемая модель, однако, схожа с той, которую я на самом деле использую (чтобы предоставить реалистичную ситуацию).

from tqdm       import tqdm

import numpy as np
import tensorflow as tf

import time
import sys
import argparse

inputData    = None
outputData   = None
batchIndices = None
opts         = None

class DataGenerator(tf.keras.utils.Sequence):

    global inputData
    global outputData
    global batchIndices

    'Генерация данных для Keras'
    def __init__(self, batchSize, shuffle):
        'Инициализация'
        self.batchIndices = batchIndices
        self.batchSize    = batchSize
        self.shuffle      = shuffle
        self.on_epoch_end()

    def __len__(self):
        'Обозначает количество батчей за эпоху'
        return int( np.floor( inputData.size / self.batchSize ) )

    def __getitem__(self, index):
        'Генерирует один батч данных'

        # Генерация данных
        X, y = self.__data_generation(self.indexes[index*self.batchSize:(index+1)*self.batchSize])

        return X, y

    def on_epoch_end(self):
        'Обновляет индексы после каждой эпохи'
        self.indexes = np.arange(inputData.size)
        if self.shuffle == True:
            np.random.shuffle(self.indexes)

    def __data_generation(self, INDX):
        'Генерирует данные, содержащие batch_size выборок'

        # Генерация данных
        X = np.expand_dims( inputData[ np.mod( batchIndices + np.reshape(INDX,(INDX.size,1)) , inputData.size ) ], axis=2)
        y = outputData[INDX,:] 

        return X, y

def main( ):

    global inputData
    global outputData
    global batchIndices
    global opts

    # Генерация данных

    print(' ')
    print('Генерация данных...')

    np.random.seed(0) # Для воспроизводимых результатов

    inputDim  = int(104)                      # Входное измерение
    outputDim = int(  2)                      # Выходное измерение
    N         = int(1049344)                  # Общее количество выборок
    M         = int(5e4)                      # Количество аномалий
    trainINDX = np.arange(N, dtype=np.uint32)

    inputData = np.sin(trainINDX) + np.random.normal(loc=0.0, scale=0.20, size=N) # Исходные данные, хранящиеся в одном массиве

    anomalyLocations = np.random.choice(N, M, replace=False)

    inputData[anomalyLocations] += 0.5

    outputData = np.zeros((N,outputDim)) # Одно-горячий закодированный целевой массив без единиц

    for i in range(N):
        if( np.any( np.logical_and( anomalyLocations >= i, anomalyLocations < np.mod(i+inputDim,N) ) ) ): 
            outputData[i,1] = 1 # устанавливаем класс #2 в единицу, если имеется хотя бы одна аномалия в диапазоне [i,i+inputDim)
        else:
            outputData[i,0] = 1 # устанавливаем класс #1 в единицу, если нет аномалий в диапазоне [i,i+inputDim)

    print('...завершено')
    print(' ')

    # Создание модели для обнаружения аномалий

    model = tf.keras.Sequential([
        tf.keras.layers.Conv1D(filters=24, kernel_size=9, strides=1, padding='valid', dilation_rate=1, activation='relu', use_bias=True, kernel_initializer="glorot_uniform", bias_initializer="zeros", input_shape=(inputDim,1)),
        tf.keras.layers.MaxPooling1D(pool_size=4, strides=None, padding='valid'),
        tf.keras.layers.Flatten(),
        tf.keras.layers.Dense(20, activation='relu', use_bias=True),
        tf.keras.layers.Dense(outputDim, activation='softmax')
    ])

    model.compile( tf.keras.optimizers.Adam(),
                   loss=tf.keras.losses.CategoricalCrossentropy(),
                   metrics=[tf.keras.metrics.CategoricalAccuracy()])

    print(' ')

    relativeIndices = np.arange(inputDim)                            # Индексы, принадлежащие одной выборке относительно текущей позиции
    batchIndices    = np.tile( relativeIndices, (opts.batchSize,1) ) # Относительные индексы, повторенные в массиве размером ( batchSize , inputDim )  
    stepsPerEpoch   = int( np.floor( N / opts.batchSize ) )          # Шаги за эпоху

    # Создание экземпляра класса DataGenerator
    generator = DataGenerator(batchSize=opts.batchSize, shuffle=True)

    # Решение, собирая данные в большой массив float32 размером ( N , inputDim ) и передавая его в model.fit

    startTime = time.time()

    X = np.expand_dims( inputData[ np.mod( np.tile(relativeIndices,(N,1)) + np.reshape(trainINDX,(N,1)) , N ) ], axis=2)
    y = outputData[trainINDX, :]

    history = model.fit(x=X, y=y, sample_weight=None, batch_size=opts.batchSize, verbose=1, callbacks=None, validation_split=None, shuffle=True, epochs=opts.epochCount)

    referenceTime = time.time() - startTime
    print(' ')
    print('Общее время решения с model.fit: %6.3f секунд' % referenceTime)
    print(' ')

    # Решение с использованием model.fit_generator  

    startTime = time.time()

    history = model.fit(x=generator, steps_per_epoch=stepsPerEpoch, verbose=1, callbacks=None, epochs=opts.epochCount, max_queue_size=1024, use_multiprocessing=False)

    generatorTime = time.time() - startTime
    print(' ')
    print('Общее время решения с model.fit_generator: %6.3f секунд (%6.2f %% больше)' % (generatorTime, 100.0 * generatorTime/referenceTime))
    print(' ')

    # Решение, собирая данные в батчи размером ( batchSize , inputDim ) и передавая их в model.train_on_batch

    startTime = time.time()

    for epoch in range(opts.epochCount):

        print(' ')
        print('Обучение эпохи # %2d ...' % (epoch+1))
        print(' ')

        np.random.shuffle(trainINDX)

        epochStartTime = time.time()

        for step in tqdm( range( stepsPerEpoch ) ):

            INDX = trainINDX[ step*opts.batchSize : (step+1)*opts.batchSize ]

            X = np.expand_dims( inputData[ np.mod( batchIndices + np.reshape(INDX,(opts.batchSize,1)) , N ) ], axis=2)
            y = outputData[INDX,:]

            history = model.train_on_batch(x=X, y=y, sample_weight=None, class_weight=None, reset_metrics=False)

        print(' ')
        print('...завершено с потерей = %9.6e, точность = %6.2f %%, %6.2f мс/шаг' % (history[0], 100.0*history[1], (1000*(time.time() - epochStartTime)/np.floor(trainINDX.size / opts.batchSize))))
        print(' ')

    batchTime = time.time() - startTime
    print(' ')
    print('Общее время решения с model.train_on_batch: %6.3f секунд (%6.2f %% больше)' % (batchTime, 100.0 * batchTime/referenceTime))
    print(' ')

parser = argparse.ArgumentParser()

parser.add_argument('--batchSize', type=int,
                default=128,
                help='Размер батча')
parser.add_argument('--epochCount', type=int,
                default=5,
                help='Количество эпох')

opts, unparsed = parser.parse_known_args()

if __name__== "__main__":
  main( )
```

Чтобы ответить на вопрос сам, я недавно обновился на Python 3.7.7 и TensorFlow 2.2.0 rc2, и вдруг все мои проблемы исчезли. Теперь,
работая 5 эпох с диаметрально меньшим размером батча 128, model.fit с явно сформированными массивами numpy занимает 126.162 секунды, model.fit с предоставленным генератором занимает 149.053 секунды, а model.train_on_batch занимает 240.698 секунды. Это с обычной версией TensorFlow без поддержки инструкций AVX2 и FMA, поддерживаемых моим процессором.