Reuse value of TensorFlow Variable between sessions without writing to disk

Viewed 1347

In sklearn, I'm used to having a model that I can run fit and then predict on. However, with TensorFlow, I'm having trouble loading the learned parameters from fit when I'm calling predict. It boils down to me not knowing how to reuse the value of a variable between sessions. For example,

import tensorflow as tf

x = tf.Variable(0.0)

# fit code
with tf.Session() as sess1:
    sess1.run(tf.global_variables_initializer())
    sess1.run(tf.assign(x, 1.0)) # at end of training, x = 1.0

# predict code
with tf.Session() as sess2:
    sess2.run(tf.global_variables_initializer())
    print(sess2.run(x)) # want this to be 1.0, but is 0.0

I can think of one workaround, but it seems really hacky, and would be annoying if there are several variables I want to reuse:

import tensorflow as tf

x = tf.Variable(0.0)

# fit code
with tf.Session() as sess1:
    sess1.run(tf.global_variables_initializer())
    sess1.run(tf.assign(x, 1.0)) # at end of training, x = 1.0
    learned_x = sess1.run(x) # remember value of learned x at end of session

# predict code
with tf.Session() as sess2:
    sess2.run(tf.global_variables_initializer())
    sess2.run(tf.assign(x, learned_x))
    print(sess2.run(x)) # prints 1.0

How do I reuse variables between sessions without writing to disk (i.e. using tf.train.Saver)? Is the workaround I wrote above the right way to do this?

1 Answers

To mimic sklearn's model, just wrap session into a single class so that you can share it between methods e.g.

class Model:
    def __init__(self):
        self.graph = self.build_graph()
        self.session = tf.Session()
        self.session.run(tf.global_variables_initializer())

    def build_graph(self):
        return {'x': tf.Variable(0.0)}

    def fit(self):
        self.session.run(tf.assign(self.graph['x'], 1.0))

    def predict(self):
        print(self.session.run(self.graph['x']))

    def close(self):
        tf.reset_default_graph()
        self.session.close()

m = Model()
m.fit()
m.predict()
m.close()

Make sure you close the session manually and handle exceptions accordingly.

Related