tensorflow - 在 while_loop 的上下文中使用 TensorArrays 来累积值

Question

下面我有一个 Tensorflow RNN Cell 的实现，旨在模拟本文中 Alex Graves 的算法 ACT:http://arxiv.org/abs/1603.08983 .

在通过 rnn.rnn 调用的序列中的单个时间步长(使用静态 sequence_length 参数，因此 rnn 是动态展开的 - 我使用的是固定批量大小 20)，我们递归调用 ACTStep，生成大小为 (1,200) 的输出，其中RNN 单元的隐藏维度为 200，我们的批次大小为 1。

使用 Tensorflow 中的 while 循环，我们迭代直到累积的停机概率足够高。所有这些工作都相当顺利，但我在 while 循环中积累状态、概率和输出时遇到问题，我们需要这样做才能创建这些的加权组合作为最终的单元输出/状态。

我曾尝试使用一个简单的列表，如下所示，但是当图形编译时失败，因为输出不在同一帧中(是否可以使用 control_flow_ops 中的“switch”函数将张量转发到它们是必需的，即在我们返回值之前的 add_n 函数？)。我也尝试过使用 TensorArray 结构，但我发现这很难使用，因为它似乎破坏了形状信息并且手动替换它没有用。我也没有找到很多关于 TensorArrays 的文档，我想大概是这样，主要用于内部 TF 使用。

关于如何正确累积 ACTStep 产生的变量的任何建议将不胜感激。

class ACTCell(RNNCell):
"""An RNN cell implementing Graves' Adaptive Computation time algorithm"""
def __init__(self, num_units, cell, epsilon, max_computation):

    self.one_minus_eps = tf.constant(1.0 - epsilon)
    self._num_units = num_units
    self.cell = cell
    self.N = tf.constant(max_computation)
@property
def input_size(self):
    return self._num_units
@property
def output_size(self):
    return self._num_units
@property
def state_size(self):
    return self._num_units

def __call__(self, inputs, state, scope=None):

    with vs.variable_scope(scope or type(self).__name__):

        # define within cell constants/ counters used to control while loop
        prob = tf.get_variable("prob", [], tf.float32,tf.constant_initializer(0.0))
        counter = tf.get_variable("counter", [],tf.float32,tf.constant_initializer(0.0))
        tf.assign(prob,0.0)
        tf.assign(counter, 0.0)

        # the predicate for stopping the while loop. Tensorflow demands that we have
        # all of the variables used in the while loop in the predicate.
        pred = lambda prob,counter,state,input,\
                      acc_state,acc_output,acc_probs:\
            tf.logical_and(tf.less(prob,self.one_minus_eps), tf.less(counter,self.N))

        acc_probs = []
        acc_outputs = []
        acc_states = []


        _,iterations,_,_,acc_states,acc_output,acc_probs = \
        control_flow_ops.while_loop(pred,
        self.ACTStep,
        [prob,counter,state,input,acc_states,acc_outputs,acc_probs])

    # TODO:fix last part of this, need to use the remainder.
    # TODO: find a way to accumulate the regulariser

    # here we take a weighted combination of the states and outputs 
    # to use as the actual output and state which is passed to the next timestep.

    next_state = tf.add_n([tf.mul(x,y) for x,y in zip(acc_probs,acc_states)])
    output = tf.add_n([tf.mul(x,y) for x,y in zip(acc_probs,acc_outputs)])


    return output, next_state

def ACTStep(self,prob,counter,state,input, acc_states,acc_outputs,acc_probs):

    output, new_state = rnn.rnn(self.cell, [input], state, scope=type(self.cell).__name__)

    prob_w = tf.get_variable("prob_w", [self.cell.input_size,1])
    prob_b = tf.get_variable("prob_b", [1])
    p = tf.nn.sigmoid(tf.matmul(prob_w,new_state) + prob_b)

    acc_states.append(new_state)
    acc_outputs.append(output)
    acc_probs.append(p)

    return [tf.add(prob,p),tf.add(counter,1.0),new_state, input,acc_states,acc_outputs,acc_probs]

Answer 1

我将在此回复前说明这不是一个完整的解决方案，而是关于如何改进您的单元格的一些评论。

首先，在您的 ACTStep 函数中，您调用 rnn.rnn对于一个时间步长(由 [input] 定义。如果您正在执行单个时间步长，那么简单地使用实际的 self.cell 调用函数可能更有效。您将在 tensorflow rnncell wrappers 中看到相同的机制

您提到您曾尝试使用 TensorArrays .你是否适本地打包和解包张量数组？这是一个 repo您可以在 model.py 下找到张量数组被正确打包和解包。

你还问了control_flow_ops中是否有函数这将需要累积所有张量。我认为您正在寻找 tf.control_dependencies

您可以在 control_dependicies 中列出所有输出张量操作，这将需要 tensorflow 计算到该点的所有张量。

另外，它看起来像您的 counter变量是可训练的。你确定要这样吗？如果您在计数器上加一，那可能不会产生正确的结果。另一方面，您可以故意保持它的可训练性，以便在最后为思考成本函数区分它。

另外我相信 Remainder 函数应该在你的脚本中:

remainder = 1.0 - tf.add_n(acc_probs[:-1])
#note that there is a -1 in the list as you do not want to grab the last probability

这是我编辑的代码版本:

class ACTCell(RNNCell):
    """An RNN cell implementing Graves' Adaptive Computation time algorithm
    Notes: https://www.evernote.com/shard/s189/sh/fd165646-b630-48b7-844c-86ad2f07fcda/c9ab960af967ef847097f21d94b0bff7

    """
    def __init__(self, num_units, cell, max_computation = 5.0, epsilon = 0.01):

        self.one_minus_eps = tf.constant(1.0 - epsilon) #episolon is 0.01 as found in the paper
        self._num_units = num_units
        self.cell = cell
        self.N = tf.constant(max_computation)

    @property
    def input_size(self):
        return self._num_units
    @property
    def output_size(self):
        return self._num_units
    @property
    def state_size(self):
        return self._num_units

    def __call__(self, inputs, state, scope=None):

        with vs.variable_scope(scope or type(self).__name__):

            # define within cell constants/ counters used to control while loop
            prob = tf.constant(0.0, shape = [batch_size]) 
            counter = tf.constant(0.0, shape = [batch_size])

            # the predicate for stopping the while loop. Tensorflow demands that we have
            # all of the variables used in the while loop in the predicate.
            pred = lambda prob,counter,state,input,acc_states,acc_output,acc_probs:\
                tf.logical_and(tf.less(prob,self.one_minus_eps), tf.less(counter,self.N))

            acc_probs, acc_outputs, acc_states  = [], [], []

            _,iterations,_,_,acc_states,acc_output,acc_probs = \
            control_flow_ops.while_loop(
            pred,
            self.ACTStep, #looks like he purposely makes the while loop here
            [prob, counter, state, input, acc_states, acc_outputs, acc_probs])

        '''mean-field updates for states and outputs'''
        next_state = tf.add_n([x*y for x,y in zip(acc_probs,acc_states)])
        output = tf.add_n([x*y for x,y in zip(acc_probs,acc_outputs)])

        remainder = 1.0 - tf.add_n(acc_probs[:-1]) #you take the last off to avoid a negative ponder cost #the problem here is we need to take the sum of all the remainders
        tf.add_to_collection("ACT_remainder", remainder) #if this doesnt work then you can do self.list based upon timesteps
        tf.add_to_collection("ACT_iterations", iterations)
        return output, next_state 

    def ACTStep(self,prob, counter, state, input, acc_states, acc_outputs, acc_probs):

        '''run rnn once'''
        output, new_state = rnn.rnn(self.cell, [input], state, scope=type(self.cell).__name__)

        prob_w = tf.get_variable("prob_w", [self.cell.input_size,1]) 
        prob_b = tf.get_variable("prob_b", [1])
        halting_probability = tf.nn.sigmoid(tf.matmul(prob_w,new_state) + prob_b) 


        acc_states.append(new_state)
        acc_outputs.append(output)
        acc_probs.append(halting_probability) 

        return [p + prob, counter + 1.0, new_state, input,acc_states,acc_outputs,acc_probs]


    def PonderCostFunction(self, time_penalty = 0.01):
        '''
        note: ponder is completely different than probability and ponder = roe

        the ponder cost function prohibits the rnn from cycling endlessly on each timestep when not much is needed
        '''
        n_iterations = tf.get_collection_ref("ACT_iterations")
        remainder = tf.get_collection_ref("ACT_remainder")
        return tf.reduce_sum(n_iterations + remainder) #completely different from probability

这是一篇复杂的论文，我一直在努力实现。我不介意与您合作在 Tensorflow 中完成它。如果您有兴趣，请在 Skype 上的 LeavesBreathe 上加我，我们可以从那里开始。