Итератор с поддержкой Tensorflow, проблемы с производительностью

Question

Я использую TF для обучения ГАН с использованием изображений.До сих пор я строил свою модель и использовал только свой тренировочный набор (который собран из файла tfrecord) и одну эталонную партию для визуализации тензорной доски, и получал приличные постоянные характеристики на Titan V (0,38 ~ 0,40 с / партия) в течение длительного временитак как я использовал инициализируемый итератор и переключался между набором обучающих данных и образцом визуализации (после каждой эпохи обучения).

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

После этого мои выступления ужасно снизились и стали нестабильными.Теперь каждый шаг тренировки занимает от 0,35 до 1,95 с, но все останавливается каждые 16, 32 или 42 итерации на случайную пару секунд (30+).Проблема в том, что это происходит не по фиксированной инструкции, а часто после того же количества итераций).

Это функция, которую я использую для построения конвейера набора данных:

def load_dataset(name, batch_size=32, prefetch_buffer=1, interleave=1, n_threads=os.cpu_count())
    # parse sample is a lambda (omitted)
    parsed_dataset = dataset.map(parse_sample, num_parallel_calls=n_threads)
    # Here there's a series of filter() that I omitted
    parsed_dataset = parsed_dataset.cache()
    parsed_dataset = parsed_dataset.shuffle(shuffle_buffer, reshuffle_each_iteration=True)
    parsed_dataset = parsed_dataset.repeat()
    parsed_dataset = parsed_dataset.batch(batch_size)
    if interleave > 1:
        parsed_dataset = parsed_dataset.interleave(lambda x: tf.data.Dataset.from_tensors(x).repeat(interleave), cycle_length=n_threads, block_length=interleave, num_parallel_calls=n_threads)
    parsed_dataset = parsed_dataset.prefetch(prefetch_buffer)
    return parsed_dataset

Мне нужно чередование, чтобы обучающий набор данных возвращал 3 идентичных пакета при последующих вызовах.

Я ничего не менял в конвейере набора данных, тем не менее я пробовал много разных значений для prefetch_buffer, shuffle_buffer и n_threads.Также попытался удалить repeat() или cache() без существенной разницы.Единственное отличие, которое я получил, заключается в частоте длительного замораживания, но я не мог понять, с каким параметром это связано.

Мой цикл обучения с использованием итерируемого кода:

train_dataset = dh.load_dataset('train',prefetch_buffer=1,shuffle_buffer=32,interleave=3)
validation_dataset = dh.load_dataset('valid',prefetch_buffer=1,shuffle_buffer=32)
tensorboard_datasets = dh.load_dataset('tboard', prefetch_buffer=1, shuffle_buffer=1)

use_dataset=tf.placeholder(tf.string, shape=[])
iterator = tf.data.Iterator.from_string_handle(use_dataset, train_dataset.output_types, train_dataset.output_shapes)
next_batch = iterator.get_next()

train_iterator = train_dataset.make_initializable_iterator()
valid_iterator = validation_dataset.make_initializable_iterator()
tboard_iterator = tensorboard_datasets.make_initializable_iterator()
# Initializers for the datasets
reset_train_iter = train_iterator.initializer
reset_validation_iter = valid_iterator.initializer
reset_tboard_iter = tboard_iterator.initializer

with tf.Session() as sess:
   # Handles to switch between datasets
   use_train_dataset = sess.run(train_iterator.string_handle())
   use_valid_dataset = sess.run(valid_iterator.string_handle())
   use_tboard_dataset = sess.run(tboard_iterator.string_handle())
    while self.layers['train']['global_step'] <= self.params['max_iterations']:
        # Initialize iterators
        sess.run([reset_train_iter, reset_validation_iter, reset_tboard_iter])
        print("Training...")
        for i in range(self.params['save_every_itn']):
            # Training of C and S
            _ = sess.run(self.layers['train']['step_C'], feed_dict={use_dataset: use_train_dataset, is_training: True})
            _ = sess.run(self.layers['train']['step_S'], feed_dict={use_dataset: use_train_dataset, is_training: True})
            # Visualizing losses for each iteration
            self.visualizer.log(sess, show='train_loss', feed_dict={use_dataset: use_train_dataset, is_training: True})
            self.visualizer.log(sess, show='test_loss',  feed_dict={use_dataset: use_valid_dataset, is_training: False})
            self.layers['train']['global_step'] += 1
        print("Checkpoint...")
        self.save(sess)
        # Show a prediction made with a reference sample:
        # Show the differences by using trained and current BatchNorm weights.
        self.visualizer.log(sess, show='train', feed_dict={use_dataset: use_tboard_dataset, is_training: True})
        self.visualizer.log(sess, show='test', feed_dict={use_dataset: use_tboard_dataset, is_training: False})

где «визуализатор» - это класс, в котором я храню весь код для управления Tensorboard.Он просто запускает правильное резюме ('*_loss' - это операция, которая регистрирует только тренировки и тестовые потери, в то время как два других также печатают изображения и другие материалы).Проблема не должна возникать, потому что я пытался удалить вызовы, но проблема с итераторами остается.

Наборы данных Train и Validation на самом деле представляют собой разные фильтры, применяемые к одному и тому же файлу .tfrecords, если это имеет значение.

Я не мог понять, почему вычисления так нестабильны. За ночь было выполнено 14 000 итераций, тогда как в предыдущей версии это было около 90 000.Топ говорит, что процесс Python использует 200% -300% нашего Xeon E5-2609 и 22 из 64 ГБ оперативной памяти при обработке данных.

Есть ли кто-нибудь, кто имеет опыт работы с такого рода итераторами?