加速 TensorFlow 2.0 渐变胶带

Question

我一直在关注卷积 VAE 的 TF 2.0 教程，位于 here。

因为它很急切，所以梯度是手动计算的，然后使用 tf.GradientTape() 手动应用。

for epoch in epochs:
  for x in x_train:
    with tf.GradientTape() as tape:
      loss = compute_loss(model, x)
    apply_gradients(tape.gradient(loss, model.trainable_variables))

该代码的问题在于它非常慢，每个 epoch 大约需要 40-50 秒。 如果我将批量大小增加很多（到 2048 左右），那么最终每个 epoch 需要大约 8 秒，但模型的性能会下降很多。

另一方面，如果我做一个更传统的模型（即使用基于惰性图的模型而不是 Eagerness），例如 here，那么每个 epoch 需要 8 秒，即使是很小的批量大小。

model.add_loss(lazy_graph_loss)
model.fit(x_train epochs=epochs)

根据这些信息，我的猜测是 TF2.0 代码的问题在于手动计算损失和梯度。

有没有什么办法可以加快TF2.0代码的速度，使其更接近正常的代码？

Answer 1

我找到了解决方案：TensorFlow 2.0 引入了functions 的概念，将 Eager 代码转换为图形代码。

用法非常简单。唯一需要更改的是所有相关函数（如compute_loss 和apply_gradients）都必须使用@tf.function 进行注释。

Answer 2

您的代码花费很长时间的主要原因是，当我们使用没有任何张量的普通 pythonic for 循环时，图的构建需要大量时间，因为直觉上我们可能认为同一个图for training 可能会在每次迭代中被重新使用，但构建的图实际上是一个链状结构，其中每个节点都是训练子图，而总数。该链中的节点数相当于第一个。循环中的迭代。简而言之，张量流解开迭代，然后构造图。因此，它在空间和时间方面单独为图使用了大量冗余。太糟糕了，仅仅在一个正常的 pythonic 循环中重复添加两个张量大约十亿次，就需要将近半个小时。

为了解决这个问题，特别是在您的情况下，我们可以借助 tf.data.Datasets api 中的 .repeat 转换，而不是编写

for i in range(epochs) :

我们可以写

For x in x_train.repeat(epochs) :
        Train here