Tensorflow With Context Manager vs. With Session

Question

I'm trying to understand how a with block behaves in the above scenarios. I assume to always have one graph and one session only. I understand that I have (at least?) 2 ways to work using a session in a with block:

Example 1 : using as_default() that creates a context manager:
Same documentation says:

Use with the with keyword to specify that calls to tf.Operation.run or tf.Tensor.eval should be executed in this session.

# Create session
sess = tf.Session()
# Enter with block on a new context manager
with sess.as_default():
  # Train: Following line should result calling tf.Operation.run
  sess.run([optimizer, loss], feed_dict={x: x_train, y: y_train)
  # Eval: Following line should result calling tf.Tensor.eval
  sess.run([loss], feed_dict={x: x_eval, y: y_eval)

Example 2 : with block on session as stated in same documentation in a lower section:

Alternatively, you can use with tf.Session(): to create a session that is automatically closed on exiting the context, including when an uncaught exception is raised.

# Enter with block on session instead of Context Manager
with tf.Session() as sess:
  # Train: Following line seems calling tf.Operation.run as per my test
  sess.run([optimizer, loss], feed_dict={x: x_train, y: y_train)
  # Eval: Following is unclear
  sess.run([loss], feed_dict={x: x_eval, y: y_eval)

I would like to understand what is the correct usage as I see GitHub examples of both cases but of course without the results. In my tests both Example 1 and 2 works for training. For evaluation it seems there is a difference that I can't understand. It exceeds my knowledge to browse the Tensorflow source. Can someone please explain?

Answer 1

They do slightly different things, so each may or may not be correct depending on the usage. tf.Session.as_default() will just ensure that the session is set as the default ones, so calls to eval and run will use that session by default:

import tensorflow as tf

sess = tf.Session()
with sess.as_default():
    print(sess is tf.get_default_session())  # True

However, as stated in the documentation, tf.Session.as_default() will not automatically close the session after the with block. If you want that, you can use the session itself as the context manager:

import tensorflow as tf

with tf.Session() as sess:
    # Do something with the session
sess.run([])  # RuntimeError: Attempted to use a closed Session.

However, although (from my point of view) not clearly documented, using a session as context manager also makes it the default session.

import tensorflow as tf

with tf.Session() as sess:
    print(sess is tf.get_default_session())  # True

What is the point of tf.Session.as_default(), then? Well, simply when you want to temporarily make a session the default one but you do not one to close it after that (in which case you should manually close it later, or use it as an outer context manager). This is probably most relevant when you have more than one open session. Consider the following case:

import tensorflow as tf

with tf.Session() as sess1, tf.Session() as sess2:
    print(sess2 is tf.get_default_session())  # True

Here sess2 is the default because its context was added later (it can be considered to be "inner" to the context created by sess1). But now maybe you want to make sess1 the default for a while. However you can not use sess1 itself as context manager again:

import tensorflow as tf

with tf.Session() as sess1, tf.Session() as sess2:
    # Do something with sess2
    with sess1:
        # RuntimeError: Session context managers are not re-entrant.
        #               Use `Session.as_default()` if you want to enter
        #               a session multiple times.

So you can switch between one and other default sessions with tf.Session.as_default():

import tensorflow as tf

with tf.Session() as sess1, tf.Session() as sess2:
    with sess1.as_default():
        # Do something with sess1
    with sess2.as_default():
        # Do something with sess2
        # This is not really needed because sess2 was the default
        # in the outer context but you can add it to be explicit
# Both sessions are closed at the end of the outer context

Of course, you can be extra explicit even with one session, if you want:

import tensorflow as tf

with tf.Session() as sess, sess.as_default():
    # ...

Personally, I have never used tf.Session.as_default() in my actual code, but then again I have rarely needed to use multiple sessions, and I prefer to use tf.Session.run() instead of relying on the default session anyway, but that is mostly a matter of personal taste I suppose.