Преобразование скользящего окна Tensorflow для больших данных

Я хотел бы накормить свою модель окнами с шагом 1 из очень длинной последовательности данных (десятки миллионов записей). Это похоже на цель, представленную в этой ветке, только моя последовательность данных может содержать несколько функций. для начала, поэтому окончательное количество функций равно n_features * window_size. то есть с двумя исходными функциями и размером окна 3 это означало бы преобразование этого:

[[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]]

to:

[[1, 2, 3, 6, 7, 8], [2, 3, 4, 7, 8, 9], [3, 4, 5, 8, 9, 10]]

Я пытался использовать нарезку с помощью map_fn или Dataset.map, примененную к последовательности индексов (согласно ответу в вышеупомянутой ветке), например:

    ti = tf.range(data.shape[0] - window_size)
    train_dataset = tf.data.Dataset.from_tensor_slices((ti, labels))

    def get_window(l, label):
         wnd = tf.reshape(data_tensor[l:(l + window_size), :], (-1, window_size * n_features))
         wnd = tf.squeeze(wnd)
         return (wnd, label)

    train_dataset = train_dataset.map(get_window)
    train_dataset = train_dataset.batch(batch_size)
    ...

В принципе это работает, но обучение происходит очень медленно, с минимальной загрузкой графического процессора (1-5%, вероятно, частично из-за того, что отображение выполняется в процессоре).

При попытке сделать то же самое с tf.map_fn построение графа становится очень долгим, с огромным использованием памяти.

Другой вариант, который я пробовал, — это преобразовать все данные заранее, прежде чем я загружу их в Tensorflow. Это работает намного быстрее (даже с учетом времени предварительной обработки, интересно, почему — разве это не должно быть той же операцией, что и отображение во время обучения?), но очень неэффективно с точки зрения памяти и хранилища, так как данные становятся window_size- складывать крупнее. Это нарушает условия моих больших наборов данных. Я подумал о том, чтобы разбить эти преобразованные раздутые наборы данных на несколько файлов («гиперпакеты») и просмотреть их последовательно для каждой эпохи, но это кажется очень неэффективным, и мне было интересно, есть ли лучший способ добиться этого простого преобразования.