Reset default graph upon exiting tf.Session() in unit tests

Viewed 519
  • At the end of each unit tests, I call tf.reset_default_graph() to clear the default graph.
  • However, when a unit test fails, the graph doesn't get cleared. That makes the next unit test fails as well.

How to clear a graph upon exiting the tf.Session() context?

Example (pytest):

import tensorflow as tf


def test_1():
    x = tf.get_variable('x', initializer=1)
    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())
        print(4 / 0)
        print(sess.run(x))


def test_2():
    x = tf.get_variable('x', initializer=1)
    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())
        print(sess.run(x))
3 Answers

I suggest to use the tools pytest offers:

@pytest.fixture(autouse=True)
def reset():
    yield
    tf.reset_default_graph()

The fixture will be automatically invoked before and after each test (flag autouse), code before/after yield is executed before/after test. This way the tests from your question will work without any modification and you follow the DRY principle, refusing to write duplicated code in each test. Another example:

@pytest.fixture(autouse=True)
def init_graph():
    with tf.Graph().as_default():
        yield

will create a new graph for each test before the test executes.

Fixtures in pytest are very powerful and can fully eliminate code repetitions when used properly. For example, the tests from your question are equivalent to:

@pytest.fixture
def x():
    return tf.get_variable('x', initializer=1)


@pytest.fixture
def session(x):
    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())
        yield sess


@pytest.fixture(autouse=True)
def init_graph():
    with tf.Graph().as_default():
        yield


def test_1(session, x):
    print(4 / 0)
    print(session.run(x))


def test_2(session, x):
    print(session.run(x))

If you want to learn more, start with pytest fixtures: explicit, modular, scalable.

Would something like this work?

import tensorflow as tf


def test_1():
    G = tf.Graph()
    with G.as_default():
        x = tf.get_variable('x', initializer=1)
        with tf.Session() as sess:
            sess.run(tf.initializers.global_variables())
            print(sess.run(x))
            print(4 / 0)


def test_2():
    G = tf.Graph()
    with G.as_default():
        x = tf.get_variable('x', initializer=1)
        with tf.Session() as sess:
            sess.run(tf.initializers.global_variables())
            print(sess.run(x))

A direct solution is to use the try ... finally clause (actually it may be better to put the clause in the code that runs the unit tests rather than in the unit tests directly):

def test_1():
    x = tf.get_variable('x', initializer=1)
    try:
       with tf.Session() as sess:
           sess.run(tf.global_variables_initializer())
           print(4 / 0)
           print(sess.run(x))
    finally:
       tf.reset_default_graph()


def test_2():
    x = tf.get_variable('x', initializer=1)
    try:
        with tf.Session() as sess:
           sess.run(tf.global_variables_initializer())
            print(sess.run(x))
    finally:
        tf.reset_default_graph()

Another clean solution is to use one graph for each unit test as shown in the previous answer. Here is an alternative solution based on this idea with a slightly simplified syntax:

def test_1():
    with tf.Graph().as_default(), tf.Session() as sess:
        x = tf.get_variable('x', initializer=1)

        sess.run(tf.global_variables_initializer())
        print(4 / 0)
        print(sess.run(x))


def test_2():
    with tf.Graph().as_default(), tf.Session() as sess:
        x = tf.get_variable('x', initializer=1)

        sess.run(tf.global_variables_initializer())
        print(sess.run(x))

Similarly to the first solution, the with statement can also be put around the code that runs the unit tests rather than being repeated in every unit test.

Related