Caching a computed value as a constant in TensorFlow

Question

Suppose I want to compute the least squares coefficients in TensorFlow using the closed form solution. Normally, I would do this like so,

beta_hat = tf.matmul(
           tf.matmul(tf.matrix_inverse(tf.matmul(tf.transpose(X), X)), tf.transpose(X)), y
)

Where X and y are TensorFlow placeholders corresponding to the covariates and target variable, respectively.

If I then wanted to perform prediction, I would do something like,

y_pred = tf.matmul(X, beta_hat)

If I were to execute,

sess.run(y_pred, feed_dict={X: data_X})

I would of course get an error that I did not provide a necessary value for the placeholder y. I would like to have the flexibility to treat beta_hat as constant after I have computed it (so that I would not need to define a new placeholder for the new covariate matrix for prediction). One way to accomplish this is,

# Make it constant.
beta_hat = sess.run(beta_hat, feed_dict={X: data_X, y: data_y})
y_pred = tf.matmul(X, beta_hat)

I was wondering if there were a more elegant way to treat the tensor as constant so that I neither need to execute the session and obtain a constant nor create a separate placeholder for incoming data to be used for prediction.

Here is some sample code that demonstrates the circumstance I'm describing.

import numpy as np
import tensorflow as tf


n, k = 100, 5
X = tf.placeholder(dtype=tf.float32, shape=[None, k])
y = tf.placeholder(dtype=tf.float32, shape=[None, 1])

beta = np.random.normal(size=(k, ))
data_X = np.random.normal(size=(n, k))

data_y = data_X.dot(beta)
data_y += np.random.normal(size=data_y.shape) / 3.0
data_y = np.atleast_2d(data_y).T

# Convert to 32-bit precision.
data_X, data_y = np.float32(data_X), np.float32(data_y)

# Compute the least squares solution.
beta_hat = tf.matmul(
    tf.matmul(tf.matrix_inverse(tf.matmul(tf.transpose(X), X)),
              tf.transpose(X)), y
)

# Launch the graph
sess = tf.Session()
sess.run(tf.initialize_all_variables())

print "True beta: {}".format(beta)
print "Est. beta: {}".format(
    sess.run(beta_hat, feed_dict={X: data_X, y: data_y}).ravel()
)

# # This would error.
# y_pred = tf.matmul(X, beta_hat)
# print "Predictions:"
# print sess.run(y_pred, feed_dict={X: data_X})

# Make it constant.
beta_hat = sess.run(beta_hat, feed_dict={X: data_X, y: data_y})

# This will no longer error.
y_pred = tf.matmul(X, beta_hat)
print "Predictions:"
print sess.run(y_pred, feed_dict={X: data_X})

Answer 1

Perhaps counter-intuitively, the simplest way to re-use beta_hat as a constant in subsequent steps would be to assign it to a tf.Variable:

n, k = 100, 5
X = tf.placeholder(dtype=tf.float32, shape=[None, k])
y = tf.placeholder(dtype=tf.float32, shape=[None, 1])

beta = np.random.normal(size=(k, ))
data_X = np.random.normal(size=(n, k))

data_y = data_X.dot(beta)
data_y += np.random.normal(size=data_y.shape) / 3.0
data_y = np.atleast_2d(data_y).T

# Convert to 32-bit precision.
data_X, data_y = np.float32(data_X), np.float32(data_y)

# Compute the least squares solution.
beta_hat = tf.matmul(
    tf.matmul(tf.matrix_inverse(tf.matmul(tf.transpose(X), X)),
              tf.transpose(X)), y
)

beta_hat_cached = tf.Variable(beta_hat)

# Launch the graph
sess = tf.Session()

print "True beta: {}".format(beta)
# Run the initializer, which computes `beta_hat` once:
sess.run(beta_hat_cached.initializer, feed_dict={X: data_X, y: data_y})
# To access the value of `beta_hat`, "run" the variable to read its contents.
print "Est. beta: {}".format(beta_hat_cached.ravel())

# Use the cached version to compute predictions.
y_pred = tf.matmul(X, beta_hat_cached)
print "Predictions:"
print sess.run(y_pred, feed_dict={X: data_X})