tf.compat.v1.while_loop

Повторять body, пока условие cond истинно.

tf.compat.v1.while_loop(
    cond,
    body,
    loop_vars,
    shape_invariants=None,
    parallel_iterations=10,
    back_prop=True,
    swap_memory=False,
    name=None,
    maximum_iterations=None,
    return_same_structure=False
)

cond — это вызываемый объект, возвращающий скалярный тензор типа boolean. body — это вызываемый объект, возвращающий (возможно, вложенный) кортеж, именованный кортеж или список тензоров с такой же арностью (длиной и структурой) и типами, что и loop_vars. loop_vars — это (возможно, вложенный) кортеж, именованный кортеж или список тензоров, передаваемый как в cond, так и в body. cond и body принимают столько же аргументов, сколько имеется loop_vars.

Помимо обычных тензоров и IndexedSlices, тело цикла может принимать и возвращать объекты TensorArray. Потоки объектов TensorArray будут корректно передаваться между циклами и при вычислении градиентов.

Обратите внимание, что while_loop вызывает cond и body ровно один раз (внутри вызова while_loop, но не во время Session.run()). while_loop объединяет фрагменты графа, созданные во время вызовов cond и body, с некоторыми дополнительными узлами графа, чтобы создать поток графа, повторяющий body до тех пор, пока cond не вернёт false.

Для обеспечения корректности tf.while_loop() строго проверяет инварианты формы для переменных цикла. Инвариант формы — это (возможно, частичная) форма, которая не изменяется между итерациями цикла. Если после итерации форма переменной цикла окажется более общей или несовместимой с её инвариантом формы, будет вызвана ошибка. Например, форма [11, None] более общая, чем форма [11, 17], а [11, 21] несовместима с [11, 17]. По умолчанию (если аргумент shape_invariants не указан) предполагается, что начальная форма каждого тензора в loop_vars остаётся неизменной во всех итерациях. Аргумент shape_invariants позволяет вызывающему коду задать менее строгий инвариант формы для каждой переменной цикла; это необходимо, если форма меняется между итерациями. Функцию tf.Tensor.set_shape также можно использовать в функции body, чтобы указать конкретную форму выходной переменной цикла. Инварианты формы для SparseTensor и IndexedSlices обрабатываются особым образом:

a) Если переменная цикла — SparseTensor, инвариант формы должен быть TensorShape([r]), где r — ранг плотного тензора, представленного разреженным тензором. Это означает, что формы трёх тензоров SparseTensor равны ([None], [None, r], [r]). ПРИМЕЧАНИЕ: в этом случае инвариант формы задаёт форму свойства SparseTensor.dense_shape. Это должна быть форма вектора.

b) Если переменная цикла — IndexedSlices, инвариант формы должен задавать инвариант формы тензора values объекта IndexedSlices. Это означает, что формы трёх тензоров IndexedSlices равны (shape, [shape[0]], [shape.ndims]).

while_loop реализует нестрогую семантику, позволяя выполнять несколько итераций параллельно. Максимальное количество параллельных итераций можно задать с помощью parallel_iterations, что позволяет пользователям в некоторой степени контролировать потребление памяти и порядок выполнения. Для корректных программ while_loop должен возвращать одинаковый результат при любом значении parallel_iterations > 0.

Во время обучения TensorFlow сохраняет тензоры, созданные при прямом проходе и необходимые для обратного распространения ошибки. Эти тензоры являются одним из основных источников потребления памяти и часто вызывают ошибки нехватки памяти при обучении на GPU. Если флаг swap_memory имеет значение true, эти тензоры переносятся из GPU в CPU. Например, это позволяет обучать модели RNN на очень длинных последовательностях и больших пакетах.

Пример:

i = tf.constant(0)
c = lambda i: tf.less(i, 10)
b = lambda i: tf.add(i, 1)
r = tf.while_loop(c, b, [i])

Пример с вложенной структурой и именованным кортежем:

import collections
Pair = collections.namedtuple('Pair', 'j, k')
ijk_0 = (tf.constant(0), Pair(tf.constant(1), tf.constant(2)))
c = lambda i, p: i < 10
b = lambda i, p: (i + 1, Pair((p.j + p.k), (p.j - p.k)))
ijk_final = tf.while_loop(c, b, ijk_0)

Пример использования shape_invariants:

i0 = tf.constant(0)
m0 = tf.ones([2, 2])
c = lambda i, m: i < 10
b = lambda i, m: [i+1, tf.concat([m, m], axis=0)]
tf.while_loop(
    c, b, loop_vars=[i0, m0],
    shape_invariants=[i0.get_shape(), tf.TensorShape([None, 2])])

Пример, демонстрирующий нестрогую семантику: в следующем примере итоговое значение счётчика i не зависит от x. Поэтому while_loop может увеличивать счётчик параллельно с обновлением x. Однако, поскольку значение счётчика цикла на каждой итерации зависит от значения на предыдущей итерации, сам счётчик цикла нельзя увеличивать параллельно. Следовательно, если нам нужно только итоговое значение счётчика (которое выводится в строке print(sess.run(i))), x никогда не будет увеличиваться, а счётчик будет обновляться в одном потоке. И наоборот, если нам нужно значение выходных данных (которое выводится в строке print(sess.run(out).shape)), счётчик может увеличиваться в отдельном потоке, а x может параллельно увеличиваться в другом потоке. В крайнем случае поток, увеличивающий счётчик, может завершить работу, прежде чем x увеличится хотя бы один раз. Единственное, что невозможно, — это чтобы поток, обновляющий x, опережал поток счётчика, поскольку поток, увеличивающий x, зависит от значения счётчика.

import tensorflow as tf

n = 10000
x = tf.constant(list(range(n)))
c = lambda i, x: i < n
b = lambda i, x: (tf.compat.v1.Print(i + 1, [i]), tf.compat.v1.Print(x + 1,
[i], "x:"))
i, out = tf.while_loop(c, b, (0, x))
with tf.compat.v1.Session() as sess:
    print(sess.run(i))  # prints [0] ... [9999]

    # The following line may increment the counter and x in parallel.
    # The counter thread may get ahead of the other thread, but not the
    # other way around. So you may see things like
    # [9996] x:[9987]
    # meaning that the counter thread is on iteration 9996,
    # while the other thread is on iteration 9987
    print(sess.run(out).shape)