矢量化波束搜索解码器在 GPU 上并不快 - Tensorflow 2

Question

我正在尝试以矢量化方式在 tf.keras.Model 上运行 RNN 波束搜索，以使其完全在 GPU 上运行。然而，尽管所有东西都像tf.function，尽我所能矢量化，但无论有没有GPU，它的运行速度都完全相同。附件是一个带有假模型的最小示例。实际上，对于 n=32，k=32，steps=128，这是我想要使用的，这需要 20 秒（每个 n=32 个样本）来解码，无论是在 CPU 上还是在 GPU 上！

我一定错过了什么。当我训练模型时，在 GPU 上，批量大小为 512 的训练迭代（128 步）需要 100 毫秒，而在 CPU 上，批量大小为 32 的训练迭代需要 1 秒。 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 性能问题。

分析

• GPU 分析器库 (cupti) 未正确加载到集群上，导致我无法在 GPU 上进行任何有用的分析。该问题已修复，因此我现在可以获得有用的 GPU 配置文件。

请注意这个非常有用的答案（网络上唯一的一个），它展示了如何分析任意 TensorFlow 2 代码，而不是 Keras 训练：

https://stackoverflow.com/a/56698035/1259675

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 中失败。

小优化

• 通过在模型使用的 LSTM 单元中使用unroll=True（此处未显示），图形变得更轻了一点，但速度并没有明显提高，因为只需要一个步骤，因此符号循环只会增加混乱。当 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) 循环。我在（上图所示）内联版本和模块化版本中都有这个：

    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 在哪里进行所有的波束搜索数学运算。不出所料，两者的表现完全相同，特别是在 AutoGraph 跟踪图表时，它们在第一次运行时都花费了 ~100/300 秒。此外，使用分析器跟踪图形会产生一个疯狂的 30-50mb 文件，该文件不会轻易加载到 Tensorboard 上，或多或少会导致它崩溃。该配置文件有数十个并行 GPU 流，每个流都有一个操作。

将其替换为 tf.while_loop 将设置时间缩短到零（back_prop=False 差别很小），并生成了一个漂亮的 500kb 图形，可以在 TensorBoard 中轻松查看并进行分析以一种有用的方式使用 4 个 GPU 流。

    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
            )

最后，真正的问题

我实际上能够以有意义的方式查看配置文件，这表明真正的问题是在 CPU 上运行的输出后处理功能。我没有怀疑它，因为它之前运行得很快，但我忽略了我所做的波束搜索修改导致每个候选者 >>>k 个序列，这大大减慢了处理速度。因此，它削减了我通过解码步骤在 GPU 上高效获得的所有好处。如果没有这种后处理，GPU 会运行 >2 次迭代/秒。将后处理（如果做得好的话会非常快）重构到 TensorFlow 中解决了这个问题。