python - 矢量化波束搜索解码器在 GPU 上并不快 - Tensorflow 2
问题描述
我正在尝试以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
以这种方式使用 a 可能并不理想,但是是否有另一种方法可以通过函数将输出张量连接回输入张量,尤其是重新连接状态?
完整的工作示例: 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
解决方案
这里发生了很多事情。我会对此发表评论,因为它可能会帮助其他人解决 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 中解决了这个问题。
推荐阅读
- hibernate - Spring boot:'spring.datasource.data'定义的资源[/src/main/resources/datadump.sql]不存在
- c# - 数据编辑窗口中收集列表中的绑定数据问题
- angular - 订阅 BehaviorSubject 直到值未定义
- c# - 如果我有一个逗号分隔的字符串,如何过滤值
- bash - 将 [nameref] 声明为指向函数的指针,失败
- swift - 将变量从 UIPageViewController 传递给 UIViewController 子
- javascript - 匹配不带小数且大于 10,000 的数字
- javascript - 如何检查文本框值是否在特定范围内
- javascript - 在 reactJS 中触发 onClick 函数后如何进入初始状态?
- regex - 是否有从 Google 表格中的框中提取特定字符的公式?