Почему циклы for такие медленные в тензорном потоке - PullRequest
Новая
       7

Почему циклы for такие медленные в тензорном потоке

2 голосов
/ 18 июня 2020

Итак, я знаю, что это как-то связано с тем, когда тензорный поток строит граф и делает это не очень хорошо… «эффективно». Вот фиктивный код, который я выполняю: 

@tf.function
def parTest(x_in):
    res = 0
    for i in range(5000):
        res += x_in + i
    return res

выполнение этой функции без использования тензорного потока занимает 0,002 секунды, однако запуск функции с использованием тензорного потока занимает от 10 до 20 секунд. Для меня это не имеет смысла, что здесь происходит? Кроме того, как мне это исправить? Фактическое значение res здесь, очевидно, можно вычислить более эффективным способом, но настоящая проблема, с которой я столкнулся, заключается в том, что у меня есть for l oop, где каждая итерация имеет множество итераций, которые могут выполняться независимо друг от друга, но тензорный поток отказывается это делать и запускает их один за другим очень медленно, как в этом фиктивном примере. Итак, как мне сказать тензорному потоку не делать этого?

1 Ответ

1 голос
/ 18 июня 2020

Циклы никогда не бывают очень эффективными в TensorFlow. Однако эта функция особенно плоха для TensorFlow, потому что она будет пытаться статически "развернуть" весь l oop. То есть он не будет «переводить» вашу функцию в tf.while_loop, а вместо этого будет буквально создавать 5000 копий операций на каждой итерации. Это очень большой график, который всегда будет работать последовательно. На самом деле я получаю предупреждение об этом в TensorFlow 2.2.0, которое указывает на эту информационную страницу: "ПРЕДУПРЕЖДЕНИЕ: обнаружен большой развернутый l oop" .

Как указано в этой ссылке , проблема в том, что TensorFlow не может (по крайней мере, на данный момент) обнаруживать циклы над произвольными итераторами, даже если они простые range, поэтому он просто запускает l oop в Python и создает соответствующие операции. Вы можете избежать этого, написав tf.while_loop самостоятельно или, благодаря AutoGraph , просто заменив range на tf.range: 

import tensorflow as tf
@tf.function
def parTest(x_in):
    res = 0
    for i in tf.range(5000):
        res += x_in + i
    return res

Тем не менее, написание собственного tf.while_loop (когда это абсолютно необходимо, поскольку векторизованные операции всегда будут быстрее) дает вам более явный контроль над такими деталями, как параметр parallel_iterations.

...