Применить dataset.repeat предварительно извлеченного набора данных

Я пытаюсь реализовать алгоритм квантования АдаРаунд, и мне нужно тренировать свои слои один за другим.

Я использую набор данных с 1024 с размером пакета 32, и мне нужно выполнить итерацию по набору данных примерно 312 эпох (или 10 000 итераций по пакетному набору данных), я заметил, что данные копируются с хоста на устройство каждую итерацию и что данные не кэшируются на графическом процессоре (несмотря на повторное использование одних и тех же данных) - графический процессор простаивает 30–40% процентов времени

Процент бездействия графического процессора

Данные по-прежнему копируются с хоста на устройство в расширенных итерациях:

Чанк memcpyH2D за одну итерацию

я пытался использовать

tf.data.experimental.prefetch_to_device

tf.data.experimental.copy_to_device

но когда я перебираю данные после prefetch_to_device или copy_to_device, тензоры сохраняются на графическом процессоре, но если я использую повтор для просмотра набора данных, тензоры сохраняются на ЦП

Я пытался использовать model.fit без dataset.repeat, но с несколькими эпохами, но получаю похожее поведение.

Я также пытался использовать model.fit с тензорами, которые хранятся на графическом процессоре, но подгонка модели преобразует его в набор данных, который возвращает данные обратно в процессор.

Фрагмент кода для воссоздания проблемы:

input_shape = (56, 56, 64)
output_shape = (54, 54, 64)
conv = tf.keras.layers.Conv2D(64, (3, 3))
mock_input = tf.keras.layers.Input(input_shape)
mock_output = conv(mock_input)
train_model = tf.keras.Model(inputs=mock_input, outputs=mock_output)

input_data = np.random.rand(1024, *input_shape)
output_data = np.random.rand(1024, *output_shape)

input_dataet = tf.data.Dataset.from_tensor_slices(input_data)
output_dataset = tf.data.Dataset.from_tensor_slices(output_data)

train_model.compile(
    optimizer='adam',
    loss='mse'
)
train_data = tf.data.Dataset.zip((input_dataet, output_dataset))
batched_train_data = train_data.batch(32).cache()
fetched_train_data = batched_train_data.prefetch(tf.data.AUTOTUNE).repeat()
with tf.profiler.experimental.Profile('logs'):
    train_model.fit(fetched_train_data, steps_per_epoch=1024, epochs=1)

Есть ли способ применить операцию dataset.repeat к графическому процессору?

• Я использую тензорный поток 2.5.2 с питоном 3.6.9.