tf.compat.v1.nn.raw_rnn

Создаёт RNN, заданный ячейкой RNN cell и функцией цикла loop_fn.

tf.compat.v1.nn.raw_rnn(
    cell, loop_fn, parallel_iterations=None, swap_memory=False, scope=None
)

Примечание: Этот метод всё ещё находится в стадии тестирования, и API может быть изменён.**

Эта функция является более примитивной версией dynamic_rnn, которая предоставляет более прямой доступ к входным данным на каждой итерации. Она также обеспечивает больший контроль над тем, когда начинать и заканчивать чтение последовательности, и что выводить для вывода.

Например, она может быть использована для реализации динамического декодера модели seq2seq.

Вместо работы с Tensor объектами, большинство операций работают непосредственно с TensorArray объектами.

Действие raw_rnn, в псевдокоде, в основном следующее:

time = tf.constant(0, dtype=tf.int32)
(finished, next_input, initial_state, emit_structure, loop_state) = loop_fn(
    time=time, cell_output=None, cell_state=None, loop_state=None)
emit_ta = TensorArray(dynamic_size=True, dtype=initial_state.dtype)
state = initial_state
while not all(finished):
  (output, cell_state) = cell(next_input, state)
  (next_finished, next_input, next_state, emit, loop_state) = loop_fn(
      time=time + 1, cell_output=output, cell_state=cell_state,
      loop_state=loop_state)
  # Emit zeros and copy forward state for minibatch entries that are finished.
  state = tf.where(finished, state, next_state)
  emit = tf.where(finished, tf.zeros_like(emit_structure), emit)
  emit_ta = emit_ta.write(time, emit)
  # If any new minibatch entries are marked as finished, mark these.
  finished = tf.logical_or(finished, next_finished)
  time += 1
return (emit_ta, state, loop_state)

с дополнительными свойствами, что вывод и состояние могут быть (возможно, вложенными) кортежами, как определено cell.output_size и cell.state_size, и в результате конечные state и emit_ta могут сами быть кортежами.

Простой пример реализации dynamic_rnn через raw_rnn выглядит так:

inputs = tf.compat.v1.placeholder(shape=(max_time, batch_size, input_depth),
                        dtype=tf.float32)
sequence_length = tf.compat.v1.placeholder(shape=(batch_size,),
dtype=tf.int32)
inputs_ta = tf.TensorArray(dtype=tf.float32, size=max_time)
inputs_ta = inputs_ta.unstack(inputs)

cell = tf.compat.v1.nn.rnn_cell.LSTMCell(num_units)

def loop_fn(time, cell_output, cell_state, loop_state):
  emit_output = cell_output  # == None for time == 0
  if cell_output is None:  # time == 0
    next_cell_state = cell.zero_state(batch_size, tf.float32)
  else:
    next_cell_state = cell_state
  elements_finished = (time >= sequence_length)
  finished = tf.reduce_all(elements_finished)
  next_input = tf.cond(
      finished,
      lambda: tf.zeros([batch_size, input_depth], dtype=tf.float32),
      lambda: inputs_ta.read(time))
  next_loop_state = None
  return (elements_finished, next_input, next_cell_state,
          emit_output, next_loop_state)

outputs_ta, final_state, _ = raw_rnn(cell, loop_fn)
outputs = outputs_ta.stack()