Ускорение градиентной ленты Tensorflow 2.0

Question

Я следовал учебному пособию TF 2.0 по сверточным VAE, расположенным здесь .

Поскольку это очень важно, градиенты вычисляются вручную, а затем применяются вручную, используя 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))

Проблема с этим кодом заключается в том, что он довольно медленный и занимает около 40-50 секунд на эпоху.Если я значительно увеличу размер пакета (примерно до 2048), то это займет около 8 секунд за эпоху, но производительность модели значительно снизится.

С другой стороны, если я это сделаюболее традиционная модель (т. е. в которой вместо рвения используется модель на основе ленивых графов), например модель here , тогда она занимает 8 секунд на эпоху даже при небольшом размере пакета.

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

Основываясь на этой информации, я думаю, что проблема с кодом TF2.0 заключается в ручном вычислении потерь и градиентов.

Есть ли способ ускорить код TF2.0чтобы он приблизился к нормальному коду?

Answer 1

Я нашел решение: TensorFlow 2.0 вводит концепцию функций , которые переводят нетерпеливый код в код графа.

Использование довольно простое.Единственное необходимое изменение заключается в том, что все соответствующие функции (например, compute_loss и apply_gradients) должны быть помечены @tf.function.

.