Декодер поиска векторизованного луча не быстрее на GPU - Tensorflow 2

Question

Я пытаюсь запустить поиск луча RNN на tf.keras.Model векторизованным способом, чтобы он полностью работал на GPU. Тем не менее, несмотря на то, что все равно tf.function, насколько я могу сделать это векторизовано, он работает точно с той же скоростью, с или без графического процессора. Прикрепленный является минимальным примером с поддельной моделью. В действительности, для n = 32, k = 32, steps = 128, с чем я бы хотел работать, для декодирования требуется 20 секунд (на n = 32 выборки), как на CPU, так и на GPU!

Я должен что-то упустить. Когда я тренирую модель, на GPU итерация обучения (128 шагов) с размером пакета 512 занимает 100 мс, а на CPU итерация обучения с размером пакета 32 занимает 1 сек c. GPU не насыщен при размере пакета 512. Я получаю, что у меня есть издержки от выполнения шагов по отдельности и выполнения операции блокировки за шаг, но с точки зрения вычислений мои издержки незначительны по сравнению с остальной частью модели.

Я также понимаю, что использование tf.keras.Model таким образом, вероятно, не идеально, но есть ли другой способ связать выходные тензоры через функцию обратно с входными тензорами и, в частности, также перепроводить состояния?

Полный рабочий пример: https://gist.github.com/meowcat/e3eaa4b8543a7c8444f4a74a9074b9ae

@tf.function
def decode_beam(states_init, scores_init, y_init, steps, k, n):    
    states = states_init
    scores = scores_init
    xstep = embed_y_to_x(y_init)

    # Keep the results in TensorArrays
    y_chain = tf.TensorArray(dtype="int32", size=steps)
    sequences_chain = tf.TensorArray(dtype="int32", size=steps)
    scores_chain = tf.TensorArray(dtype="float32", size=steps)


    for i in range(steps):
        # model_decode is the trained model with 3.5 million trainable params.
        # Run a single step of the RNN model.
        y, states = model_decode([xstep, states])
        # Add scores of step n to previous scores
        # (I left out the sequence end killer for this demo)
        scores_y = tf.expand_dims(tf.reshape(scores, y.shape[:-1]), 2) + tm.log(y)
        # Reshape into (n,k,tokens) and find the best k sequences to continue for each of n candidates
        scores_y = tf.reshape(scores_y, [n, -1])
        top_k = tm.top_k(scores_y, k, sorted=False)
        # Transform the indices. I was using tf.unravel_index but
        # `tf.debugging.set_log_device_placement(True)` indicated that this would be placed on the CPU
        # thus I rewrote it
        top_k_index = tf.reshape(
                top_k[1] + tf.reshape(tf.range(n), (-1, 1)) * scores_y.shape[1], [-1])
        ysequence = top_k_index // y.shape[2]
        ymax = top_k_index % y.shape[2]
        # this gives us two (n*k,) tensors with parent sequence (ysequence) 
        # and chosen character (ymax) per sequence.
        # For continuation, pick the states, and "return" the scores
        states = tf.gather(states, ysequence)
        scores = tf.reshape(top_k[0], [-1])
        # Write the results into the TensorArrays,
        # and embed for the next step
        xstep = embed_y_to_x(ymax)
        y_chain = y_chain.write(i, ymax)
        sequences_chain = sequences_chain.write(i, ysequence)
        scores_chain = scores_chain.write(i, scores)
    # Done: Stack up the results and return them
    sequences_final = sequences_chain.stack()
    y_final = y_chain.stack()
    scores_final = scores_chain.stack()

    return sequences_final, y_final, scores_final

Answer 1

Здесь много чего произошло. Я прокомментирую это, потому что это может помочь другим решить проблемы производительности TensorFlow.

Профилирование

• Библиотека профилировщика графических процессоров (cupti) неправильно загружалась в кластере, останавливая меня делать любое полезное профилирование на GPU. Это было исправлено, поэтому теперь я получаю полезные профили графического процессора.

Обратите внимание на этот очень полезный ответ (единственный в Интернете), который показывает, как профилировать произвольный код TensorFlow 2, а не обучение Keras. :

{ ссылка }

logdir = "log"
writer = tf.summary.create_file_writer(logdir)
tf.summary.trace_on(graph=True, profiler=True)

# run any @tf.function decorated functions here

sequences, y, scores = decode_beam_steps(
    y_init, states_init, scores_init, 
    steps = steps, k = k, n = n, pad_mask = pad_mask)  

with writer.as_default():
    tf.summary.trace_export(name="model_trace", step=0, profiler_outdir=logdir)
tf.summary.trace_off()

Обратите внимание, что старая версия Chromium необходима для просмотра результатов профилирования, начиная с того времени (4-17-20 ) это не работает в текущем Chrome / Chromium.

Небольшие оптимизации

• График стал немного светлее, но не значительно быстрее благодаря использованию unroll=True в используемых ячейках LSTM моделью (не показана здесь), поскольку требуется только один шаг, поэтому символ c l oop только добавляет беспорядок. Это значительно сокращает время первой итерации указанной выше функции, когда AutoGraph строит график. Обратите внимание, что это время огромно (см. Ниже).

  unroll=False (по умолчанию) строится за 300 секунд, unroll=True строится за 100 секунд. Обратите внимание, что сама производительность остается неизменной (15-20 секунд / итерация для n = 32, k = 32).

implementation=1 сделал его немного медленнее, поэтому я остался со значением по умолчанию implementation=2.

Используя tf.while_loop вместо использования AutoGraph

• for i in range(steps) л oop. У меня было это как в (выше показанном) встроенном варианте, так и в модульном:

    for i in range(steps):
        ystep, states = model_decode([xstep, states])
        ymax, ysequence, states, scores = model_beam_step(
            ystep, states, scores, k, n, pad_mask)
        xstep = model_rtox(ymax)
        y_chain = y_chain.write(i, ymax)
        sequences_chain = sequences_chain.write(i, ysequence)
        scores_chain = scores_chain.write(i, scores)

, где model_beam_step делает всю математику поиска луча. Неудивительно, что оба работали точно так же плохо , и, в частности, оба заняли ~ 100/300 секунд при первом запуске, когда AutoGraph отслеживал график. Кроме того, при трассировке графика с помощью профилировщика вы получите сумасшедший файл размером 30-50 МБ, который будет нелегко загрузить на Tensorboard и более или менее обработать его sh. В профиле были десятки параллельных потоков графического процессора, каждая из которых выполнялась по одной операции.

Подставляя это значение в tf.while_loop, сократилось время установки до нуля (back_prop=False делает очень небольшую разницу), и создает хороший график размером 500 КБ, который можно легко просмотреть в TensorBoard и профилировать полезным способом с помощью 4 потоков графического процессора.

    beam_steps_cond = lambda i, y_, seq_, sc_, xstep, states, scores: i < steps
    def decode_beam_steps_body(i, y_, seq_, sc_, xstep, states, scores):
        y, states = model_decode([xstep, states])
        ymax, ysequence, states, scores = model_beam_step(
                y, states, scores, k, n, pad_mask)
        xstep = model_rtox(ymax)
        y_ = y_.write(i, ymax)
        seq_ = seq_.write(i, ysequence)
        sc_= sc_.write(i, scores)
        i = i + 1
        return i, y_, seq_, sc_, xstep, states, scores

    _, y_chain, sequences_chain, scores_chain, _, _, _ = \
        tf.while_loop(
            cond = beam_steps_cond,
            body = decode_beam_steps_body,
            loop_vars = [i, y_chain, sequences_chain, scores_chain,
                         xstep, states, scores],
            back_prop = False
            )

Наконец, реальная проблема

В том, что я действительно смог осмысленный взгляд на профиль показал мне, что реальная проблема заключалась в функции постобработки вывода, работающей на процессоре. Я не подозревал об этом, потому что раньше он работал быстро, но я проигнорировал, что внесенная мной модификация поиска луча приводит к >>> k последовательностям на кандидата, что значительно замедляет обработку. Таким образом, это сокращало все преимущества, которые я мог получить, работая на GPU с шагом декодирования. Без этой постобработки графический процессор запускает> 2 итерации / сек c. Рефакторинг постобработки (которая выполняется очень быстро, если все сделано правильно) в TensorFlow решило проблему.