Ограничение количества элементов в tf.data.Dataset

Question

tl; dr; Можно ли ограничить количество элементов в tf.data.Dataset?

A имеет цикл обучения и оценки, который обрабатывает все данные dataset. Это не идеально для тестирования, так как весь набор данных занимает вечность. Я могу протестировать этот код, создав набор данных Mock или ограничив число элементов dataset, чтобы код проходил, скажем, только через первые 10 точек данных. Как я могу сделать второй?

Спасибо

Answer 1

Самый простой способ взять только фиксированное количество элементов n из Dataset - это использовать Dataset.take(n). Например:

large_dataset = ...
small_dataset = large_dataset.take(10)