在递归循环期间分配给 TensorFlow 变量

Question

In Tensorflow 1.9, I want to create a network and then recursively feed the output (the prediction) of the network back into the input of the network. During this loop, I want to store the predictions made by the network in a list.

Here is my attempt:

    # Define the number of steps over which to loop the network
    num_steps = 5

    # Define the network weights
    weights_1 = np.random.uniform(0, 1, [1, 10]).astype(np.float32)
    weights_2 = np.random.uniform(0, 1, [10, 1]).astype(np.float32)

    # Create a variable to store the predictions, one for each loop
    predictions = tf.Variable(np.zeros([num_steps, 1]), dtype=np.float32)

    # Define the initial prediction to feed into the loop
    initial_prediction = np.array([[0.1]], dtype=np.float32)
    x = initial_prediction

    # Loop through the predictions
    for step_num in range(num_steps):
        x = tf.matmul(x, weights_1)
        x = tf.matmul(x, weights_2)
        predictions[step_num-1].assign(x)

    # Define the final prediction
    final_prediction = x

    # Start a session
    sess = tf.Session()
    sess.run(tf.global_variables_initializer())

    # Make the predictions
    last_pred, all_preds = sess.run([final_prediction, predictions])
    print(last_pred)
    print(all_preds)

And this prints out:

[[48.8769]]

[[0.]
 [0.]
 [0.]
 [0.]
 [0.]]

So whilst the value of final_prediction appears correct, the value of predictions is not what I would expect. It seems that predictions is never actually assigned to, despite the line predictions[step_num-1].assign(x) .

Please can somebody explain to me why this isn't working, and what I should be doing instead? Thanks!

Answer 1

This happens because assign ist just a TF op like any other, and as such is only executed if needed. Since nothing on the path to final_prediction relies on the assign op, and predictions is just a variable, the assignment is never executed.

I think the most straightforward solution would be to replace the line

predictions[step_num-1].assign(x)

by

x = predictions[step_num-1].assign(x)

This works because assign also returns the value it is assigning. Now, to compute final_prediction TF actually needs to "go through" the assign op so the assignments should be carried out.

Another option would be to use tf.control_dependencies which is a way to "force" TF to compute specific ops when it is computing other ones. However in this case it could be a bit icky because the op we want to force ( assign ) depends on values that are being computed within the loop and I'm not sure about the order in which TF does stuff in this case. The following should work:

for step_num in range(num_steps):
    x = tf.matmul(x, weights_1)
    x = tf.matmul(x, weights_2)
    with tf.control_dependencies([predictions[step_num-1].assign(x)]):
        x = tf.identity(x)

We use tf.identity as a noop just to have something to wrap with control_dependencies . I think this is the more flexible option between the two. However it comes with some caveats discussed in the docs .