具有可变长度训练数据的 Tensorflow 2 Keras 模型性能不佳

Question

我使用的是 Tensorflow 2.2.0-gpu，我有一个简单的 Keras 模型，它由几个密集层和一个线性输出组成（参考下面的代码）。我在可变长度样本上训练模型，当我运行代码时，我收到关于 tf.function 回溯的警告。根据我的阅读，函数跟踪很昂贵，因此性能很差。这是代码，在我的机器上运行大约需要 330 秒。

#import tensorflow as tf
#tf.compat.v1.disable_eager_execution()

import numpy as np
import timeit
from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras import optimizers

def main():
  state_input = keras.Input((2,))
  hidden1     = layers.Dense(units = 64, activation = "relu")(state_input)
  hidden2     = layers.Dense(units = 128, activation = "relu")(hidden1)
  hidden3     = layers.Dense(units = 128, activation = "relu")(hidden2)
  output      = layers.Dense(units = 2, activation = "linear")(hidden3)

  model = keras.Model(inputs = state_input, outputs = output)
  opt   = optimizers.Adam(lr = 1e-4)

  model.compile(optimizer = opt, loss = "mean_squared_error")

  np.random.seed(0)

  def train():
    for i in range(2000):
      print(i)

      num_samples = np.random.randint(int(1e4), int(1e5))
      x = np.random.rand(num_samples, 2)
      y = np.random.rand(num_samples, 2)

      model.train_on_batch(x, y)

  print(timeit.timeit(train, number=1))

if __name__ == "__main__":
  main()

如果我使用 tf.compat.v1.disable_eager_execution()（代码中的第 2 行）禁用急切执行，那么相同的代码将在大约 30 秒内运行。这与我在 Tensorflow 1 下看到的性能相似。

有没有一种方法可以改变我的模型，从而获得与禁用急切执行时相似的性能？即，是否可以更改模型以使每次调用都不会发生函数回溯？

作为参考，这是调用train_on_batch 时生成的警告：

WARNING:tensorflow:10 out of the last 11 calls to <function Model.make_train_function.<locals>.train_function at 0x7f68f3724158> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings is likely due to passing python objects instead of tensors. Also, tf.function has experimental_relax_shapes=True option that relaxes argument shapes that can avoid unnecessary retracing. Please refer to https://www.tensorflow.org/tutorials/customization/performance#python_or_tensor_args and https://www.tensorflow.org/api_docs/python/tf/function for more details.

Answer 1

通过使用带有签名的tf.function 并手动应用渐变，我能够在不禁用急切模式的情况下提高性能。 （参考 Tensorflow 的 Better performance with tf.function 文章。）这显着提高了性能，但是当完全禁用急切执行时性能仍然更好。

import tensorflow as tf
import numpy as np
import timeit
from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras import optimizers
from tensorflow.keras import losses

def main():
  state_input = keras.Input((2,))
  hidden1     = layers.Dense(units = 64, activation = "relu")(state_input)
  hidden2     = layers.Dense(units = 128, activation = "relu")(hidden1)
  hidden3     = layers.Dense(units = 128, activation = "relu")(hidden2)
  output      = layers.Dense(units = 2, activation = "linear")(hidden3)

  model = keras.Model(inputs = state_input, outputs = output)
  opt   = optimizers.Adam(lr = 1e-4)
  loss  = losses.MeanSquaredError()

  np.random.seed(0)

  @tf.function(input_signature=[
    tf.TensorSpec(shape=(None, 2), dtype=tf.float32),
    tf.TensorSpec(shape=(None, 2), dtype=tf.float32)
  ])
  def fit(x, y):
    with tf.GradientTape() as tape:
      preds = model(x)
      losses = loss(preds, y)
    grad = tape.gradient(losses, model.trainable_variables)
    opt.apply_gradients(zip(grad, model.trainable_variables))

  def train():
    for i in range(2000):
      print(i)

      num_samples = np.random.randint(int(1e4), int(1e5))
      x = np.random.rand(num_samples, 2)
      y = x * 2

      fit(x, y)

  print(timeit.timeit(train, number=1))

  print('test')
  print(model.predict(np.array([[.2, .4], [.6, .8]])))

if __name__ == "__main__":
  main()

但老实说，这很丑。

这是一个很好的问题，为什么 TF2 比 TF1 慢：Why is TensorFlow 2 much slower than TensorFlow 1? 这提供了一些基准。

我的实际代码明显比问题中提供的简单 sn-p 复杂得多，大约是启用急切执行（默认）的速度的 1/10。虽然使用带有签名的tf.function 确实可以加快代码速度，但它仍然不如简单地禁用急切执行快（另外，再次使用tf.function 和GradientTape 非常糟糕）。

最后我只是禁用了急切执行。如果有人提出更好的答案，我很乐意接受。