tf.while_loop

Повторять body, пока условие cond истинно. (устаревшие значения аргументов)

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

Используется в блокнотах

Примечание: Эта операция автоматически используется в tf.function для преобразования циклов for и while в Python, если переменная цикла является tf.Tensor, кроме случаев, когда в аргументах tf.function явно указано autograph=False. Например, следующие варианты эквивалентны:

@tf.function
def sumSquare(n):
  i, result = tf.constant(0), tf.constant(0)
  while i < n: # AutoGraph converts while-loop to tf.while_loop().
    result += i * i
    i += 1
  return result
sumSquare(10).numpy()
285

@tf.function
def sumSquare2(n):
  i, result = tf.constant(0), tf.constant(0)
  c = lambda i, _: tf.less(i, n)
  b = lambda i, result: (i + 1, result + i * i)
  return tf.while_loop(c, b, [i, result])[1]
sumSquare2(10).numpy()
285

Дополнительную информацию см. в руководстве по tf.function и AutoGraph.

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])[0]
r.numpy()
10

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

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)[1]
ijk_final[0].numpy(), ijk_final[1].numpy()
(32, 64)

Пример использования 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])])[1]
<tf.Tensor: shape=(2048, 2), dtype=float32, numpy=...>

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

with tf.compat.v1.Session() as sess:
  n = 10
  c = lambda i, x: i < n
  b = lambda i, x: (
      tf.compat.v1.Print(i + 1, [i], "Updating i based on i == "),
      # Let x depend on i
      tf.compat.v1.Print(x + i, [i], "Updating x based on i == "))

  # Make x to be a big matrix so its updating thread would run slowly
  x = tf.zeros([1000, 100], dtype=tf.int32)
  counter = tf.constant(0)
  counter_out, x_out = tf.while_loop(c, b, (counter, x))

  # The following line may increment the counter and x in parallel.
  # The counter thread may get ahead of the x thread, but not the
  # other way around. For example, the log may contain these messages:
  #

... # Обновление i на основе i == [9] ... # Обновление x на основе i == [3] ... # ... # meaning that the counter(i) thread is on iteration 9, ... # while the x thread is on iteration 3. ... print(sess.run(x_out).shape) (1000, 100)