Using TensorArrays in a while_loop context to accumulate values

Question

Below I have an implementation of RNN Cell Tensorflow designed to emulate the Alex Graves ACT algorithm in this article: http://arxiv.org/abs/1603.08983 .

At the same time, in the sequence called through rnn.rnn (with the static parameter sequence_length, so rnn is expanded dynamically - I use a fixed lot size of 20), we call ACTStep recursively, producing outputs of size (1200), where the hidden size of the RNN cell is 200, and the batch size is 1.

Using the while loop in Tensorflow, we repeat until the accumulated probability of termination is high enough. All this works quite smoothly, but I have problems with the accumulation of states, probabilities and outputs in the while loop, which we need to do to create weighted combinations of them as the final result / state of the cell.

I tried using a simple list, as shown below, but this fails when the graph is compiled, because the outputs are not in the same frame (is it possible to use the "switch" function in control_flow_ops to forward the tensors to the point where they are needed, i.e. add_n function before we return the values?). I also tried using the TensorArray structure, but I find it difficult to use, as it seems to destroy the form information, and replacing it manually did not work. I also could not find much documentation on TensorArrays, presumably as they are, mainly for internal use of TF.

Any advice on how to properly accumulate the variables produced by ACTStep would be greatly appreciated.

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

I will preface this answer that this is NOT a complete solution, but rather a comment on how to improve your cell.

To get started, in your ACTStep function, you call rnn.rnn for one temporary value (as defined by [input] ). If you are executing a single timecode, it is probably more efficient to just use the actual self.cell you will see the same mechanism used in the rnncell wrappers tensor stream

You mentioned that you tried to use TensorArrays . Did you pack and unpack tensor diagrams properly? Here is the repo where you will find in the model.py section the tensor diagrams are packed and unpacked properly.

You also asked if there is a function in control_flow_ops that would require the accumulation of all tensors. I think you are looking for tf.control_dependencies

You can list all your output tensors in control_dependicies, and to calculate all the tensors at this point, you need a tensor flow.

Also, it looks like your counter variable is trainable. Are you sure you want this to be so? If you add a plus one to your counter, this probably will not give the correct result. On the other hand, you could keep it up to date in order to differentiate it at the end for a function of the cost of thinking.

I also think that the Remainder function should be in your script:

 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 

Here is my version of your code:

 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 

This is a complex document to implement that I worked on myself. I would not mind working with you to do this in Tensorflow. If you are interested, please add me to LeavesBreathe via Skype and we can leave from there.