在单元测试中退出 tf.Session() 时重置默认图

Question

• 在每个单元测试结束时，我调用tf.reset_default_graph()清除默认图形。
• 但是，当单元测试失败时，图表不会被清除。这也会导致下一个单元测试失败。

退出tf.Session()上下文时如何清除图形？

示例（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))

Answer 1

我建议使用pytest提供的工具：

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

在每次测试之前和之后都会自动调用夹具（标志autouse），yield在测试之前/之后执行之前/之后的代码。这样，您的问题中的测试将无需任何修改即可运行，并且您遵循 DRY 原则，拒绝在每个测试中编写重复的代码。另一个例子：

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

将在测试执行之前为每个测试创建一个新图。

Fixtures inpytest非常强大，如果使用得当，可以完全消除代码重复。例如，您问题中的测试等效于：

@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))

如果您想了解更多，请从pytest 固定装置开始：显式、模块化、可扩展。