Tensorflow: подсчитайте количество примеров в файле TFRecord — без использования устаревшего `tf.python_io.tf_record_iterator`

Я получил следующий код для работы без устаревшей команды. Надеюсь, это поможет другим.

Используя API набора данных, я настраиваю итератор, а затем перебираю его. Не уверен, что это самый быстрый, но он работает. УБЕДИТЕСЬ, ЧТО РАЗМЕР ПАКЕТА И ПОВТОР УСТАНОВЛЕНЫ НА 1, иначе код вернет количество пакетов, а не количество примеров в наборе данных.

count_test = tf.data.TFRecordDataset('testing.tfrecord')
count_test = count_test.map(_parse_image_function)
count_test = count_test.repeat(1)
count_test = count_test.batch(1)
test_counter = count_test.make_one_shot_iterator()

c = 0
for ex in test_counter:
    c += 1
f"There are {c} testing records"

Это, казалось, работало достаточно хорошо даже на относительно большом файле.

В классе reduce есть метод Dataset. Они приводят пример подсчета записей по методу:

# generate the dataset (batch size and repeat must be 1, maybe avoid dataset manipulation like map and shard)
ds = tf.data.Dataset.range(5) 
# count the examples by reduce
cnt = ds.reduce(np.int64(0), lambda x, _: x + 1)

## produces 5

Не знаю, быстрее ли этот метод, чем цикл for @krishnab.

Следующее работает для меня с использованием TensorFlow версии 2.1 (используя код, найденный в этот ответ):

def count_tfrecord_examples(
        tfrecords_dir: str,
) -> int:
    """
    Counts the total number of examples in a collection of TFRecord files.

    :param tfrecords_dir: directory that is assumed to contain only TFRecord files
    :return: the total number of examples in the collection of TFRecord files
        found in the specified directory
    """

    count = 0
    for file_name in os.listdir(tfrecords_dir):
        tfrecord_path = os.path.join(tfrecords_dir, file_name)
        count += sum(1 for _ in tf.data.TFRecordDataset(tfrecord_path))

    return count