Tensorflow: How to save the model in the memory while training

Question

My training process in tensorflow involves switching between two models. While using tf.saver and restore the model from the hard disk is really time-consuming (In my code, the switching is frequent), therefore, I wonder if there is a way to store the model parameters in the memory and restore them just from memory. My model is rather small, which can be definitely stored in RAM. There is one answer from stackoverflow. Storing tensorflow models in memory However, I do not quite understand how this can work. Does anyone know how to achieve this goal? Thank you.

Answer 1

You should just use two separate graphs like this:

g1 = tf.Graph()
g2 = tf.Graph()

with g1.as_default():
  # build your 1st model
  sess1 = tf.Session(graph=g1)
  # do some work with sess1 on g1
  sess1.run(...)

with g2.as_default():
  # build your 2nd model
  sess2 = tf.Session(graph=g2)
  # do some work with sess2 on g2
  sess2.run(...)

with g1.as_default():
  # do some more work with sess1 on g1 
  sess1.run(...)

with g2.as_default():
  # do some more work with sess2 on g2
  sess2.run(...)

sess1.close()
sess2.close()

You don't actually need the with statements, once you've created sess1 and sess2 you can just use them, they'll refer to the correct graph, but it's probably good form to set the default graph whenever you're working with that graph while you're still getting used to how TF juggles global variables.