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 — это вызываемый объект, возвращающий булеву скалярную тензор. 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 не вернет ложь.

Для корректности 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, инвариант формы должен быть инвариантом формы тензора значений IndexedSlices. Это означает, что формы трех тензоров IndexedSlices равны (shape, [shape[0]], [shape.ndims]).

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

Для обучения TensorFlow сохраняет тензоры, которые генерируются при прямом выводе и необходимы при обратном распространении. Эти тензоры являются основным источником потребления памяти и часто вызывают ошибки OOM при обучении на 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)